Rollup merge of #75519 - ssomers:btree_splitpoint_cleanup, r=Mark-Simulacrum

[rust.git] / library / alloc / src / collections / btree / node.rs

diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
index 5a6a71c5fb687a0725fad418f737f499eb2a0c90..acc2ae73572ba757aa801bafde1ae2e376efc2b5 100644 (file)
--- a/library/alloc/src/collections/btree/node.rs
+++ b/library/alloc/src/collections/btree/node.rs
@@ -416,7 +416,7 @@ pub unsafe fn deallocate_and_ascend(
 impl<'a, K, V, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
     /// Unsafely asserts to the compiler some static information about whether this
     /// node is a `Leaf` or an `Internal`.
-    unsafe fn cast_unchecked<NewType>(&mut self) -> NodeRef<marker::Mut<'_>, K, V, NewType> {
+    unsafe fn cast_unchecked<NewType>(self) -> NodeRef<marker::Mut<'a>, K, V, NewType> {
         NodeRef { height: self.height, node: self.node, root: self.root, _marker: PhantomData }
     }
 
@@ -724,7 +724,7 @@ fn clone(&self) -> Self {
 }
 
 impl<Node, Type> Handle<Node, Type> {
-    /// Retrieves the node that contains the edge of key/value pair this handle points to.
+    /// Retrieves the node that contains the edge or key/value pair this handle points to.
     pub fn into_node(self) -> Node {
         self.node
     }
@@ -1044,6 +1044,16 @@ pub fn into_kv(self) -> (&'a K, &'a V) {
 }
 
 impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::KV> {
+    pub fn into_key_mut(self) -> &'a mut K {
+        let keys = self.node.into_key_slice_mut();
+        unsafe { keys.get_unchecked_mut(self.idx) }
+    }
+
+    pub fn into_val_mut(self) -> &'a mut V {
+        let vals = self.node.into_val_slice_mut();
+        unsafe { vals.get_unchecked_mut(self.idx) }
+    }
+
     pub fn into_kv_mut(self) -> (&'a mut K, &'a mut V) {
         unsafe {
             let (keys, vals) = self.node.into_slices_mut();
@@ -1108,7 +1118,7 @@ pub fn split(mut self) -> (NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, K, V, R
     }
 
     /// Removes the key/value pair pointed to by this handle and returns it, along with the edge
-    /// between the now adjacent key/value pairs (if any) to the left and right of this handle.
+    /// that the key/value pair collapsed into.
     pub fn remove(
         mut self,
     ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
@@ -1166,7 +1176,7 @@ pub fn can_merge(&self) -> bool {
     /// to by this handle, and the node immediately to the right of this handle into one new
     /// child of the underlying node, returning an edge referencing that new child.
     ///
-    /// Assumes that this edge `.can_merge()`.
+    /// Panics unless this edge `.can_merge()`.
     pub fn merge(
         mut self,
     ) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge> {
@@ -1174,10 +1184,9 @@ pub fn merge(
         let self2 = unsafe { ptr::read(&self) };
         let mut left_node = self1.left_edge().descend();
         let left_len = left_node.len();
-        let mut right_node = self2.right_edge().descend();
+        let right_node = self2.right_edge().descend();
         let right_len = right_node.len();
 
-        // necessary for correctness, but in a private module
         assert!(left_len + right_len < CAPACITY);
 
         unsafe {
@@ -1208,28 +1217,25 @@ pub fn merge(
 
             (*left_node.as_leaf_mut()).len += right_len as u16 + 1;
 
-            let layout = if self.node.height > 1 {
+            if self.node.height > 1 {
+                // SAFETY: the height of the nodes being merged is one below the height
+                // of the node of this edge, thus above zero, so they are internal.
+                let mut left_node = left_node.cast_unchecked();
+                let right_node = right_node.cast_unchecked();
                 ptr::copy_nonoverlapping(
-                    right_node.cast_unchecked().as_internal().edges.as_ptr(),
-                    left_node
-                        .cast_unchecked()
-                        .as_internal_mut()
-                        .edges
-                        .as_mut_ptr()
-                        .add(left_len + 1),
+                    right_node.reborrow().as_internal().edges.as_ptr(),
+                    left_node.reborrow_mut().as_internal_mut().edges.as_mut_ptr().add(left_len + 1),
                     right_len + 1,
                 );
 
                 for i in left_len + 1..left_len + right_len + 2 {
-                    Handle::new_edge(left_node.cast_unchecked().reborrow_mut(), i)
-                        .correct_parent_link();
+                    Handle::new_edge(left_node.reborrow_mut(), i).correct_parent_link();
                 }
 
-                Layout::new::<InternalNode<K, V>>()
+                Global.dealloc(right_node.node.cast(), Layout::new::<InternalNode<K, V>>());
             } else {
-                Layout::new::<LeafNode<K, V>>()
-            };
-            Global.dealloc(right_node.node.cast(), layout);
+                Global.dealloc(right_node.node.cast(), Layout::new::<LeafNode<K, V>>());
+            }
 
             Handle::new_edge(self.node, self.idx)
         }
@@ -1242,8 +1248,8 @@ pub fn steal_left(&mut self) {
         unsafe {
             let (k, v, edge) = self.reborrow_mut().left_edge().descend().pop();
 
-            let k = mem::replace(self.reborrow_mut().into_kv_mut().0, k);
-            let v = mem::replace(self.reborrow_mut().into_kv_mut().1, v);
+            let k = mem::replace(self.kv_mut().0, k);
+            let v = mem::replace(self.kv_mut().1, v);
 
             match self.reborrow_mut().right_edge().descend().force() {
                 ForceResult::Leaf(mut leaf) => leaf.push_front(k, v),
@@ -1259,8 +1265,8 @@ pub fn steal_right(&mut self) {
         unsafe {
             let (k, v, edge) = self.reborrow_mut().right_edge().descend().pop_front();
 
-            let k = mem::replace(self.reborrow_mut().into_kv_mut().0, k);
-            let v = mem::replace(self.reborrow_mut().into_kv_mut().1, v);
+            let k = mem::replace(self.kv_mut().0, k);
+            let v = mem::replace(self.kv_mut().1, v);
 
             match self.reborrow_mut().left_edge().descend().force() {
                 ForceResult::Leaf(mut leaf) => leaf.push(k, v),
@@ -1288,7 +1294,7 @@ pub fn bulk_steal_left(&mut self, count: usize) {
                 let left_kv = left_node.reborrow_mut().into_kv_pointers_mut();
                 let right_kv = right_node.reborrow_mut().into_kv_pointers_mut();
                 let parent_kv = {
-                    let kv = self.reborrow_mut().into_kv_mut();
+                    let kv = self.kv_mut();
                     (kv.0 as *mut K, kv.1 as *mut V)
                 };
 
@@ -1345,7 +1351,7 @@ pub fn bulk_steal_right(&mut self, count: usize) {
                 let left_kv = left_node.reborrow_mut().into_kv_pointers_mut();
                 let right_kv = right_node.reborrow_mut().into_kv_pointers_mut();
                 let parent_kv = {
-                    let kv = self.reborrow_mut().into_kv_mut();
+                    let kv = self.kv_mut();
                     (kv.0 as *mut K, kv.1 as *mut V)
                 };