Auto merge of #84904 - ssomers:btree_drop_kv_in_place, r=Mark-Simulacrum

BTree: no longer copy keys and values before dropping them

When dropping BTreeMap or BTreeSet instances, keys-value pairs are up to now each copied and then dropped, at least according to source code. This is because the code for dropping and for iterators is shared.

This PR postpones the treatment of doomed key-value pairs from the intermediate functions `deallocating_next`(`_back`) to the last minute, so the we can drop the keys and values in place. According to the library/alloc benchmarks, this does make a difference, (and a positive difference with an `#[inline]` on `drop_key_val`). It does not change anything for #81444 though.

r? `@Mark-Simulacrum`

3 files changed, 95 insertions, 45 deletions

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index 0a46387c34e..30194aa446f 100644
--- a/library/alloc/src/collections/btree/map.rs
+++ b/library/alloc/src/collections/btree/map.rs
@@ -1439,7 +1439,10 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
 impl<K, V> Drop for Dropper<K, V> {
     fn drop(&mut self) {
         // Similar to advancing a non-fusing iterator.
-        fn next_or_end<K, V>(this: &mut Dropper<K, V>) -> Option<(K, V)> {
+        fn next_or_end<K, V>(
+            this: &mut Dropper<K, V>,
+        ) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>>
+        {
             if this.remaining_length == 0 {
                 unsafe { ptr::read(&this.front).deallocating_end() }
                 None
@@ -1455,13 +1458,15 @@ impl<K, V> Drop for Dropper<K, V> {
             fn drop(&mut self) {
                 // Continue the same loop we perform below. This only runs when unwinding, so we
                 // don't have to care about panics this time (they'll abort).
-                while let Some(_pair) = next_or_end(&mut self.0) {}
+                while let Some(kv) = next_or_end(&mut self.0) {
+                    kv.drop_key_val();
+                }
             }
         }
 
-        while let Some(pair) = next_or_end(self) {
+        while let Some(kv) = next_or_end(self) {
             let guard = DropGuard(self);
-            drop(pair);
+            kv.drop_key_val();
             mem::forget(guard);
         }
     }
@@ -1485,7 +1490,9 @@ impl<K, V> Iterator for IntoIter<K, V> {
             None
         } else {
             self.length -= 1;
-            Some(unsafe { self.range.front.as_mut().unwrap().deallocating_next_unchecked() })
+            let front = self.range.front.as_mut().unwrap();
+            let kv = unsafe { front.deallocating_next_unchecked() };
+            Some(kv.into_key_val())
         }
     }
 
@@ -1501,7 +1508,9 @@ impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
             None
         } else {
             self.length -= 1;
-            Some(unsafe { self.range.back.as_mut().unwrap().deallocating_next_back_unchecked() })
+            let back = self.range.back.as_mut().unwrap();
+            let kv = unsafe { back.deallocating_next_back_unchecked() };
+            Some(kv.into_key_val())
         }
     }
 }
diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs
index 4399feaccc9..563c070dd0f 100644
--- a/library/alloc/src/collections/btree/navigate.rs
+++ b/library/alloc/src/collections/btree/navigate.rs
@@ -237,25 +237,27 @@ impl<BorrowType: marker::BorrowType, K, V>
 
 impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     /// Given a leaf edge handle into a dying tree, returns the next leaf edge
-    /// on the right side, and the key-value pair in between, which is either
-    /// in the same leaf node, in an ancestor node, or non-existent.
+    /// on the right side, and the key-value pair in between, if they exist.
     ///
-    /// This method also deallocates any node(s) it reaches the end of. This
-    /// implies that if no more key-value pair exists, the entire remainder of
-    /// the tree will have been deallocated and there is nothing left to return.
+    /// If the given edge is the last one in a leaf, this method deallocates
+    /// the leaf, as well as any ancestor nodes whose last edge was reached.
+    /// This implies that if no more key-value pair follows, the entire tree
+    /// will have been deallocated and there is nothing left to return.
     ///
     /// # Safety
-    /// The given edge must not have been previously returned by counterpart
-    /// `deallocating_next_back`.
-    unsafe fn deallocating_next(self) -> Option<(Self, (K, V))> {
+    /// - The given edge must not have been previously returned by counterpart
+    ///   `deallocating_next_back`.
+    /// - The returned KV handle is only valid to access the key and value,
+    ///   and only valid until the next call to this method or counterpart
+    ///   `deallocating_next_back`.
+    pub unsafe fn deallocating_next(
+        self,
+    ) -> Option<(Self, Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>)>
+    {
         let mut edge = self.forget_node_type();
         loop {
             edge = match edge.right_kv() {
-                Ok(kv) => {
-                    let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
-                    let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
-                    return Some((kv.next_leaf_edge(), (k, v)));
-                }
+                Ok(kv) => return Some((unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)),
                 Err(last_edge) => match unsafe { last_edge.into_node().deallocate_and_ascend() } {
                     Some(parent_edge) => parent_edge.forget_node_type(),
                     None => return None,
@@ -265,25 +267,27 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     }
 
     /// Given a leaf edge handle into a dying tree, returns the next leaf edge
-    /// on the left side, and the key-value pair in between, which is either
-    /// in the same leaf node, in an ancestor node, or non-existent.
+    /// on the left side, and the key-value pair in between, if they exist.
     ///
-    /// This method also deallocates any node(s) it reaches the end of. This
-    /// implies that if no more key-value pair exists, the entire remainder of
-    /// the tree will have been deallocated and there is nothing left to return.
+    /// If the given edge is the first one in a leaf, this method deallocates
+    /// the leaf, as well as any ancestor nodes whose first edge was reached.
+    /// This implies that if no more key-value pair follows, the entire tree
+    /// will have been deallocated and there is nothing left to return.
     ///
     /// # Safety
-    /// The given edge must not have been previously returned by counterpart
-    /// `deallocating_next`.
-    unsafe fn deallocating_next_back(self) -> Option<(Self, (K, V))> {
+    /// - The given edge must not have been previously returned by counterpart
+    ///   `deallocating_next`.
+    /// - The returned KV handle is only valid to access the key and value,
+    ///   and only valid until the next call to this method or counterpart
+    ///   `deallocating_next`.
+    unsafe fn deallocating_next_back(
+        self,
+    ) -> Option<(Self, Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>)>
+    {
         let mut edge = self.forget_node_type();
         loop {
             edge = match edge.left_kv() {
-                Ok(kv) => {
-                    let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
-                    let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
-                    return Some((kv.next_back_leaf_edge(), (k, v)));
-                }
+                Ok(kv) => return Some((unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)),
                 Err(last_edge) => match unsafe { last_edge.into_node().deallocate_and_ascend() } {
                     Some(parent_edge) => parent_edge.forget_node_type(),
                     None => return None,
@@ -373,13 +377,15 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     ///
     /// # Safety
     /// - There must be another KV in the direction travelled.
-    /// - That KV was not previously returned by counterpart `next_back_unchecked`
-    ///   on any copy of the handles being used to traverse the tree.
+    /// - That KV was not previously returned by counterpart
+    ///   `deallocating_next_back_unchecked` on any copy of the handles
+    ///   being used to traverse the tree.
     ///
     /// The only safe way to proceed with the updated handle is to compare it, drop it,
-    /// call this method again subject to its safety conditions, or call counterpart
-    /// `next_back_unchecked` subject to its safety conditions.
-    pub unsafe fn deallocating_next_unchecked(&mut self) -> (K, V) {
+    /// or call this method or counterpart `deallocating_next_back_unchecked` again.
+    pub unsafe fn deallocating_next_unchecked(
+        &mut self,
+    ) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
         super::mem::replace(self, |leaf_edge| unsafe {
             leaf_edge.deallocating_next().unwrap_unchecked()
         })
@@ -391,13 +397,15 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     ///
     /// # Safety
     /// - There must be another KV in the direction travelled.
-    /// - That leaf edge was not previously returned by counterpart `next_unchecked`
-    ///   on any copy of the handles being used to traverse the tree.
+    /// - That leaf edge was not previously returned by counterpart
+    ///   `deallocating_next_unchecked` on any copy of the handles
+    ///   being used to traverse the tree.
     ///
     /// The only safe way to proceed with the updated handle is to compare it, drop it,
-    /// call this method again subject to its safety conditions, or call counterpart
-    /// `next_unchecked` subject to its safety conditions.
-    pub unsafe fn deallocating_next_back_unchecked(&mut self) -> (K, V) {
+    /// or call this method or counterpart `deallocating_next_unchecked` again.
+    pub unsafe fn deallocating_next_back_unchecked(
+        &mut self,
+    ) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
         super::mem::replace(self, |leaf_edge| unsafe {
             leaf_edge.deallocating_next_back().unwrap_unchecked()
         })
diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
index af403496e38..3c453529ba8 100644
--- a/library/alloc/src/collections/btree/node.rs
+++ b/library/alloc/src/collections/btree/node.rs
@@ -422,14 +422,14 @@ impl<'a, K, V, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
         NodeRef { height: self.height, node: self.node, _marker: PhantomData }
     }
 
-    /// Borrows exclusive access to the leaf portion of any leaf or internal node.
+    /// Borrows exclusive access to the leaf portion of a leaf or internal node.
     fn as_leaf_mut(&mut self) -> &mut LeafNode<K, V> {
         let ptr = Self::as_leaf_ptr(self);
         // SAFETY: we have exclusive access to the entire node.
         unsafe { &mut *ptr }
     }
 
-    /// Offers exclusive access to the leaf portion of any leaf or internal node.
+    /// Offers exclusive access to the leaf portion of a leaf or internal node.
     fn into_leaf_mut(mut self) -> &'a mut LeafNode<K, V> {
         let ptr = Self::as_leaf_ptr(&mut self);
         // SAFETY: we have exclusive access to the entire node.
@@ -437,6 +437,15 @@ impl<'a, K, V, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
     }
 }
 
+impl<K, V, Type> NodeRef<marker::Dying, K, V, Type> {
+    /// Borrows exclusive access to the leaf portion of a dying leaf or internal node.
+    fn as_leaf_dying(&mut self) -> &mut LeafNode<K, V> {
+        let ptr = Self::as_leaf_ptr(self);
+        // SAFETY: we have exclusive access to the entire node.
+        unsafe { &mut *ptr }
+    }
+}
+
 impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
     /// Borrows exclusive access to an element of the key storage area.
     ///
@@ -1040,13 +1049,37 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>
         }
     }
 
-    /// Replace the key and value that the KV handle refers to.
+    /// Replaces the key and value that the KV handle refers to.
     pub fn replace_kv(&mut self, k: K, v: V) -> (K, V) {
         let (key, val) = self.kv_mut();
         (mem::replace(key, k), mem::replace(val, v))
     }
 }
 
+impl<K, V, NodeType> Handle<NodeRef<marker::Dying, K, V, NodeType>, marker::KV> {
+    /// Extracts the key and value that the KV handle refers to.
+    pub fn into_key_val(mut self) -> (K, V) {
+        debug_assert!(self.idx < self.node.len());
+        let leaf = self.node.as_leaf_dying();
+        unsafe {
+            let key = leaf.keys.get_unchecked_mut(self.idx).assume_init_read();
+            let val = leaf.vals.get_unchecked_mut(self.idx).assume_init_read();
+            (key, val)
+        }
+    }
+
+    /// Drops the key and value that the KV handle refers to.
+    #[inline]
+    pub fn drop_key_val(mut self) {
+        debug_assert!(self.idx < self.node.len());
+        let leaf = self.node.as_leaf_dying();
+        unsafe {
+            leaf.keys.get_unchecked_mut(self.idx).assume_init_drop();
+            leaf.vals.get_unchecked_mut(self.idx).assume_init_drop();
+        }
+    }
+}
+
 impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::KV> {
     /// Helps implementations of `split` for a particular `NodeType`,
     /// by taking care of leaf data.