Rollup merge of #78056 - ssomers:btree_chop_up_1, r=dtolnay

BTreeMap: split off most code of remove and split_off

Putting map.rs on a diet, in addition to #77851.
r? @dtolnay

4 files changed, 239 insertions, 222 deletions

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index 92cbce96054..20c6ebd2292 100644
--- a/library/alloc/src/collections/btree/map.rs
+++ b/library/alloc/src/collections/btree/map.rs
@@ -13,8 +13,6 @@ use super::node::{self, marker, ForceResult::*, Handle, NodeRef};
 use super::search::{self, SearchResult::*};
 use super::unwrap_unchecked;
 
-use UnderflowResult::*;
-
 mod entry;
 pub use entry::{Entry, OccupiedEntry, VacantEntry};
 use Entry::*;
@@ -1154,40 +1152,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
 
         let mut right = Self::new();
         let right_root = Self::ensure_is_owned(&mut right.root);
-        for _ in 0..left_root.height() {
-            right_root.push_internal_level();
-        }
-
-        {
-            let mut left_node = left_root.node_as_mut();
-            let mut right_node = right_root.node_as_mut();
-
-            loop {
-                let mut split_edge = match search::search_node(left_node, key) {
-                    // key is going to the right tree
-                    Found(handle) => handle.left_edge(),
-                    GoDown(handle) => handle,
-                };
 
-                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();
+        left_root.split_off(right_root, key);
 
         if left_root.height() < right_root.height() {
             self.length = left_root.node_as_ref().calc_length();
@@ -2250,193 +2216,6 @@ impl<K, V> BTreeMap<K, V> {
     }
 }
 
-impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
-    /// Removes a key/value-pair from the map, and returns that pair, as well as
-    /// the leaf edge corresponding to that former pair.
-    fn remove_kv_tracking<F: FnOnce()>(
-        self,
-        handle_emptied_internal_root: F,
-    ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
-        let (old_kv, mut pos, was_internal) = match self.force() {
-            Leaf(leaf) => {
-                let (old_kv, pos) = leaf.remove();
-                (old_kv, pos, false)
-            }
-            Internal(mut internal) => {
-                // Replace the location freed in the internal node with an
-                // adjacent KV, and remove that adjacent KV from its leaf.
-                // Always choose the adjacent KV on the left side because
-                // it is typically faster to pop an element from the end
-                // of the KV arrays without needing to shift other elements.
-
-                let key_loc = internal.kv_mut().0 as *mut K;
-                let val_loc = internal.kv_mut().1 as *mut V;
-
-                let to_remove = internal.left_edge().descend().last_leaf_edge().left_kv().ok();
-                let to_remove = unsafe { unwrap_unchecked(to_remove) };
-
-                let (kv, pos) = to_remove.remove();
-
-                let old_key = unsafe { mem::replace(&mut *key_loc, kv.0) };
-                let old_val = unsafe { mem::replace(&mut *val_loc, kv.1) };
-
-                ((old_key, old_val), pos, true)
-            }
-        };
-
-        // Handle underflow
-        let mut cur_node = unsafe { ptr::read(&pos).into_node().forget_type() };
-        let mut at_leaf = true;
-        while cur_node.len() < node::MIN_LEN {
-            match handle_underfull_node(cur_node) {
-                AtRoot => break,
-                Merged(edge, merged_with_left, offset) => {
-                    // If we merged with our right sibling then our tracked
-                    // position has not changed. However if we merged with our
-                    // left sibling then our tracked position is now dangling.
-                    if at_leaf && merged_with_left {
-                        let idx = pos.idx() + offset;
-                        let node = match unsafe { ptr::read(&edge).descend().force() } {
-                            Leaf(leaf) => leaf,
-                            Internal(_) => unreachable!(),
-                        };
-                        pos = unsafe { Handle::new_edge(node, idx) };
-                    }
-
-                    let parent = edge.into_node();
-                    if parent.len() == 0 {
-                        // The parent that was just emptied must be the root,
-                        // because nodes on a lower level would not have been
-                        // left with a single child.
-                        handle_emptied_internal_root();
-                        break;
-                    } else {
-                        cur_node = parent.forget_type();
-                        at_leaf = false;
-                    }
-                }
-                Stole(stole_from_left) => {
-                    // Adjust the tracked position if we stole from a left sibling
-                    if stole_from_left && at_leaf {
-                        // SAFETY: This is safe since we just added an element to our node.
-                        unsafe {
-                            pos.move_next_unchecked();
-                        }
-                    }
-                    break;
-                }
-            }
-        }
-
-        // If we deleted from an internal node then we need to compensate for
-        // the earlier swap and adjust the tracked position to point to the
-        // next element.
-        if was_internal {
-            pos = unsafe { unwrap_unchecked(pos.next_kv().ok()).next_leaf_edge() };
-        }
-
-        (old_kv, pos)
-    }
-}
-
-impl<K, V> node::Root<K, V> {
-    /// Removes empty levels on the top, but keep an empty leaf if the entire tree is empty.
-    fn fix_top(&mut self) {
-        while self.height() > 0 && self.node_as_ref().len() == 0 {
-            self.pop_internal_level();
-        }
-    }
-
-    fn fix_right_border(&mut self) {
-        self.fix_top();
-
-        {
-            let mut cur_node = self.node_as_mut();
-
-            while let Internal(node) = cur_node.force() {
-                let mut last_kv = node.last_kv();
-
-                if last_kv.can_merge() {
-                    cur_node = last_kv.merge().descend();
-                } else {
-                    let right_len = last_kv.reborrow().right_edge().descend().len();
-                    // `MINLEN + 1` to avoid readjust if merge happens on the next level.
-                    if right_len < node::MIN_LEN + 1 {
-                        last_kv.bulk_steal_left(node::MIN_LEN + 1 - right_len);
-                    }
-                    cur_node = last_kv.right_edge().descend();
-                }
-            }
-        }
-
-        self.fix_top();
-    }
-
-    /// The symmetric clone of `fix_right_border`.
-    fn fix_left_border(&mut self) {
-        self.fix_top();
-
-        {
-            let mut cur_node = self.node_as_mut();
-
-            while let Internal(node) = cur_node.force() {
-                let mut first_kv = node.first_kv();
-
-                if first_kv.can_merge() {
-                    cur_node = first_kv.merge().descend();
-                } else {
-                    let left_len = first_kv.reborrow().left_edge().descend().len();
-                    if left_len < node::MIN_LEN + 1 {
-                        first_kv.bulk_steal_right(node::MIN_LEN + 1 - left_len);
-                    }
-                    cur_node = first_kv.left_edge().descend();
-                }
-            }
-        }
-
-        self.fix_top();
-    }
-}
-
-enum UnderflowResult<'a, K, V> {
-    AtRoot,
-    Merged(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge>, bool, usize),
-    Stole(bool),
-}
-
-fn handle_underfull_node<'a, K: 'a, V: 'a>(
-    node: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
-) -> UnderflowResult<'_, K, V> {
-    let parent = match node.ascend() {
-        Ok(parent) => parent,
-        Err(_) => return AtRoot,
-    };
-
-    // Prefer the left KV if it exists. Merging with the left side is faster,
-    // since merging happens towards the left and `node` has fewer elements.
-    // Stealing from the left side is faster, since we can pop from the end of
-    // the KV arrays.
-    let (is_left, mut handle) = match parent.left_kv() {
-        Ok(left) => (true, left),
-        Err(parent) => {
-            let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) };
-            (false, right)
-        }
-    };
-
-    if handle.can_merge() {
-        let offset = if is_left { handle.reborrow().left_edge().descend().len() + 1 } else { 0 };
-        Merged(handle.merge(), is_left, offset)
-    } else {
-        if is_left {
-            handle.steal_left();
-        } else {
-            handle.steal_right();
-        }
-        Stole(is_left)
-    }
-}
-
 impl<K: Ord, V, I: Iterator<Item = (K, V)>> Iterator for MergeIter<K, V, I> {
     type Item = (K, V);
 
diff --git a/library/alloc/src/collections/btree/mod.rs b/library/alloc/src/collections/btree/mod.rs
index ecbdacda4b6..bcc50ed5615 100644
--- a/library/alloc/src/collections/btree/mod.rs
+++ b/library/alloc/src/collections/btree/mod.rs
@@ -2,8 +2,10 @@ mod borrow;
 pub mod map;
 mod navigate;
 mod node;
+mod remove;
 mod search;
 pub mod set;
+mod split;
 
 #[doc(hidden)]
 trait Recover<Q: ?Sized> {
diff --git a/library/alloc/src/collections/btree/remove.rs b/library/alloc/src/collections/btree/remove.rs
new file mode 100644
index 00000000000..9733b7d6084
--- /dev/null
+++ b/library/alloc/src/collections/btree/remove.rs
@@ -0,0 +1,132 @@
+use super::node::{self, marker, ForceResult, Handle, NodeRef};
+use super::unwrap_unchecked;
+use core::mem;
+use core::ptr;
+
+impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
+    /// Removes a key/value-pair from the map, and returns that pair, as well as
+    /// the leaf edge corresponding to that former pair.
+    pub fn remove_kv_tracking<F: FnOnce()>(
+        self,
+        handle_emptied_internal_root: F,
+    ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
+        let (old_kv, mut pos, was_internal) = match self.force() {
+            ForceResult::Leaf(leaf) => {
+                let (old_kv, pos) = leaf.remove();
+                (old_kv, pos, false)
+            }
+            ForceResult::Internal(mut internal) => {
+                // Replace the location freed in the internal node with an
+                // adjacent KV, and remove that adjacent KV from its leaf.
+                // Always choose the adjacent KV on the left side because
+                // it is typically faster to pop an element from the end
+                // of the KV arrays without needing to shift other elements.
+
+                let key_loc = internal.kv_mut().0 as *mut K;
+                let val_loc = internal.kv_mut().1 as *mut V;
+
+                let to_remove = internal.left_edge().descend().last_leaf_edge().left_kv().ok();
+                let to_remove = unsafe { unwrap_unchecked(to_remove) };
+
+                let (kv, pos) = to_remove.remove();
+
+                let old_key = unsafe { mem::replace(&mut *key_loc, kv.0) };
+                let old_val = unsafe { mem::replace(&mut *val_loc, kv.1) };
+
+                ((old_key, old_val), pos, true)
+            }
+        };
+
+        // Handle underflow
+        let mut cur_node = unsafe { ptr::read(&pos).into_node().forget_type() };
+        let mut at_leaf = true;
+        while cur_node.len() < node::MIN_LEN {
+            match handle_underfull_node(cur_node) {
+                UnderflowResult::AtRoot => break,
+                UnderflowResult::Merged(edge, merged_with_left, offset) => {
+                    // If we merged with our right sibling then our tracked
+                    // position has not changed. However if we merged with our
+                    // left sibling then our tracked position is now dangling.
+                    if at_leaf && merged_with_left {
+                        let idx = pos.idx() + offset;
+                        let node = match unsafe { ptr::read(&edge).descend().force() } {
+                            ForceResult::Leaf(leaf) => leaf,
+                            ForceResult::Internal(_) => unreachable!(),
+                        };
+                        pos = unsafe { Handle::new_edge(node, idx) };
+                    }
+
+                    let parent = edge.into_node();
+                    if parent.len() == 0 {
+                        // The parent that was just emptied must be the root,
+                        // because nodes on a lower level would not have been
+                        // left with a single child.
+                        handle_emptied_internal_root();
+                        break;
+                    } else {
+                        cur_node = parent.forget_type();
+                        at_leaf = false;
+                    }
+                }
+                UnderflowResult::Stole(stole_from_left) => {
+                    // Adjust the tracked position if we stole from a left sibling
+                    if stole_from_left && at_leaf {
+                        // SAFETY: This is safe since we just added an element to our node.
+                        unsafe {
+                            pos.move_next_unchecked();
+                        }
+                    }
+                    break;
+                }
+            }
+        }
+
+        // If we deleted from an internal node then we need to compensate for
+        // the earlier swap and adjust the tracked position to point to the
+        // next element.
+        if was_internal {
+            pos = unsafe { unwrap_unchecked(pos.next_kv().ok()).next_leaf_edge() };
+        }
+
+        (old_kv, pos)
+    }
+}
+
+enum UnderflowResult<'a, K, V> {
+    AtRoot,
+    Merged(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge>, bool, usize),
+    Stole(bool),
+}
+
+fn handle_underfull_node<'a, K: 'a, V: 'a>(
+    node: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
+) -> UnderflowResult<'_, K, V> {
+    let parent = match node.ascend() {
+        Ok(parent) => parent,
+        Err(_) => return UnderflowResult::AtRoot,
+    };
+
+    // Prefer the left KV if it exists. Merging with the left side is faster,
+    // since merging happens towards the left and `node` has fewer elements.
+    // Stealing from the left side is faster, since we can pop from the end of
+    // the KV arrays.
+    let (is_left, mut handle) = match parent.left_kv() {
+        Ok(left) => (true, left),
+        Err(parent) => {
+            let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) };
+            (false, right)
+        }
+    };
+
+    if handle.can_merge() {
+        let offset = if is_left { handle.reborrow().left_edge().descend().len() + 1 } else { 0 };
+        UnderflowResult::Merged(handle.merge(), is_left, offset)
+    } else {
+        if is_left {
+            handle.steal_left();
+        } else {
+            handle.steal_right();
+        }
+        UnderflowResult::Stole(is_left)
+    }
+}
diff --git a/library/alloc/src/collections/btree/split.rs b/library/alloc/src/collections/btree/split.rs
new file mode 100644
index 00000000000..0e6e213f6e8
--- /dev/null
+++ b/library/alloc/src/collections/btree/split.rs
@@ -0,0 +1,104 @@
+use super::node::{self, ForceResult::*, Root};
+use super::search::{self, 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)
+    where
+        K: Borrow<Q>,
+    {
+        debug_assert!(right_root.height() == 0);
+        debug_assert!(right_root.node_as_ref().len() == 0);
+
+        let left_root = self;
+        for _ in 0..left_root.height() {
+            right_root.push_internal_level();
+        }
+
+        {
+            let mut left_node = left_root.node_as_mut();
+            let mut right_node = right_root.node_as_mut();
+
+            loop {
+                let mut split_edge = match search::search_node(left_node, key) {
+                    // key is going to the right tree
+                    Found(handle) => handle.left_edge(),
+                    GoDown(handle) => handle,
+                };
+
+                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();
+    }
+
+    /// Removes empty levels on the top, but keeps an empty leaf if the entire tree is empty.
+    fn fix_top(&mut self) {
+        while self.height() > 0 && self.node_as_ref().len() == 0 {
+            self.pop_internal_level();
+        }
+    }
+
+    fn fix_right_border(&mut self) {
+        self.fix_top();
+
+        {
+            let mut cur_node = self.node_as_mut();
+
+            while let Internal(node) = cur_node.force() {
+                let mut last_kv = node.last_kv();
+
+                if last_kv.can_merge() {
+                    cur_node = last_kv.merge().descend();
+                } else {
+                    let right_len = last_kv.reborrow().right_edge().descend().len();
+                    // `MINLEN + 1` to avoid readjust if merge happens on the next level.
+                    if right_len < node::MIN_LEN + 1 {
+                        last_kv.bulk_steal_left(node::MIN_LEN + 1 - right_len);
+                    }
+                    cur_node = last_kv.right_edge().descend();
+                }
+            }
+        }
+
+        self.fix_top();
+    }
+
+    /// The symmetric clone of `fix_right_border`.
+    fn fix_left_border(&mut self) {
+        self.fix_top();
+
+        {
+            let mut cur_node = self.node_as_mut();
+
+            while let Internal(node) = cur_node.force() {
+                let mut first_kv = node.first_kv();
+
+                if first_kv.can_merge() {
+                    cur_node = first_kv.merge().descend();
+                } else {
+                    let left_len = first_kv.reborrow().left_edge().descend().len();
+                    if left_len < node::MIN_LEN + 1 {
+                        first_kv.bulk_steal_right(node::MIN_LEN + 1 - left_len);
+                    }
+                    cur_node = first_kv.left_edge().descend();
+                }
+            }
+        }
+
+        self.fix_top();
+    }
+}