From 51d247c2cfbd3aea896c650e5cf25eb0a8e14615 Mon Sep 17 00:00:00 2001 From: Sidney Cammeresi Date: Wed, 7 May 2025 17:00:09 -0700 Subject: [PATCH 1/4] Add Range parameter to `BTreeMap::extract_if` and `BTreeSet::extract_if` This change was requested in the btree_extract_if tracking issue: https://github.com/rust-lang/rust/issues/70530#issuecomment-2486566328 --- library/alloc/src/collections/btree/map.rs | 52 +++++++++++++++++----- library/alloc/src/collections/btree/set.rs | 24 +++++++--- 2 files changed, 57 insertions(+), 19 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index ea81645aa6491..e956511238b4c 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1151,7 +1151,7 @@ impl BTreeMap { K: Ord, F: FnMut(&K, &mut V) -> bool, { - self.extract_if(|k, v| !f(k, v)).for_each(drop); + self.extract_if(.., |k, v| !f(k, v)).for_each(drop); } /// Moves all elements from `other` into `self`, leaving `other` empty. @@ -1429,27 +1429,30 @@ impl BTreeMap { /// assert_eq!(odds.keys().copied().collect::>(), [1, 3, 5, 7]); /// ``` #[unstable(feature = "btree_extract_if", issue = "70530")] - pub fn extract_if(&mut self, pred: F) -> ExtractIf<'_, K, V, F, A> + pub fn extract_if(&mut self, range: R, pred: F) -> ExtractIf<'_, K, V, R, F, A> where K: Ord, + R: RangeBounds, F: FnMut(&K, &mut V) -> bool, { - let (inner, alloc) = self.extract_if_inner(); + let (inner, alloc) = self.extract_if_inner(range); ExtractIf { pred, inner, alloc } } - pub(super) fn extract_if_inner(&mut self) -> (ExtractIfInner<'_, K, V>, A) + pub(super) fn extract_if_inner(&mut self, range: R) -> (ExtractIfInner<'_, K, V, R>, A) where K: Ord, + R: RangeBounds, { if let Some(root) = self.root.as_mut() { let (root, dormant_root) = DormantMutRef::new(root); - let front = root.borrow_mut().first_leaf_edge(); + let first = root.borrow_mut().lower_bound(SearchBound::from_range(range.start_bound())); ( ExtractIfInner { length: &mut self.length, dormant_root: Some(dormant_root), - cur_leaf_edge: Some(front), + cur_leaf_edge: Some(first), + range, }, (*self.alloc).clone(), ) @@ -1459,6 +1462,7 @@ impl BTreeMap { length: &mut self.length, dormant_root: None, cur_leaf_edge: None, + range, }, (*self.alloc).clone(), ) @@ -1917,18 +1921,19 @@ pub struct ExtractIf< 'a, K, V, + R, F, #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + Clone = Global, > { pred: F, - inner: ExtractIfInner<'a, K, V>, + inner: ExtractIfInner<'a, K, V, R>, /// The BTreeMap will outlive this IntoIter so we don't care about drop order for `alloc`. alloc: A, } /// Most of the implementation of ExtractIf are generic over the type /// of the predicate, thus also serving for BTreeSet::ExtractIf. -pub(super) struct ExtractIfInner<'a, K, V> { +pub(super) struct ExtractIfInner<'a, K, V, R> { /// Reference to the length field in the borrowed map, updated live. length: &'a mut usize, /// Buried reference to the root field in the borrowed map. @@ -1938,10 +1943,13 @@ pub(super) struct ExtractIfInner<'a, K, V> { /// Empty if the map has no root, if iteration went beyond the last leaf edge, /// or if a panic occurred in the predicate. cur_leaf_edge: Option, K, V, marker::Leaf>, marker::Edge>>, + /// Range over which iteration was requested. We don't need the left side, but we + /// can't extract the right side without requiring K: Clone. + range: R, } #[unstable(feature = "btree_extract_if", issue = "70530")] -impl fmt::Debug for ExtractIf<'_, K, V, F, A> +impl fmt::Debug for ExtractIf<'_, K, V, R, F, A> where K: fmt::Debug, V: fmt::Debug, @@ -1953,8 +1961,10 @@ where } #[unstable(feature = "btree_extract_if", issue = "70530")] -impl Iterator for ExtractIf<'_, K, V, F, A> +impl Iterator for ExtractIf<'_, K, V, R, F, A> where + K: PartialOrd, + R: RangeBounds, F: FnMut(&K, &mut V) -> bool, { type Item = (K, V); @@ -1968,7 +1978,7 @@ where } } -impl<'a, K, V> ExtractIfInner<'a, K, V> { +impl<'a, K, V, R> ExtractIfInner<'a, K, V, R> { /// Allow Debug implementations to predict the next element. pub(super) fn peek(&self) -> Option<(&K, &V)> { let edge = self.cur_leaf_edge.as_ref()?; @@ -1978,10 +1988,22 @@ impl<'a, K, V> ExtractIfInner<'a, K, V> { /// Implementation of a typical `ExtractIf::next` method, given the predicate. pub(super) fn next(&mut self, pred: &mut F, alloc: A) -> Option<(K, V)> where + K: PartialOrd, + R: RangeBounds, F: FnMut(&K, &mut V) -> bool, { while let Ok(mut kv) = self.cur_leaf_edge.take()?.next_kv() { let (k, v) = kv.kv_mut(); + + // On creation, we navigated directly to the left bound, so we need only check the + // right bound here to decide whether to stop. + match self.range.end_bound() { + Bound::Included(ref end) if (*k).le(end) => (), + Bound::Excluded(ref end) if (*k).lt(end) => (), + Bound::Unbounded => (), + _ => return None, + } + if pred(k, v) { *self.length -= 1; let (kv, pos) = kv.remove_kv_tracking( @@ -2013,7 +2035,13 @@ impl<'a, K, V> ExtractIfInner<'a, K, V> { } #[unstable(feature = "btree_extract_if", issue = "70530")] -impl FusedIterator for ExtractIf<'_, K, V, F> where F: FnMut(&K, &mut V) -> bool {} +impl FusedIterator for ExtractIf<'_, K, V, R, F> +where + K: PartialOrd, + R: RangeBounds, + F: FnMut(&K, &mut V) -> bool, +{ +} #[stable(feature = "btree_range", since = "1.17.0")] impl<'a, K, V> Iterator for Range<'a, K, V> { diff --git a/library/alloc/src/collections/btree/set.rs b/library/alloc/src/collections/btree/set.rs index 343934680b87a..ec840ae64873c 100644 --- a/library/alloc/src/collections/btree/set.rs +++ b/library/alloc/src/collections/btree/set.rs @@ -1109,7 +1109,7 @@ impl BTreeSet { T: Ord, F: FnMut(&T) -> bool, { - self.extract_if(|v| !f(v)).for_each(drop); + self.extract_if(.., |v| !f(v)).for_each(drop); } /// Moves all elements from `other` into `self`, leaving `other` empty. @@ -1214,12 +1214,13 @@ impl BTreeSet { /// assert_eq!(odds.into_iter().collect::>(), vec![1, 3, 5, 7]); /// ``` #[unstable(feature = "btree_extract_if", issue = "70530")] - pub fn extract_if<'a, F>(&'a mut self, pred: F) -> ExtractIf<'a, T, F, A> + pub fn extract_if<'a, F, R>(&'a mut self, range: R, pred: F) -> ExtractIf<'a, T, R, F, A> where T: Ord, + R: RangeBounds, F: 'a + FnMut(&T) -> bool, { - let (inner, alloc) = self.map.extract_if_inner(); + let (inner, alloc) = self.map.extract_if_inner(range); ExtractIf { pred, inner, alloc } } @@ -1554,17 +1555,18 @@ impl<'a, T, A: Allocator + Clone> IntoIterator for &'a BTreeSet { pub struct ExtractIf< 'a, T, + R, F, #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator + Clone = Global, > { pred: F, - inner: super::map::ExtractIfInner<'a, T, SetValZST>, + inner: super::map::ExtractIfInner<'a, T, SetValZST, R>, /// The BTreeMap will outlive this IntoIter so we don't care about drop order for `alloc`. alloc: A, } #[unstable(feature = "btree_extract_if", issue = "70530")] -impl fmt::Debug for ExtractIf<'_, T, F, A> +impl fmt::Debug for ExtractIf<'_, T, R, F, A> where T: fmt::Debug, A: Allocator + Clone, @@ -1577,8 +1579,10 @@ where } #[unstable(feature = "btree_extract_if", issue = "70530")] -impl<'a, T, F, A: Allocator + Clone> Iterator for ExtractIf<'_, T, F, A> +impl<'a, T, R, F, A: Allocator + Clone> Iterator for ExtractIf<'_, T, R, F, A> where + T: PartialOrd, + R: RangeBounds, F: 'a + FnMut(&T) -> bool, { type Item = T; @@ -1595,7 +1599,13 @@ where } #[unstable(feature = "btree_extract_if", issue = "70530")] -impl FusedIterator for ExtractIf<'_, T, F, A> where F: FnMut(&T) -> bool {} +impl FusedIterator for ExtractIf<'_, T, R, F, A> +where + T: PartialOrd, + R: RangeBounds, + F: FnMut(&T) -> bool, +{ +} #[stable(feature = "rust1", since = "1.0.0")] impl Extend for BTreeSet { From 1ae96fcd792f000da741e4c73effb7bfb89cee1d Mon Sep 17 00:00:00 2001 From: Sidney Cammeresi Date: Wed, 7 May 2025 17:05:44 -0700 Subject: [PATCH 2/4] Update tests with Range parameter to `BTreeMap::extract_if` etc. --- .../alloc/src/collections/btree/map/tests.rs | 50 +++++++++---------- .../alloc/src/collections/btree/set/tests.rs | 12 ++--- library/alloctests/benches/btree/map.rs | 12 ++--- library/alloctests/benches/btree/set.rs | 8 +-- library/alloctests/tests/autotraits.rs | 14 +++++- src/tools/miri/tests/pass/btreemap.rs | 2 +- .../lit-pattern-matching-with-methods.rs | 4 +- 7 files changed, 56 insertions(+), 46 deletions(-) diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index 5975134382e7d..10add43892635 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -944,7 +944,7 @@ mod test_extract_if { #[test] fn empty() { let mut map: BTreeMap = BTreeMap::new(); - map.extract_if(|_, _| unreachable!("there's nothing to decide on")).for_each(drop); + map.extract_if(.., |_, _| unreachable!("there's nothing to decide on")).for_each(drop); assert_eq!(map.height(), None); map.check(); } @@ -954,7 +954,7 @@ mod test_extract_if { fn consumed_keeping_all() { let pairs = (0..3).map(|i| (i, i)); let mut map = BTreeMap::from_iter(pairs); - assert!(map.extract_if(|_, _| false).eq(iter::empty())); + assert!(map.extract_if(.., |_, _| false).eq(iter::empty())); map.check(); } @@ -963,7 +963,7 @@ mod test_extract_if { fn consumed_removing_all() { let pairs = (0..3).map(|i| (i, i)); let mut map = BTreeMap::from_iter(pairs.clone()); - assert!(map.extract_if(|_, _| true).eq(pairs)); + assert!(map.extract_if(.., |_, _| true).eq(pairs)); assert!(map.is_empty()); map.check(); } @@ -974,7 +974,7 @@ mod test_extract_if { let pairs = (0..3).map(|i| (i, i)); let mut map = BTreeMap::from_iter(pairs); assert!( - map.extract_if(|_, v| { + map.extract_if(.., |_, v| { *v += 6; false }) @@ -991,7 +991,7 @@ mod test_extract_if { let pairs = (0..3).map(|i| (i, i)); let mut map = BTreeMap::from_iter(pairs); assert!( - map.extract_if(|_, v| { + map.extract_if(.., |_, v| { *v += 6; true }) @@ -1005,7 +1005,7 @@ mod test_extract_if { fn underfull_keeping_all() { let pairs = (0..3).map(|i| (i, i)); let mut map = BTreeMap::from_iter(pairs); - map.extract_if(|_, _| false).for_each(drop); + map.extract_if(.., |_, _| false).for_each(drop); assert!(map.keys().copied().eq(0..3)); map.check(); } @@ -1015,7 +1015,7 @@ mod test_extract_if { let pairs = (0..3).map(|i| (i, i)); for doomed in 0..3 { let mut map = BTreeMap::from_iter(pairs.clone()); - map.extract_if(|i, _| *i == doomed).for_each(drop); + map.extract_if(.., |i, _| *i == doomed).for_each(drop); assert_eq!(map.len(), 2); map.check(); } @@ -1026,7 +1026,7 @@ mod test_extract_if { let pairs = (0..3).map(|i| (i, i)); for sacred in 0..3 { let mut map = BTreeMap::from_iter(pairs.clone()); - map.extract_if(|i, _| *i != sacred).for_each(drop); + map.extract_if(.., |i, _| *i != sacred).for_each(drop); assert!(map.keys().copied().eq(sacred..=sacred)); map.check(); } @@ -1036,7 +1036,7 @@ mod test_extract_if { fn underfull_removing_all() { let pairs = (0..3).map(|i| (i, i)); let mut map = BTreeMap::from_iter(pairs); - map.extract_if(|_, _| true).for_each(drop); + map.extract_if(.., |_, _| true).for_each(drop); assert!(map.is_empty()); map.check(); } @@ -1045,7 +1045,7 @@ mod test_extract_if { fn height_0_keeping_all() { let pairs = (0..node::CAPACITY).map(|i| (i, i)); let mut map = BTreeMap::from_iter(pairs); - map.extract_if(|_, _| false).for_each(drop); + map.extract_if(.., |_, _| false).for_each(drop); assert!(map.keys().copied().eq(0..node::CAPACITY)); map.check(); } @@ -1055,7 +1055,7 @@ mod test_extract_if { let pairs = (0..node::CAPACITY).map(|i| (i, i)); for doomed in 0..node::CAPACITY { let mut map = BTreeMap::from_iter(pairs.clone()); - map.extract_if(|i, _| *i == doomed).for_each(drop); + map.extract_if(.., |i, _| *i == doomed).for_each(drop); assert_eq!(map.len(), node::CAPACITY - 1); map.check(); } @@ -1066,7 +1066,7 @@ mod test_extract_if { let pairs = (0..node::CAPACITY).map(|i| (i, i)); for sacred in 0..node::CAPACITY { let mut map = BTreeMap::from_iter(pairs.clone()); - map.extract_if(|i, _| *i != sacred).for_each(drop); + map.extract_if(.., |i, _| *i != sacred).for_each(drop); assert!(map.keys().copied().eq(sacred..=sacred)); map.check(); } @@ -1076,7 +1076,7 @@ mod test_extract_if { fn height_0_removing_all() { let pairs = (0..node::CAPACITY).map(|i| (i, i)); let mut map = BTreeMap::from_iter(pairs); - map.extract_if(|_, _| true).for_each(drop); + map.extract_if(.., |_, _| true).for_each(drop); assert!(map.is_empty()); map.check(); } @@ -1084,7 +1084,7 @@ mod test_extract_if { #[test] fn height_0_keeping_half() { let mut map = BTreeMap::from_iter((0..16).map(|i| (i, i))); - assert_eq!(map.extract_if(|i, _| *i % 2 == 0).count(), 8); + assert_eq!(map.extract_if(.., |i, _| *i % 2 == 0).count(), 8); assert_eq!(map.len(), 8); map.check(); } @@ -1093,7 +1093,7 @@ mod test_extract_if { fn height_1_removing_all() { let pairs = (0..MIN_INSERTS_HEIGHT_1).map(|i| (i, i)); let mut map = BTreeMap::from_iter(pairs); - map.extract_if(|_, _| true).for_each(drop); + map.extract_if(.., |_, _| true).for_each(drop); assert!(map.is_empty()); map.check(); } @@ -1103,7 +1103,7 @@ mod test_extract_if { let pairs = (0..MIN_INSERTS_HEIGHT_1).map(|i| (i, i)); for doomed in 0..MIN_INSERTS_HEIGHT_1 { let mut map = BTreeMap::from_iter(pairs.clone()); - map.extract_if(|i, _| *i == doomed).for_each(drop); + map.extract_if(.., |i, _| *i == doomed).for_each(drop); assert_eq!(map.len(), MIN_INSERTS_HEIGHT_1 - 1); map.check(); } @@ -1114,7 +1114,7 @@ mod test_extract_if { let pairs = (0..MIN_INSERTS_HEIGHT_1).map(|i| (i, i)); for sacred in 0..MIN_INSERTS_HEIGHT_1 { let mut map = BTreeMap::from_iter(pairs.clone()); - map.extract_if(|i, _| *i != sacred).for_each(drop); + map.extract_if(.., |i, _| *i != sacred).for_each(drop); assert!(map.keys().copied().eq(sacred..=sacred)); map.check(); } @@ -1125,7 +1125,7 @@ mod test_extract_if { let pairs = (0..MIN_INSERTS_HEIGHT_2).map(|i| (i, i)); for doomed in (0..MIN_INSERTS_HEIGHT_2).step_by(12) { let mut map = BTreeMap::from_iter(pairs.clone()); - map.extract_if(|i, _| *i == doomed).for_each(drop); + map.extract_if(.., |i, _| *i == doomed).for_each(drop); assert_eq!(map.len(), MIN_INSERTS_HEIGHT_2 - 1); map.check(); } @@ -1136,7 +1136,7 @@ mod test_extract_if { let pairs = (0..MIN_INSERTS_HEIGHT_2).map(|i| (i, i)); for sacred in (0..MIN_INSERTS_HEIGHT_2).step_by(12) { let mut map = BTreeMap::from_iter(pairs.clone()); - map.extract_if(|i, _| *i != sacred).for_each(drop); + map.extract_if(.., |i, _| *i != sacred).for_each(drop); assert!(map.keys().copied().eq(sacred..=sacred)); map.check(); } @@ -1146,7 +1146,7 @@ mod test_extract_if { fn height_2_removing_all() { let pairs = (0..MIN_INSERTS_HEIGHT_2).map(|i| (i, i)); let mut map = BTreeMap::from_iter(pairs); - map.extract_if(|_, _| true).for_each(drop); + map.extract_if(.., |_, _| true).for_each(drop); assert!(map.is_empty()); map.check(); } @@ -1162,7 +1162,7 @@ mod test_extract_if { map.insert(b.spawn(Panic::InDrop), ()); map.insert(c.spawn(Panic::Never), ()); - catch_unwind(move || map.extract_if(|dummy, _| dummy.query(true)).for_each(drop)) + catch_unwind(move || map.extract_if(.., |dummy, _| dummy.query(true)).for_each(drop)) .unwrap_err(); assert_eq!(a.queried(), 1); @@ -1185,7 +1185,7 @@ mod test_extract_if { map.insert(c.spawn(Panic::InQuery), ()); catch_unwind(AssertUnwindSafe(|| { - map.extract_if(|dummy, _| dummy.query(true)).for_each(drop) + map.extract_if(.., |dummy, _| dummy.query(true)).for_each(drop) })) .unwrap_err(); @@ -1214,7 +1214,7 @@ mod test_extract_if { map.insert(c.spawn(Panic::InQuery), ()); { - let mut it = map.extract_if(|dummy, _| dummy.query(true)); + let mut it = map.extract_if(.., |dummy, _| dummy.query(true)); catch_unwind(AssertUnwindSafe(|| while it.next().is_some() {})).unwrap_err(); // Iterator behavior after a panic is explicitly unspecified, // so this is just the current implementation: @@ -1658,7 +1658,7 @@ fn assert_sync() { } fn extract_if(v: &mut BTreeMap) -> impl Sync + '_ { - v.extract_if(|_, _| false) + v.extract_if(.., |_, _| false) } fn iter(v: &BTreeMap) -> impl Sync + '_ { @@ -1727,7 +1727,7 @@ fn assert_send() { } fn extract_if(v: &mut BTreeMap) -> impl Send + '_ { - v.extract_if(|_, _| false) + v.extract_if(.., |_, _| false) } fn iter(v: &BTreeMap) -> impl Send + '_ { diff --git a/library/alloc/src/collections/btree/set/tests.rs b/library/alloc/src/collections/btree/set/tests.rs index d538ef707eb93..85c9a98c461d9 100644 --- a/library/alloc/src/collections/btree/set/tests.rs +++ b/library/alloc/src/collections/btree/set/tests.rs @@ -368,8 +368,8 @@ fn test_extract_if() { let mut x = BTreeSet::from([1]); let mut y = BTreeSet::from([1]); - x.extract_if(|_| true).for_each(drop); - y.extract_if(|_| false).for_each(drop); + x.extract_if(.., |_| true).for_each(drop); + y.extract_if(.., |_| false).for_each(drop); assert_eq!(x.len(), 0); assert_eq!(y.len(), 1); } @@ -385,7 +385,7 @@ fn test_extract_if_drop_panic_leak() { set.insert(b.spawn(Panic::InDrop)); set.insert(c.spawn(Panic::Never)); - catch_unwind(move || set.extract_if(|dummy| dummy.query(true)).for_each(drop)).ok(); + catch_unwind(move || set.extract_if(.., |dummy| dummy.query(true)).for_each(drop)).ok(); assert_eq!(a.queried(), 1); assert_eq!(b.queried(), 1); @@ -406,7 +406,7 @@ fn test_extract_if_pred_panic_leak() { set.insert(b.spawn(Panic::InQuery)); set.insert(c.spawn(Panic::InQuery)); - catch_unwind(AssertUnwindSafe(|| set.extract_if(|dummy| dummy.query(true)).for_each(drop))) + catch_unwind(AssertUnwindSafe(|| set.extract_if(.., |dummy| dummy.query(true)).for_each(drop))) .ok(); assert_eq!(a.queried(), 1); @@ -605,7 +605,7 @@ fn assert_sync() { } fn extract_if(v: &mut BTreeSet) -> impl Sync + '_ { - v.extract_if(|_| false) + v.extract_if(.., |_| false) } fn difference(v: &BTreeSet) -> impl Sync + '_ { @@ -644,7 +644,7 @@ fn assert_send() { } fn extract_if(v: &mut BTreeSet) -> impl Send + '_ { - v.extract_if(|_| false) + v.extract_if(.., |_| false) } fn difference(v: &BTreeSet) -> impl Send + '_ { diff --git a/library/alloctests/benches/btree/map.rs b/library/alloctests/benches/btree/map.rs index 20f02dc3a968a..778065fd96575 100644 --- a/library/alloctests/benches/btree/map.rs +++ b/library/alloctests/benches/btree/map.rs @@ -386,7 +386,7 @@ pub fn clone_slim_100_and_clear(b: &mut Bencher) { #[bench] pub fn clone_slim_100_and_drain_all(b: &mut Bencher) { let src = slim_map(100); - b.iter(|| src.clone().extract_if(|_, _| true).count()) + b.iter(|| src.clone().extract_if(.., |_, _| true).count()) } #[bench] @@ -394,7 +394,7 @@ pub fn clone_slim_100_and_drain_half(b: &mut Bencher) { let src = slim_map(100); b.iter(|| { let mut map = src.clone(); - assert_eq!(map.extract_if(|i, _| i % 2 == 0).count(), 100 / 2); + assert_eq!(map.extract_if(.., |i, _| i % 2 == 0).count(), 100 / 2); assert_eq!(map.len(), 100 / 2); }) } @@ -457,7 +457,7 @@ pub fn clone_slim_10k_and_clear(b: &mut Bencher) { #[bench] pub fn clone_slim_10k_and_drain_all(b: &mut Bencher) { let src = slim_map(10_000); - b.iter(|| src.clone().extract_if(|_, _| true).count()) + b.iter(|| src.clone().extract_if(.., |_, _| true).count()) } #[bench] @@ -465,7 +465,7 @@ pub fn clone_slim_10k_and_drain_half(b: &mut Bencher) { let src = slim_map(10_000); b.iter(|| { let mut map = src.clone(); - assert_eq!(map.extract_if(|i, _| i % 2 == 0).count(), 10_000 / 2); + assert_eq!(map.extract_if(.., |i, _| i % 2 == 0).count(), 10_000 / 2); assert_eq!(map.len(), 10_000 / 2); }) } @@ -528,7 +528,7 @@ pub fn clone_fat_val_100_and_clear(b: &mut Bencher) { #[bench] pub fn clone_fat_val_100_and_drain_all(b: &mut Bencher) { let src = fat_val_map(100); - b.iter(|| src.clone().extract_if(|_, _| true).count()) + b.iter(|| src.clone().extract_if(.., |_, _| true).count()) } #[bench] @@ -536,7 +536,7 @@ pub fn clone_fat_val_100_and_drain_half(b: &mut Bencher) { let src = fat_val_map(100); b.iter(|| { let mut map = src.clone(); - assert_eq!(map.extract_if(|i, _| i % 2 == 0).count(), 100 / 2); + assert_eq!(map.extract_if(.., |i, _| i % 2 == 0).count(), 100 / 2); assert_eq!(map.len(), 100 / 2); }) } diff --git a/library/alloctests/benches/btree/set.rs b/library/alloctests/benches/btree/set.rs index 5aa395b4d52ac..027c86a89a51a 100644 --- a/library/alloctests/benches/btree/set.rs +++ b/library/alloctests/benches/btree/set.rs @@ -69,7 +69,7 @@ pub fn clone_100_and_clear(b: &mut Bencher) { #[bench] pub fn clone_100_and_drain_all(b: &mut Bencher) { let src = slim_set(100); - b.iter(|| src.clone().extract_if(|_| true).count()) + b.iter(|| src.clone().extract_if(.., |_| true).count()) } #[bench] @@ -77,7 +77,7 @@ pub fn clone_100_and_drain_half(b: &mut Bencher) { let src = slim_set(100); b.iter(|| { let mut set = src.clone(); - assert_eq!(set.extract_if(|i| i % 2 == 0).count(), 100 / 2); + assert_eq!(set.extract_if(.., |i| i % 2 == 0).count(), 100 / 2); assert_eq!(set.len(), 100 / 2); }) } @@ -140,7 +140,7 @@ pub fn clone_10k_and_clear(b: &mut Bencher) { #[bench] pub fn clone_10k_and_drain_all(b: &mut Bencher) { let src = slim_set(10_000); - b.iter(|| src.clone().extract_if(|_| true).count()) + b.iter(|| src.clone().extract_if(.., |_| true).count()) } #[bench] @@ -148,7 +148,7 @@ pub fn clone_10k_and_drain_half(b: &mut Bencher) { let src = slim_set(10_000); b.iter(|| { let mut set = src.clone(); - assert_eq!(set.extract_if(|i| i % 2 == 0).count(), 10_000 / 2); + assert_eq!(set.extract_if(.., |i| i % 2 == 0).count(), 10_000 / 2); assert_eq!(set.len(), 10_000 / 2); }) } diff --git a/library/alloctests/tests/autotraits.rs b/library/alloctests/tests/autotraits.rs index 6b82deeac8acb..ad0a103859696 100644 --- a/library/alloctests/tests/autotraits.rs +++ b/library/alloctests/tests/autotraits.rs @@ -1,3 +1,5 @@ +use std::ops::Range; + fn require_sync(_: T) {} fn require_send_sync(_: T) {} @@ -55,7 +57,13 @@ fn test_btree_map() { require_send_sync(async { let _v = None::< - alloc::collections::btree_map::ExtractIf<'_, &u32, &u32, fn(&&u32, &mut &u32) -> bool>, + alloc::collections::btree_map::ExtractIf< + '_, + &u32, + &u32, + Range, + fn(&&u32, &mut &u32) -> bool, + >, >; async {}.await; }); @@ -144,7 +152,9 @@ fn test_btree_set() { }); require_send_sync(async { - let _v = None:: bool>>; + let _v = None::< + alloc::collections::btree_set::ExtractIf<'_, &u32, Range, fn(&&u32) -> bool>, + >; async {}.await; }); diff --git a/src/tools/miri/tests/pass/btreemap.rs b/src/tools/miri/tests/pass/btreemap.rs index 1213f81a6f128..1d65e69bf726f 100644 --- a/src/tools/miri/tests/pass/btreemap.rs +++ b/src/tools/miri/tests/pass/btreemap.rs @@ -50,7 +50,7 @@ pub fn main() { test_all_refs(&mut 13, b.values_mut()); // Test forgetting the extractor. - let mut d = b.extract_if(|_, i| *i < 30); + let mut d = b.extract_if(.., |_, i| *i < 30); d.next().unwrap(); mem::forget(d); } diff --git a/tests/ui/closures/2229_closure_analysis/run_pass/lit-pattern-matching-with-methods.rs b/tests/ui/closures/2229_closure_analysis/run_pass/lit-pattern-matching-with-methods.rs index 7a4d7d9a81ec4..afb16cf58e871 100644 --- a/tests/ui/closures/2229_closure_analysis/run_pass/lit-pattern-matching-with-methods.rs +++ b/tests/ui/closures/2229_closure_analysis/run_pass/lit-pattern-matching-with-methods.rs @@ -14,14 +14,14 @@ fn main() { map.insert("c", ()); { - let mut it = map.extract_if(|_, _| true); + let mut it = map.extract_if(.., |_, _| true); catch_unwind(AssertUnwindSafe(|| while it.next().is_some() {})).unwrap_err(); let result = catch_unwind(AssertUnwindSafe(|| it.next())); assert!(matches!(result, Ok(None))); } { - let mut it = map.extract_if(|_, _| true); + let mut it = map.extract_if(.., |_, _| true); catch_unwind(AssertUnwindSafe(|| while let Some(_) = it.next() {})).unwrap_err(); let result = catch_unwind(AssertUnwindSafe(|| it.next())); assert!(matches!(result, Ok(None))); From 38c37eb3cbd87894cdab588b5002fa07803feff7 Mon Sep 17 00:00:00 2001 From: Sidney Cammeresi Date: Wed, 7 May 2025 17:06:53 -0700 Subject: [PATCH 3/4] Update docs for new Range parameter to `BTreeMap::extract_if` etc. --- library/alloc/src/collections/btree/map.rs | 10 ++++++++-- library/alloc/src/collections/btree/set.rs | 10 ++++++++-- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index e956511238b4c..52b98291ff93f 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1397,7 +1397,7 @@ impl BTreeMap { } } - /// Creates an iterator that visits all elements (key-value pairs) in + /// Creates an iterator that visits elements (key-value pairs) in the specified range in /// ascending key order and uses a closure to determine if an element /// should be removed. /// @@ -1423,10 +1423,16 @@ impl BTreeMap { /// use std::collections::BTreeMap; /// /// let mut map: BTreeMap = (0..8).map(|x| (x, x)).collect(); - /// let evens: BTreeMap<_, _> = map.extract_if(|k, _v| k % 2 == 0).collect(); + /// let evens: BTreeMap<_, _> = map.extract_if(.., |k, _v| k % 2 == 0).collect(); /// let odds = map; /// assert_eq!(evens.keys().copied().collect::>(), [0, 2, 4, 6]); /// assert_eq!(odds.keys().copied().collect::>(), [1, 3, 5, 7]); + /// + /// let mut map: BTreeMap = (0..8).map(|x| (x, x)).collect(); + /// let low: BTreeMap<_, _> = map.extract_if(0..4, |_k, _v| true).collect(); + /// let high = map; + /// assert_eq!(low.keys().copied().collect::>(), [0, 1, 2, 3]); + /// assert_eq!(high.keys().copied().collect::>(), [4, 5, 6, 7]); /// ``` #[unstable(feature = "btree_extract_if", issue = "70530")] pub fn extract_if(&mut self, range: R, pred: F) -> ExtractIf<'_, K, V, R, F, A> diff --git a/library/alloc/src/collections/btree/set.rs b/library/alloc/src/collections/btree/set.rs index ec840ae64873c..780bd8b0dd143 100644 --- a/library/alloc/src/collections/btree/set.rs +++ b/library/alloc/src/collections/btree/set.rs @@ -1187,7 +1187,7 @@ impl BTreeSet { BTreeSet { map: self.map.split_off(value) } } - /// Creates an iterator that visits all elements in ascending order and + /// Creates an iterator that visits elements in the specified range in ascending order and /// uses a closure to determine if an element should be removed. /// /// If the closure returns `true`, the element is removed from the set and @@ -1208,10 +1208,16 @@ impl BTreeSet { /// use std::collections::BTreeSet; /// /// let mut set: BTreeSet = (0..8).collect(); - /// let evens: BTreeSet<_> = set.extract_if(|v| v % 2 == 0).collect(); + /// let evens: BTreeSet<_> = set.extract_if(.., |v| v % 2 == 0).collect(); /// let odds = set; /// assert_eq!(evens.into_iter().collect::>(), vec![0, 2, 4, 6]); /// assert_eq!(odds.into_iter().collect::>(), vec![1, 3, 5, 7]); + /// + /// let mut map: BTreeSet = (0..8).collect(); + /// let low: BTreeSet<_> = map.extract_if(0..4, |_v| true).collect(); + /// let high = map; + /// assert_eq!(low.into_iter().collect::>(), [0, 1, 2, 3]); + /// assert_eq!(high.into_iter().collect::>(), [4, 5, 6, 7]); /// ``` #[unstable(feature = "btree_extract_if", issue = "70530")] pub fn extract_if<'a, F, R>(&'a mut self, range: R, pred: F) -> ExtractIf<'a, T, R, F, A> From 8656d9e619f7b8ee3bde8bc7ee57a7847c0a5b10 Mon Sep 17 00:00:00 2001 From: Sidney Cammeresi Date: Wed, 7 May 2025 17:07:27 -0700 Subject: [PATCH 4/4] Unit test for Range parameter of `BTreeMap::extract_if` --- .../alloc/src/collections/btree/map/tests.rs | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index 10add43892635..79879d31d3dfd 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -968,6 +968,30 @@ mod test_extract_if { map.check(); } + #[test] + fn consumed_removing_some() { + let pairs = (0..3).map(|i| (i, i)); + let map = BTreeMap::from_iter(pairs); + for x in 0..3 { + for y in 0..3 { + let mut map = map.clone(); + assert!(map.extract_if(x..y, |_, _| true).eq((x..y).map(|i| (i, i)))); + for i in 0..3 { + assert_ne!(map.contains_key(&i), (x..y).contains(&i)); + } + } + } + for x in 0..3 { + for y in 0..2 { + let mut map = map.clone(); + assert!(map.extract_if(x..=y, |_, _| true).eq((x..=y).map(|i| (i, i)))); + for i in 0..3 { + assert_ne!(map.contains_key(&i), (x..=y).contains(&i)); + } + } + } + } + // Explicitly consumes the iterator and modifies values through it. #[test] fn mutating_and_keeping() {