BTree: no longer copy keys and values before dropping them

ssomers · ssomers · commit 728204b40efe · 2021-05-07T10:53:53.000+02:00
diff --git 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
@@ -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
@@ -422,21 +422,30 @@ 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.
         unsafe { &mut *ptr }
     }
 }
 
+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.

Original file line number	Diff line number	Diff line change
`@@ -1439,7 +1439,10 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {`
`1439`	`1439`	`impl<K, V> Drop for Dropper<K, V> {`
`1440`	`1440`	`fn drop(&mut self) {`
`1441`	`1441`	`// Similar to advancing a non-fusing iterator.`
`1442`		`- fn next_or_end<K, V>(this: &mut Dropper<K, V>) -> Option<(K, V)> {`
	`1442`	`+ fn next_or_end<K, V>(`
	`1443`	`+ this: &mut Dropper<K, V>,`
	`1444`	`+ ) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>>`
	`1445`	`+ {`
`1443`	`1446`	`if this.remaining_length == 0 {`
`1444`	`1447`	`unsafe { ptr::read(&this.front).deallocating_end() }`
`1445`	`1448`	`None`
`@@ -1455,13 +1458,15 @@ impl<K, V> Drop for Dropper<K, V> {`
`1455`	`1458`	`fn drop(&mut self) {`
`1456`	`1459`	`// Continue the same loop we perform below. This only runs when unwinding, so we`
`1457`	`1460`	`// don't have to care about panics this time (they'll abort).`
`1458`		`- while let Some(_pair) = next_or_end(&mut self.0) {}`
	`1461`	`+ while let Some(kv) = next_or_end(&mut self.0) {`
	`1462`	`+ kv.drop_key_val();`
	`1463`	`+ }`
`1459`	`1464`	`}`
`1460`	`1465`	`}`
`1461`	`1466`
`1462`		`- while let Some(pair) = next_or_end(self) {`
	`1467`	`+ while let Some(kv) = next_or_end(self) {`
`1463`	`1468`	`let guard = DropGuard(self);`
`1464`		`- drop(pair);`
	`1469`	`+ kv.drop_key_val();`
`1465`	`1470`	`mem::forget(guard);`
`1466`	`1471`	`}`
`1467`	`1472`	`}`
`@@ -1485,7 +1490,9 @@ impl<K, V> Iterator for IntoIter<K, V> {`
`1485`	`1490`	`None`
`1486`	`1491`	`} else {`
`1487`	`1492`	`self.length -= 1;`
`1488`		`- Some(unsafe { self.range.front.as_mut().unwrap().deallocating_next_unchecked() })`
	`1493`	`+ let front = self.range.front.as_mut().unwrap();`
	`1494`	`+ let kv = unsafe { front.deallocating_next_unchecked() };`
	`1495`	`+ Some(kv.into_key_val())`
`1489`	`1496`	`}`
`1490`	`1497`	`}`
`1491`	`1498`
`@@ -1501,7 +1508,9 @@ impl<K, V> DoubleEndedIterator for IntoIter<K, V> {`
`1501`	`1508`	`None`
`1502`	`1509`	`} else {`
`1503`	`1510`	`self.length -= 1;`
`1504`		`- Some(unsafe { self.range.back.as_mut().unwrap().deallocating_next_back_unchecked() })`
	`1511`	`+ let back = self.range.back.as_mut().unwrap();`
	`1512`	`+ let kv = unsafe { back.deallocating_next_back_unchecked() };`
	`1513`	`+ Some(kv.into_key_val())`
`1505`	`1514`	`}`
`1506`	`1515`	`}`
`1507`	`1516`	`}`