BTree: add split_off_range methods

Author: ssomers

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

@@ -1204,6 +1204,66 @@ impl<K, V> BTreeMap<K, V> {
         BTreeMap { root: Some(right_root), length: right_len }
     }
 
+    /// Splits the collection into two. Returns a new collection with all keys in the given range.
+    ///
+    /// # 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`.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![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 = a.split_off_range(&3..&33);
+    ///
+    /// assert_eq!(a.len(), 3);
+    /// assert_eq!(b.len(), 2);
+    ///
+    /// assert_eq!(a[&1], "a");
+    /// assert_eq!(a[&2], "b");
+    /// assert_eq!(a[&41], "e");
+    ///
+    /// assert_eq!(b[&3], "c");
+    /// assert_eq!(b[&17], "d");
+    /// ```
+    #[unstable(feature = "btree_drain", issue = "81074")]
+    pub fn split_off_range<T: ?Sized, R>(&mut self, range: R) -> Self
+    where
+        T: Ord,
+        K: Borrow<T> + Ord,
+        R: RangeBounds<T>,
+    {
+        if self.is_empty() {
+            return Self::new();
+        }
+
+        let total_num = self.length;
+        let left_root = self.root.as_mut().unwrap(); // unwrap succeeds because not empty
+
+        let mut right_root = left_root.split_off_range(range);
+        right_root.fix_both_borders();
+
+        let (new_left_len, right_len) = Root::calc_split_length(total_num, &left_root, &right_root);
+        self.length = new_left_len;
+
+        BTreeMap { root: Some(right_root), length: right_len }
+    }
+
     /// 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

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

@@ -17,6 +17,8 @@ use std::ops::RangeBounds;
 use std::panic::{catch_unwind, AssertUnwindSafe};
 use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
 
+mod split_off_range;
+
 // Capacity of a tree with a single level,
 // i.e., a tree who's root is a leaf node at height 0.
 const NODE_CAPACITY: usize = node::CAPACITY;
@@ -31,6 +33,12 @@ const MIN_INSERTS_HEIGHT_1: usize = NODE_CAPACITY + 1;
 // It's not the minimum size: removing an element from such a tree does not always reduce height.
 const MIN_INSERTS_HEIGHT_2: usize = 89;
 
+// Like MIN_INSERTS_HEIGHT_2, with an additional internal level.
+const MIN_INSERTS_HEIGHT_3: usize = 628;
+
+// Like MIN_INSERTS_HEIGHT_3, with an additional internal level.
+const MIN_INSERTS_HEIGHT_4: usize = 4401;
+
 // Gathers all references from a mutable iterator and makes sure Miri notices if
 // using them is dangerous.
 fn test_all_refs<'a, T: 'a>(dummy: &mut T, iter: impl Iterator<Item = &'a mut T>) {
@@ -173,6 +181,26 @@ fn test_levels() {
     // - 5 elements in right child's last grandchild
     assert_eq!(map.height(), Some(2));
     assert_eq!(map.len(), MIN_INSERTS_HEIGHT_2, "{}", map.dump_keys());
+
+    if cfg!(miri) {
+        // Miri is too slow
+        return;
+    }
+    while map.height() == Some(2) {
+        let last_key = *map.last_key_value().unwrap().0;
+        map.insert(last_key + 1, ());
+    }
+    map.check();
+    assert_eq!(map.height(), Some(3));
+    assert_eq!(map.len(), MIN_INSERTS_HEIGHT_3);
+
+    while map.height() == Some(3) {
+        let last_key = *map.last_key_value().unwrap().0;
+        map.insert(last_key + 1, ());
+    }
+    map.check();
+    assert_eq!(map.height(), Some(4));
+    assert_eq!(map.len(), MIN_INSERTS_HEIGHT_4);
 }
 
 // Ensures the testing infrastructure usually notices order violations.