Auto merge of #81090 - ssomers:btree_drainy_refactor_2, r=Mark-Simulacrum

BTreeMap: offer merge in variants with more clarity

r? `@Mark-Simulacrum`

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

@@ -1282,25 +1282,25 @@ pub fn into_right_child(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrIn
         self.right_child
     }
 
-    /// Returns `true` if it is valid to call `.merge()` in the balancing context,
+    /// Returns whether merging is possible, i.e., whether there is enough room
-    /// i.e., whether there is enough room in a node to hold the combination of
+    /// in a node to combine the central KV with both adjacent child nodes.
-    /// both adjacent child nodes, along with the key-value pair in the parent.
     pub fn can_merge(&self) -> bool {
         self.left_child.len() + 1 + self.right_child.len() <= CAPACITY
     }
 }
 
 impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
-    /// Merges the parent's key-value pair and both adjacent child nodes into
+    /// Performs a merge and lets a closure decide what to return.
-    /// the left node and returns an edge handle in that expanded left node.
+    fn do_merge<
-    /// If `track_edge_idx` is given some value, the returned edge corresponds
+        F: FnOnce(
-    /// to where the edge in that child node ended up,
+            NodeRef<marker::Mut<'a>, K, V, marker::Internal>,
-    ///
+            NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
-    /// Panics unless we `.can_merge()`.
+        ) -> R,
-    pub fn merge(
+        R,
+    >(
         self,
-        track_edge_idx: Option<LeftOrRight<usize>>,
+        result: F,
-    ) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::Edge> {
+    ) -> R {
         let Handle { node: mut parent_node, idx: parent_idx, _marker } = self.parent;
         let old_parent_len = parent_node.len();
         let mut left_node = self.left_child;
@@ -1310,11 +1310,6 @@ pub fn merge(
         let new_left_len = old_left_len + 1 + right_len;
 
         assert!(new_left_len <= CAPACITY);
-        assert!(match track_edge_idx {
-            None => true,
-            Some(LeftOrRight::Left(idx)) => idx <= old_left_len,
-            Some(LeftOrRight::Right(idx)) => idx <= right_len,
-        });
 
         unsafe {
             *left_node.len_mut() = new_left_len as u16;
@@ -1353,14 +1348,47 @@ pub fn merge(
             } else {
                 Global.deallocate(right_node.node.cast(), Layout::new::<LeafNode<K, V>>());
             }
-
-            let new_idx = match track_edge_idx {
-                None => 0,
-                Some(LeftOrRight::Left(idx)) => idx,
-                Some(LeftOrRight::Right(idx)) => old_left_len + 1 + idx,
-            };
-            Handle::new_edge(left_node, new_idx)
         }
+        result(parent_node, left_node)
+    }
+
+    /// Merges the parent's key-value pair and both adjacent child nodes into
+    /// the left child node and returns the shrunk parent node.
+    ///
+    /// Panics unless we `.can_merge()`.
+    pub fn merge_tracking_parent(self) -> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
+        self.do_merge(|parent, _child| parent)
+    }
+
+    /// Merges the parent's key-value pair and both adjacent child nodes into
+    /// the left child node and returns that child node.
+    ///
+    /// Panics unless we `.can_merge()`.
+    pub fn merge_tracking_child(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
+        self.do_merge(|_parent, child| child)
+    }
+
+    /// Merges the parent's key-value pair and both adjacent child nodes into
+    /// the left child node and returns the edge handle in that child node
+    /// where the tracked child edge ended up,
+    ///
+    /// Panics unless we `.can_merge()`.
+    pub fn merge_tracking_child_edge(
+        self,
+        track_edge_idx: LeftOrRight<usize>,
+    ) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::Edge> {
+        let old_left_len = self.left_child.len();
+        let right_len = self.right_child.len();
+        assert!(match track_edge_idx {
+            LeftOrRight::Left(idx) => idx <= old_left_len,
+            LeftOrRight::Right(idx) => idx <= right_len,
+        });
+        let child = self.merge_tracking_child();
+        let new_idx = match track_edge_idx {
+            LeftOrRight::Left(idx) => idx,
+            LeftOrRight::Right(idx) => old_left_len + 1 + idx,
+        };
+        unsafe { Handle::new_edge(child, new_idx) }
     }
 
     /// Removes a key-value pair from the left child and places it in the key-value storage

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

@@ -33,7 +33,7 @@ fn remove_leaf_kv<F: FnOnce()>(
                 Ok(Left(left_parent_kv)) => {
                     debug_assert!(left_parent_kv.right_child_len() == MIN_LEN - 1);
                     if left_parent_kv.can_merge() {
-                        left_parent_kv.merge(Some(Right(idx)))
+                        left_parent_kv.merge_tracking_child_edge(Right(idx))
                     } else {
                         debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
                         left_parent_kv.steal_left(idx)
@@ -42,7 +42,7 @@ fn remove_leaf_kv<F: FnOnce()>(
                 Ok(Right(right_parent_kv)) => {
                     debug_assert!(right_parent_kv.left_child_len() == MIN_LEN - 1);
                     if right_parent_kv.can_merge() {
-                        right_parent_kv.merge(Some(Left(idx)))
+                        right_parent_kv.merge_tracking_child_edge(Left(idx))
                     } else {
                         debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
                         right_parent_kv.steal_right(idx)
@@ -124,9 +124,8 @@ fn handle_underfull_node_locally(
             Ok(Left(left_parent_kv)) => {
                 debug_assert_eq!(left_parent_kv.right_child_len(), MIN_LEN - 1);
                 if left_parent_kv.can_merge() {
-                    let pos = left_parent_kv.merge(None);
+                    let parent = left_parent_kv.merge_tracking_parent();
-                    let parent_edge = unsafe { unwrap_unchecked(pos.into_node().ascend().ok()) };
+                    Some(parent)
-                    Some(parent_edge.into_node())
                 } else {
                     debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
                     left_parent_kv.steal_left(0);
@@ -136,9 +135,8 @@ fn handle_underfull_node_locally(
             Ok(Right(right_parent_kv)) => {
                 debug_assert_eq!(right_parent_kv.left_child_len(), MIN_LEN - 1);
                 if right_parent_kv.can_merge() {
-                    let pos = right_parent_kv.merge(None);
+                    let parent = right_parent_kv.merge_tracking_parent();
-                    let parent_edge = unsafe { unwrap_unchecked(pos.into_node().ascend().ok()) };
+                    Some(parent)
-                    Some(parent_edge.into_node())
                 } else {
                     debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
                     right_parent_kv.steal_right(0);

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

@@ -66,7 +66,7 @@ fn fix_right_border(&mut self) {
                 let mut last_kv = node.last_kv().consider_for_balancing();
 
                 if last_kv.can_merge() {
-                    cur_node = last_kv.merge(None).into_node();
+                    cur_node = last_kv.merge_tracking_child();
                 } else {
                     let right_len = last_kv.right_child_len();
                     // `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
@@ -93,7 +93,7 @@ fn fix_left_border(&mut self) {
                 let mut first_kv = node.first_kv().consider_for_balancing();
 
                 if first_kv.can_merge() {
-                    cur_node = first_kv.merge(None).into_node();
+                    cur_node = first_kv.merge_tracking_child();
                 } else {
                     let left_len = first_kv.left_child_len();
                     // `MIN_LEN + 1` to avoid readjust if merge happens on the next level.