BTreeMap: lightly refactor the split_off implementation

2 files changed, 67 insertions, 42 deletions

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index ecc28731874..1a62fc682e6 100644
--- a/library/alloc/src/collections/btree/map.rs
+++ b/library/alloc/src/collections/btree/map.rs
@@ -21,6 +21,14 @@ use Entry::*;
 /// We might temporarily have fewer elements during methods.
 pub(super) const MIN_LEN: usize = node::MIN_LEN_AFTER_SPLIT;
 
+// A tree in a `BTreeMap` is a tree in the `node` module with addtional invariants:
+// - Keys must appear in ascending order (according to the key's type).
+// - If the root node is internal, it must contain at least 1 element.
+// - Every non-root node contains at least MIN_LEN elements.
+//
+// An empty map may be represented both by the absense of a root node or by a
+// root node that is an empty leaf.
+
 /// A map based on a B-Tree.
 ///
 /// B-Trees represent a fundamental compromise between cache-efficiency and actually minimizing
@@ -1100,20 +1108,12 @@ impl<K: Ord, V> BTreeMap<K, V> {
         let total_num = self.len();
         let left_root = self.root.as_mut().unwrap(); // unwrap succeeds because not empty
 
-        let mut right = Self::new();
-        let right_root = Self::ensure_is_owned(&mut right.root);
-
-        left_root.split_off(right_root, key);
+        let right_root = left_root.split_off(key);
 
-        if left_root.height() < right_root.height() {
-            self.length = left_root.reborrow().calc_length();
-            right.length = total_num - self.len();
-        } else {
-            right.length = right_root.reborrow().calc_length();
-            self.length = total_num - right.len();
-        }
+        let (new_left_len, right_len) = Root::calc_split_length(total_num, &left_root, &right_root);
+        self.length = new_left_len;
 
-        right
+        BTreeMap { root: Some(right_root), length: right_len }
     }
 
     /// Creates an iterator which uses a closure to determine if an element should be removed.
diff --git a/library/alloc/src/collections/btree/split.rs b/library/alloc/src/collections/btree/split.rs
index 62c5e3a46d7..1921982464a 100644
--- a/library/alloc/src/collections/btree/split.rs
+++ b/library/alloc/src/collections/btree/split.rs
@@ -4,46 +4,71 @@ use super::search::SearchResult::*;
 use core::borrow::Borrow;
 
 impl<K, V> Root<K, V> {
-    pub fn split_off<Q: ?Sized + Ord>(&mut self, right_root: &mut Self, key: &Q)
+    /// Calculates the length of both trees that result from splitting up
+    /// a given number of distinct key-value pairs.
+    pub fn calc_split_length(
+        total_num: usize,
+        root_a: &Root<K, V>,
+        root_b: &Root<K, V>,
+    ) -> (usize, usize) {
+        let (length_a, length_b);
+        if root_a.height() < root_b.height() {
+            length_a = root_a.reborrow().calc_length();
+            length_b = total_num - length_a;
+            debug_assert_eq!(length_b, root_b.reborrow().calc_length());
+        } else {
+            length_b = root_b.reborrow().calc_length();
+            length_a = total_num - length_b;
+            debug_assert_eq!(length_a, root_a.reborrow().calc_length());
+        }
+        (length_a, length_b)
+    }
+
+    /// Split off a tree with key-value pairs at and after the given key.
+    /// The result is meaningful only if the tree is ordered by key,
+    /// and if the ordering of `Q` corresponds to that of `K`.
+    /// If `self` respects all `BTreeMap` tree invariants, then both
+    /// `self` and the returned tree will respect those invariants.
+    pub fn split_off<Q: ?Sized + Ord>(&mut self, key: &Q) -> Self
     where
         K: Borrow<Q>,
     {
-        debug_assert!(right_root.height() == 0);
-        debug_assert!(right_root.len() == 0);
-
         let left_root = self;
-        for _ in 0..left_root.height() {
-            right_root.push_internal_level();
-        }
-
-        {
-            let mut left_node = left_root.borrow_mut();
-            let mut right_node = right_root.borrow_mut();
-
-            loop {
-                let mut split_edge = match left_node.search_node(key) {
-                    // key is going to the right tree
-                    Found(kv) => kv.left_edge(),
-                    GoDown(edge) => edge,
-                };
-
-                split_edge.move_suffix(&mut right_node);
-
-                match (split_edge.force(), right_node.force()) {
-                    (Internal(edge), Internal(node)) => {
-                        left_node = edge.descend();
-                        right_node = node.first_edge().descend();
-                    }
-                    (Leaf(_), Leaf(_)) => {
-                        break;
-                    }
-                    _ => unreachable!(),
+        let mut right_root = Root::new_pillar(left_root.height());
+        let mut left_node = left_root.borrow_mut();
+        let mut right_node = right_root.borrow_mut();
+
+        loop {
+            let mut split_edge = match left_node.search_node(key) {
+                // key is going to the right tree
+                Found(kv) => kv.left_edge(),
+                GoDown(edge) => edge,
+            };
+
+            split_edge.move_suffix(&mut right_node);
+
+            match (split_edge.force(), right_node.force()) {
+                (Internal(edge), Internal(node)) => {
+                    left_node = edge.descend();
+                    right_node = node.first_edge().descend();
                 }
+                (Leaf(_), Leaf(_)) => break,
+                _ => unreachable!(),
             }
         }
 
         left_root.fix_right_border();
         right_root.fix_left_border();
+        right_root
+    }
+
+    /// Creates a tree consisting of empty nodes.
+    fn new_pillar(height: usize) -> Self {
+        let mut root = Root::new();
+        for _ in 0..height {
+            root.push_internal_level();
+        }
+        root
     }
 
     /// Removes empty levels on the top, but keeps an empty leaf if the entire tree is empty.