@@ -1151,7 +1151,7 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
11511151 K : Ord ,
11521152 F : FnMut ( & K , & mut V ) -> bool ,
11531153 {
1154- self . extract_if ( |k, v| !f ( k, v) ) . for_each ( drop) ;
1154+ self . extract_if ( .. , |k, v| !f ( k, v) ) . for_each ( drop) ;
11551155 }
11561156
11571157 /// Moves all elements from `other` into `self`, leaving `other` empty.
@@ -1397,11 +1397,13 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
13971397 }
13981398 }
13991399
1400- /// Creates an iterator that visits all elements (key-value pairs) in
1401- /// ascending key order and uses a closure to determine if an element should
1402- /// be removed. If the closure returns `true`, the element is removed from
1403- /// the map and yielded. If the closure returns `false`, or panics, the
1404- /// element remains in the map and will not be yielded.
1400+ /// Creates an iterator that visits elements (key-value pairs) in the specified range in
1401+ /// ascending key order and uses a closure to determine if an element
1402+ /// should be removed.
1403+ ///
1404+ /// If the closure returns `true`, the element is removed from the map and
1405+ /// yielded. If the closure returns `false`, or panics, the element remains
1406+ /// in the map and will not be yielded.
14051407 ///
14061408 /// The iterator also lets you mutate the value of each element in the
14071409 /// closure, regardless of whether you choose to keep or remove it.
@@ -1421,33 +1423,42 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
14211423 /// use std::collections::BTreeMap;
14221424 ///
14231425 /// let mut map: BTreeMap<i32, i32> = (0..8).map(|x| (x, x)).collect();
1424- /// let evens: BTreeMap<_, _> = map.extract_if(|k, _v| k % 2 == 0).collect();
1426+ /// let evens: BTreeMap<_, _> = map.extract_if(.., |k, _v| k % 2 == 0).collect();
14251427 /// let odds = map;
14261428 /// assert_eq!(evens.keys().copied().collect::<Vec<_>>(), [0, 2, 4, 6]);
14271429 /// assert_eq!(odds.keys().copied().collect::<Vec<_>>(), [1, 3, 5, 7]);
1430+ ///
1431+ /// let mut map: BTreeMap<i32, i32> = (0..8).map(|x| (x, x)).collect();
1432+ /// let low: BTreeMap<_, _> = map.extract_if(0..4, |_k, _v| true).collect();
1433+ /// let high = map;
1434+ /// assert_eq!(low.keys().copied().collect::<Vec<_>>(), [0, 1, 2, 3]);
1435+ /// assert_eq!(high.keys().copied().collect::<Vec<_>>(), [4, 5, 6, 7]);
14281436 /// ```
14291437 #[ unstable( feature = "btree_extract_if" , issue = "70530" ) ]
1430- pub fn extract_if < F > ( & mut self , pred : F ) -> ExtractIf < ' _ , K , V , F , A >
1438+ pub fn extract_if < F , R > ( & mut self , range : R , pred : F ) -> ExtractIf < ' _ , K , V , R , F , A >
14311439 where
14321440 K : Ord ,
1441+ R : RangeBounds < K > ,
14331442 F : FnMut ( & K , & mut V ) -> bool ,
14341443 {
1435- let ( inner, alloc) = self . extract_if_inner ( ) ;
1444+ let ( inner, alloc) = self . extract_if_inner ( range ) ;
14361445 ExtractIf { pred, inner, alloc }
14371446 }
14381447
1439- pub ( super ) fn extract_if_inner ( & mut self ) -> ( ExtractIfInner < ' _ , K , V > , A )
1448+ pub ( super ) fn extract_if_inner < R > ( & mut self , range : R ) -> ( ExtractIfInner < ' _ , K , V , R > , A )
14401449 where
14411450 K : Ord ,
1451+ R : RangeBounds < K > ,
14421452 {
14431453 if let Some ( root) = self . root . as_mut ( ) {
14441454 let ( root, dormant_root) = DormantMutRef :: new ( root) ;
1445- let front = root. borrow_mut ( ) . first_leaf_edge ( ) ;
1455+ let first = root. borrow_mut ( ) . lower_bound ( SearchBound :: from_range ( range . start_bound ( ) ) ) ;
14461456 (
14471457 ExtractIfInner {
14481458 length : & mut self . length ,
14491459 dormant_root : Some ( dormant_root) ,
1450- cur_leaf_edge : Some ( front) ,
1460+ cur_leaf_edge : Some ( first) ,
1461+ range,
14511462 } ,
14521463 ( * self . alloc ) . clone ( ) ,
14531464 )
@@ -1457,6 +1468,7 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
14571468 length : & mut self . length ,
14581469 dormant_root : None ,
14591470 cur_leaf_edge : None ,
1471+ range,
14601472 } ,
14611473 ( * self . alloc ) . clone ( ) ,
14621474 )
@@ -1915,18 +1927,19 @@ pub struct ExtractIf<
19151927 ' a ,
19161928 K ,
19171929 V ,
1930+ R ,
19181931 F ,
19191932 #[ unstable( feature = "allocator_api" , issue = "32838" ) ] A : Allocator + Clone = Global ,
19201933> {
19211934 pred : F ,
1922- inner : ExtractIfInner < ' a , K , V > ,
1935+ inner : ExtractIfInner < ' a , K , V , R > ,
19231936 /// The BTreeMap will outlive this IntoIter so we don't care about drop order for `alloc`.
19241937 alloc : A ,
19251938}
19261939
19271940/// Most of the implementation of ExtractIf are generic over the type
19281941/// of the predicate, thus also serving for BTreeSet::ExtractIf.
1929- pub ( super ) struct ExtractIfInner < ' a , K , V > {
1942+ pub ( super ) struct ExtractIfInner < ' a , K , V , R > {
19301943 /// Reference to the length field in the borrowed map, updated live.
19311944 length : & ' a mut usize ,
19321945 /// Buried reference to the root field in the borrowed map.
@@ -1936,10 +1949,13 @@ pub(super) struct ExtractIfInner<'a, K, V> {
19361949 /// Empty if the map has no root, if iteration went beyond the last leaf edge,
19371950 /// or if a panic occurred in the predicate.
19381951 cur_leaf_edge : Option < Handle < NodeRef < marker:: Mut < ' a > , K , V , marker:: Leaf > , marker:: Edge > > ,
1952+ /// Range over which iteration was requested. We don't need the left side, but we
1953+ /// can't extract the right side without requiring K: Clone.
1954+ range : R ,
19391955}
19401956
19411957#[ unstable( feature = "btree_extract_if" , issue = "70530" ) ]
1942- impl < K , V , F , A > fmt:: Debug for ExtractIf < ' _ , K , V , F , A >
1958+ impl < K , V , R , F , A > fmt:: Debug for ExtractIf < ' _ , K , V , R , F , A >
19431959where
19441960 K : fmt:: Debug ,
19451961 V : fmt:: Debug ,
@@ -1951,8 +1967,10 @@ where
19511967}
19521968
19531969#[ unstable( feature = "btree_extract_if" , issue = "70530" ) ]
1954- impl < K , V , F , A : Allocator + Clone > Iterator for ExtractIf < ' _ , K , V , F , A >
1970+ impl < K , V , R , F , A : Allocator + Clone > Iterator for ExtractIf < ' _ , K , V , R , F , A >
19551971where
1972+ K : PartialOrd ,
1973+ R : RangeBounds < K > ,
19561974 F : FnMut ( & K , & mut V ) -> bool ,
19571975{
19581976 type Item = ( K , V ) ;
@@ -1966,7 +1984,7 @@ where
19661984 }
19671985}
19681986
1969- impl < ' a , K , V > ExtractIfInner < ' a , K , V > {
1987+ impl < ' a , K , V , R > ExtractIfInner < ' a , K , V , R > {
19701988 /// Allow Debug implementations to predict the next element.
19711989 pub ( super ) fn peek ( & self ) -> Option < ( & K , & V ) > {
19721990 let edge = self . cur_leaf_edge . as_ref ( ) ?;
@@ -1976,10 +1994,22 @@ impl<'a, K, V> ExtractIfInner<'a, K, V> {
19761994 /// Implementation of a typical `ExtractIf::next` method, given the predicate.
19771995 pub ( super ) fn next < F , A : Allocator + Clone > ( & mut self , pred : & mut F , alloc : A ) -> Option < ( K , V ) >
19781996 where
1997+ K : PartialOrd ,
1998+ R : RangeBounds < K > ,
19791999 F : FnMut ( & K , & mut V ) -> bool ,
19802000 {
19812001 while let Ok ( mut kv) = self . cur_leaf_edge . take ( ) ?. next_kv ( ) {
19822002 let ( k, v) = kv. kv_mut ( ) ;
2003+
2004+ // On creation, we navigated directly to the left bound, so we need only check the
2005+ // right bound here to decide whether to stop.
2006+ match self . range . end_bound ( ) {
2007+ Bound :: Included ( ref end) if ( * k) . le ( end) => ( ) ,
2008+ Bound :: Excluded ( ref end) if ( * k) . lt ( end) => ( ) ,
2009+ Bound :: Unbounded => ( ) ,
2010+ _ => return None ,
2011+ }
2012+
19832013 if pred ( k, v) {
19842014 * self . length -= 1 ;
19852015 let ( kv, pos) = kv. remove_kv_tracking (
@@ -2011,7 +2041,13 @@ impl<'a, K, V> ExtractIfInner<'a, K, V> {
20112041}
20122042
20132043#[ unstable( feature = "btree_extract_if" , issue = "70530" ) ]
2014- impl < K , V , F > FusedIterator for ExtractIf < ' _ , K , V , F > where F : FnMut ( & K , & mut V ) -> bool { }
2044+ impl < K , V , R , F > FusedIterator for ExtractIf < ' _ , K , V , R , F >
2045+ where
2046+ K : PartialOrd ,
2047+ R : RangeBounds < K > ,
2048+ F : FnMut ( & K , & mut V ) -> bool ,
2049+ {
2050+ }
20152051
20162052#[ stable( feature = "btree_range" , since = "1.17.0" ) ]
20172053impl < ' a , K , V > Iterator for Range < ' a , K , V > {
0 commit comments