@@ -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