@@ -1151,7 +1151,7 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
1151
1151
K : Ord ,
1152
1152
F : FnMut ( & K , & mut V ) -> bool ,
1153
1153
{
1154
- self . extract_if ( |k, v| !f ( k, v) ) . for_each ( drop) ;
1154
+ self . extract_if ( .. , |k, v| !f ( k, v) ) . for_each ( drop) ;
1155
1155
}
1156
1156
1157
1157
/// Moves all elements from `other` into `self`, leaving `other` empty.
@@ -1397,11 +1397,13 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
1397
1397
}
1398
1398
}
1399
1399
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.
1405
1407
///
1406
1408
/// The iterator also lets you mutate the value of each element in the
1407
1409
/// 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> {
1421
1423
/// use std::collections::BTreeMap;
1422
1424
///
1423
1425
/// 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();
1425
1427
/// let odds = map;
1426
1428
/// assert_eq!(evens.keys().copied().collect::<Vec<_>>(), [0, 2, 4, 6]);
1427
1429
/// 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]);
1428
1436
/// ```
1429
1437
#[ 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 >
1431
1439
where
1432
1440
K : Ord ,
1441
+ R : RangeBounds < K > ,
1433
1442
F : FnMut ( & K , & mut V ) -> bool ,
1434
1443
{
1435
- let ( inner, alloc) = self . extract_if_inner ( ) ;
1444
+ let ( inner, alloc) = self . extract_if_inner ( range ) ;
1436
1445
ExtractIf { pred, inner, alloc }
1437
1446
}
1438
1447
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 )
1440
1449
where
1441
1450
K : Ord ,
1451
+ R : RangeBounds < K > ,
1442
1452
{
1443
1453
if let Some ( root) = self . root . as_mut ( ) {
1444
1454
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 ( ) ) ) ;
1446
1456
(
1447
1457
ExtractIfInner {
1448
1458
length : & mut self . length ,
1449
1459
dormant_root : Some ( dormant_root) ,
1450
- cur_leaf_edge : Some ( front) ,
1460
+ cur_leaf_edge : Some ( first) ,
1461
+ range,
1451
1462
} ,
1452
1463
( * self . alloc ) . clone ( ) ,
1453
1464
)
@@ -1457,6 +1468,7 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
1457
1468
length : & mut self . length ,
1458
1469
dormant_root : None ,
1459
1470
cur_leaf_edge : None ,
1471
+ range,
1460
1472
} ,
1461
1473
( * self . alloc ) . clone ( ) ,
1462
1474
)
@@ -1915,18 +1927,19 @@ pub struct ExtractIf<
1915
1927
' a ,
1916
1928
K ,
1917
1929
V ,
1930
+ R ,
1918
1931
F ,
1919
1932
#[ unstable( feature = "allocator_api" , issue = "32838" ) ] A : Allocator + Clone = Global ,
1920
1933
> {
1921
1934
pred : F ,
1922
- inner : ExtractIfInner < ' a , K , V > ,
1935
+ inner : ExtractIfInner < ' a , K , V , R > ,
1923
1936
/// The BTreeMap will outlive this IntoIter so we don't care about drop order for `alloc`.
1924
1937
alloc : A ,
1925
1938
}
1926
1939
1927
1940
/// Most of the implementation of ExtractIf are generic over the type
1928
1941
/// 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 > {
1930
1943
/// Reference to the length field in the borrowed map, updated live.
1931
1944
length : & ' a mut usize ,
1932
1945
/// Buried reference to the root field in the borrowed map.
@@ -1936,10 +1949,13 @@ pub(super) struct ExtractIfInner<'a, K, V> {
1936
1949
/// Empty if the map has no root, if iteration went beyond the last leaf edge,
1937
1950
/// or if a panic occurred in the predicate.
1938
1951
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 ,
1939
1955
}
1940
1956
1941
1957
#[ 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 >
1943
1959
where
1944
1960
K : fmt:: Debug ,
1945
1961
V : fmt:: Debug ,
@@ -1951,8 +1967,10 @@ where
1951
1967
}
1952
1968
1953
1969
#[ 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 >
1955
1971
where
1972
+ K : PartialOrd ,
1973
+ R : RangeBounds < K > ,
1956
1974
F : FnMut ( & K , & mut V ) -> bool ,
1957
1975
{
1958
1976
type Item = ( K , V ) ;
@@ -1966,7 +1984,7 @@ where
1966
1984
}
1967
1985
}
1968
1986
1969
- impl < ' a , K , V > ExtractIfInner < ' a , K , V > {
1987
+ impl < ' a , K , V , R > ExtractIfInner < ' a , K , V , R > {
1970
1988
/// Allow Debug implementations to predict the next element.
1971
1989
pub ( super ) fn peek ( & self ) -> Option < ( & K , & V ) > {
1972
1990
let edge = self . cur_leaf_edge . as_ref ( ) ?;
@@ -1976,10 +1994,22 @@ impl<'a, K, V> ExtractIfInner<'a, K, V> {
1976
1994
/// Implementation of a typical `ExtractIf::next` method, given the predicate.
1977
1995
pub ( super ) fn next < F , A : Allocator + Clone > ( & mut self , pred : & mut F , alloc : A ) -> Option < ( K , V ) >
1978
1996
where
1997
+ K : PartialOrd ,
1998
+ R : RangeBounds < K > ,
1979
1999
F : FnMut ( & K , & mut V ) -> bool ,
1980
2000
{
1981
2001
while let Ok ( mut kv) = self . cur_leaf_edge . take ( ) ?. next_kv ( ) {
1982
2002
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
+
1983
2013
if pred ( k, v) {
1984
2014
* self . length -= 1 ;
1985
2015
let ( kv, pos) = kv. remove_kv_tracking (
@@ -2011,7 +2041,13 @@ impl<'a, K, V> ExtractIfInner<'a, K, V> {
2011
2041
}
2012
2042
2013
2043
#[ 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
+ }
2015
2051
2016
2052
#[ stable( feature = "btree_range" , since = "1.17.0" ) ]
2017
2053
impl < ' a , K , V > Iterator for Range < ' a , K , V > {
0 commit comments