BTree: add drain methods

Author: ssomers

library/alloc/benches/btree/map.rs

@@ -384,17 +384,33 @@ pub fn clone_slim_100_and_clear(b: &mut Bencher) {
 }
 
 #[bench]
-pub fn clone_slim_100_and_drain_all(b: &mut Bencher) {
+pub fn clone_slim_100_and_drain_range_all(b: &mut Bencher) {
+    let src = slim_map(100);
+    b.iter(|| src.clone().drain(..).count())
+}
+
+#[bench]
+pub fn clone_slim_100_and_drain_filter_all(b: &mut Bencher) {
     let src = slim_map(100);
     b.iter(|| src.clone().drain_filter(|_, _| true).count())
 }
 
 #[bench]
-pub fn clone_slim_100_and_drain_half(b: &mut Bencher) {
+pub fn clone_slim_100_and_drain_range_half(b: &mut Bencher) {
+    let src = slim_map(100);
+    b.iter(|| {
+        let mut map = src.clone();
+        assert_eq!(map.drain(25..75).count(), 100 / 2);
+        assert_eq!(map.len(), 100 / 2);
+    })
+}
+
+#[bench]
+pub fn clone_slim_100_and_drain_filter_half(b: &mut Bencher) {
     let src = slim_map(100);
     b.iter(|| {
         let mut map = src.clone();
-        assert_eq!(map.drain_filter(|i, _| i % 2 == 0).count(), 100 / 2);
+        assert_eq!(map.drain_filter(|i, _| (25..75).contains(i)).count(), 100 / 2);
         assert_eq!(map.len(), 100 / 2);
     })
 }
@@ -455,17 +471,33 @@ pub fn clone_slim_10k_and_clear(b: &mut Bencher) {
 }
 
 #[bench]
-pub fn clone_slim_10k_and_drain_all(b: &mut Bencher) {
+pub fn clone_slim_10k_and_drain_range_all(b: &mut Bencher) {
+    let src = slim_map(10_000);
+    b.iter(|| src.clone().drain(..).count())
+}
+
+#[bench]
+pub fn clone_slim_10k_and_drain_filter_all(b: &mut Bencher) {
     let src = slim_map(10_000);
     b.iter(|| src.clone().drain_filter(|_, _| true).count())
 }
 
 #[bench]
-pub fn clone_slim_10k_and_drain_half(b: &mut Bencher) {
+pub fn clone_slim_10k_and_drain_range_half(b: &mut Bencher) {
+    let src = slim_map(10_000);
+    b.iter(|| {
+        let mut map = src.clone();
+        assert_eq!(map.drain(2500..7500).count(), 10_000 / 2);
+        assert_eq!(map.len(), 10_000 / 2);
+    })
+}
+
+#[bench]
+pub fn clone_slim_10k_and_drain_filter_half(b: &mut Bencher) {
     let src = slim_map(10_000);
     b.iter(|| {
         let mut map = src.clone();
-        assert_eq!(map.drain_filter(|i, _| i % 2 == 0).count(), 10_000 / 2);
+        assert_eq!(map.drain_filter(|i, _| (2500..7500).contains(i)).count(), 10_000 / 2);
         assert_eq!(map.len(), 10_000 / 2);
     })
 }
@@ -526,17 +558,33 @@ pub fn clone_fat_val_100_and_clear(b: &mut Bencher) {
 }
 
 #[bench]
-pub fn clone_fat_val_100_and_drain_all(b: &mut Bencher) {
+pub fn clone_fat_val_100_and_drain_range_all(b: &mut Bencher) {
+    let src = fat_val_map(100);
+    b.iter(|| src.clone().drain(..).count())
+}
+
+#[bench]
+pub fn clone_fat_val_100_and_drain_filter_all(b: &mut Bencher) {
     let src = fat_val_map(100);
     b.iter(|| src.clone().drain_filter(|_, _| true).count())
 }
 
 #[bench]
-pub fn clone_fat_val_100_and_drain_half(b: &mut Bencher) {
+pub fn clone_fat_val_100_and_drain_range_half(b: &mut Bencher) {
+    let src = fat_val_map(100);
+    b.iter(|| {
+        let mut map = src.clone();
+        assert_eq!(map.drain(25..75).count(), 100 / 2);
+        assert_eq!(map.len(), 100 / 2);
+    })
+}
+
+#[bench]
+pub fn clone_fat_val_100_and_drain_filter_half(b: &mut Bencher) {
     let src = fat_val_map(100);
     b.iter(|| {
         let mut map = src.clone();
-        assert_eq!(map.drain_filter(|i, _| i % 2 == 0).count(), 100 / 2);
+        assert_eq!(map.drain_filter(|i, _| (25..75).contains(i)).count(), 100 / 2);
         assert_eq!(map.len(), 100 / 2);
     })
 }

library/alloc/benches/lib.rs

@@ -1,6 +1,7 @@
 // Disabling on android for the time being
 // See https://github.com/rust-lang/rust/issues/73535#event-3477699747
 #![cfg(not(target_os = "android"))]
+#![feature(btree_drain)]
 #![feature(btree_drain_filter)]
 #![feature(map_first_last)]
 #![feature(repr_simd)]

library/alloc/src/collections/btree/map.rs

@@ -1264,6 +1264,66 @@ impl<K, V> BTreeMap<K, V> {
         BTreeMap { root: Some(right_root), length: right_len }
     }
 
+    /// Removes at once all keys within the range from the map, returning the
+    /// removed key-value pairs as an iterator in ascending key order. If the
+    /// iterator is dropped before being fully consumed, it drops the remaining
+    /// removed key-value pairs.
+    ///
+    /// The returned iterator keeps a mutable borrow on the map to allow
+    /// optimizing its implementation.
+    ///
+    /// # Panics
+    ///
+    /// Panics if range `start > end`.
+    /// Panics if range `start == end` and both bounds are `Excluded`.
+    /// May panic if the [`Ord`] implementation of type `T` is ill-defined,
+    /// either because it does not form a total order or because it does not
+    /// correspond to the [`Ord`] implementation of type `K`.
+    ///
+    /// # Leaking
+    ///
+    /// If the returned iterator goes out of scope without being dropped (due to
+    /// [`mem::forget`], for example), the map may have lost and leaked
+    /// key-value pairs arbitrarily, including key-value pairs outside the range.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(btree_drain)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut a = BTreeMap::new();
+    /// a.insert(1, "a");
+    /// a.insert(2, "b");
+    /// a.insert(3, "c");
+    /// a.insert(17, "d");
+    /// a.insert(41, "e");
+    ///
+    /// let b: Vec<_> = a.drain(3..33).collect();
+    /// assert_eq!(b, vec![(3, "c"), (17, "d")]);
+    /// assert_eq!(a.len(), 3);
+    /// ```
+    #[unstable(feature = "btree_drain", issue = "81074")]
+    pub fn drain<T: ?Sized, R>(&mut self, range: R) -> Drain<'_, K, V>
+    where
+        T: Ord,
+        K: Borrow<T> + Ord,
+        R: RangeBounds<T>,
+    {
+        let inner = if let Some(left_root) = self.root.as_mut() {
+            let total_num = self.length;
+            let right_root = left_root.split_off_range(range);
+            let (new_left_len, right_len) =
+                Root::calc_split_length(total_num, &left_root, &right_root);
+            self.length = new_left_len;
+            let right_range = right_root.into_dying().full_range();
+            IntoIter { range: right_range, length: right_len }
+        } else {
+            IntoIter { range: LazyLeafRange::none(), length: 0 }
+        };
+        Drain { inner, _marker: PhantomData }
+    }
+
     /// Creates an iterator that visits all elements (key-value pairs) in
     /// ascending key order and uses a closure to determine if an element should
     /// be removed. If the closure returns `true`, the element is removed from
@@ -1712,6 +1772,37 @@ impl<K, V> Clone for Values<'_, K, V> {
     }
 }
 
+/// An iterator produced by calling `drain` on BTreeMap.
+#[unstable(feature = "btree_drain", issue = "81074")]
+#[derive(Debug)]
+pub struct Drain<'a, K, V> {
+    inner: IntoIter<K, V>,
+    _marker: PhantomData<&'a mut BTreeMap<K, V>>,
+}
+
+#[unstable(feature = "btree_drain", issue = "81074")]
+impl<K, V> Iterator for Drain<'_, K, V> {
+    type Item = (K, V);
+
+    fn next(&mut self) -> Option<(K, V)> {
+        self.inner.next()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.inner.size_hint()
+    }
+}
+
+#[stable(feature = "rust1", since = "1.0.0")]
+impl<K, V> ExactSizeIterator for Drain<'_, K, V> {
+    fn len(&self) -> usize {
+        self.inner.len()
+    }
+}
+
+#[stable(feature = "fused", since = "1.26.0")]
+impl<K, V> FusedIterator for Drain<'_, K, V> {}
+
 /// An iterator produced by calling `drain_filter` on BTreeMap.
 #[unstable(feature = "btree_drain_filter", issue = "70530")]
 pub struct DrainFilter<'a, K, V, F>

library/alloc/src/collections/btree/map/tests.rs

@@ -17,6 +17,7 @@ use std::ops::RangeBounds;
 use std::panic::{catch_unwind, AssertUnwindSafe};
 use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
 
+mod drain;
 mod split_off_range;
 
 // Minimum number of elements to insert, to guarantee a tree with 2 levels,