BTreeMap: disentangle Drop implementation from IntoIter

Author: ssomers

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

@@ -138,8 +138,8 @@ pub struct BTreeMap<K, V> {
 #[stable(feature = "btree_drop", since = "1.7.0")]
 unsafe impl<#[may_dangle] K, #[may_dangle] V> Drop for BTreeMap<K, V> {
     fn drop(&mut self) {
-        unsafe {
-            drop(ptr::read(self).into_iter());
+        if let Some(root) = self.root.take() {
+            Dropper { front: root.into_dying().first_leaf_edge(), remaining_length: self.length };
         }
     }
 }
@@ -325,6 +325,14 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for IntoIter<K, V> {
     }
 }
 
+/// A simplified version of `IntoIter` that is not double-ended and has only one
+/// purpose: to drop the remainder of an `IntoIter`. Therefore it also serves to
+/// drop an entire tree without the need to look up a `back` leaf edge.
+struct Dropper<K, V> {
+    front: Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
+    remaining_length: usize,
+}
+
 /// An iterator over the keys of a `BTreeMap`.
 ///
 /// This `struct` is created by the [`keys`] method on [`BTreeMap`]. See its
@@ -1373,42 +1381,42 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
     }
 }
 
-#[stable(feature = "btree_drop", since = "1.7.0")]
-impl<K, V> Drop for IntoIter<K, V> {
+impl<K, V> Drop for Dropper<K, V> {
     fn drop(&mut self) {
-        struct DropGuard<'a, K, V>(&'a mut IntoIter<K, V>);
+        // Similar to advancing a non-fusing iterator.
+        fn next_or_end<K, V>(this: &mut Dropper<K, V>) -> Option<(K, V)> {
+            if this.remaining_length == 0 {
+                unsafe { ptr::read(&this.front).deallocating_end() }
+                None
+            } else {
+                this.remaining_length -= 1;
+                Some(unsafe { this.front.next_unchecked() })
+            }
+        }
+
+        struct DropGuard<'a, K, V>(&'a mut Dropper<K, V>);
 
         impl<'a, K, V> Drop for DropGuard<'a, 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(_) = self.0.next() {}
-
-                unsafe {
-                    let mut node =
-                        unwrap_unchecked(ptr::read(&self.0.front)).into_node().forget_type();
-                    while let Some(parent) = node.deallocate_and_ascend() {
-                        node = parent.into_node().forget_type();
-                    }
-                }
+                while let Some(_pair) = next_or_end(&mut self.0) {}
             }
         }
 
-        while let Some(pair) = self.next() {
+        while let Some(pair) = next_or_end(self) {
             let guard = DropGuard(self);
             drop(pair);
             mem::forget(guard);
         }
+    }
+}
 
-        unsafe {
-            if let Some(front) = ptr::read(&self.front) {
-                let mut node = front.into_node().forget_type();
-                // Most of the nodes have been deallocated while traversing
-                // but one pile from a leaf up to the root is left standing.
-                while let Some(parent) = node.deallocate_and_ascend() {
-                    node = parent.into_node().forget_type();
-                }
-            }
+#[stable(feature = "btree_drop", since = "1.7.0")]
+impl<K, V> Drop for IntoIter<K, V> {
+    fn drop(&mut self) {
+        if let Some(front) = self.front.take() {
+            Dropper { front, remaining_length: self.length };
         }
     }
 }

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

@@ -290,37 +290,71 @@ impl<BorrowType: marker::BorrowType, K, V>
     }
 }
 
-macro_rules! def_next_kv_uncheched_dealloc {
-    { unsafe fn $name:ident : $adjacent_kv:ident } => {
-        /// Given a leaf edge handle into an owned tree, returns a handle to the next KV,
-        /// while deallocating any node left behind yet leaving the corresponding edge
-        /// in its parent node dangling.
-        ///
-        /// # Safety
-        /// - The leaf edge must not be the last one in the direction travelled.
-        /// - The node carrying the next KV returned must not have been deallocated by a
-        ///   previous call on any handle obtained for this tree.
-        unsafe fn $name <K, V>(
-            leaf_edge: Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
-        ) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
-            let mut edge = leaf_edge.forget_node_type();
-            loop {
-                edge = match edge.$adjacent_kv() {
-                    Ok(internal_kv) => return internal_kv,
-                    Err(last_edge) => {
-                        unsafe {
-                            let parent_edge = last_edge.into_node().deallocate_and_ascend();
-                            unwrap_unchecked(parent_edge).forget_node_type()
-                        }
-                    }
+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.
+    ///
+    /// 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.
+    ///
+    /// # Safety
+    /// The next KV must not have been previously returned by counterpart `deallocating_next_back`.
+    unsafe fn deallocating_next(self) -> Option<(Self, (K, V))> {
+        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)));
                 }
+                Err(last_edge) => match unsafe { last_edge.into_node().deallocate_and_ascend() } {
+                    Some(parent_edge) => parent_edge.forget_node_type(),
+                    None => return None,
+                },
             }
         }
-    };
-}
+    }
 
-def_next_kv_uncheched_dealloc! {unsafe fn next_kv_unchecked_dealloc: right_kv}
-def_next_kv_uncheched_dealloc! {unsafe fn next_back_kv_unchecked_dealloc: left_kv}
+    /// 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.
+    ///
+    /// 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.
+    ///
+    /// # Safety
+    /// The next KV must not have been previously returned by counterpart `deallocating_next`.
+    unsafe fn deallocating_next_back(self) -> Option<(Self, (K, V))> {
+        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)));
+                }
+                Err(last_edge) => match unsafe { last_edge.into_node().deallocate_and_ascend() } {
+                    Some(parent_edge) => parent_edge.forget_node_type(),
+                    None => return None,
+                },
+            }
+        }
+    }
+
+    /// Deallocates a pile of nodes from the leaf up to the root.
+    /// This is useful when `deallocating_next` and `deallocating_next_back`
+    /// have been nibbling at both sides of the same tree.
+    pub fn deallocating_end(self) {
+        let mut edge = self.forget_node_type();
+        while let Some(parent_edge) = unsafe { edge.into_node().deallocate_and_ascend() } {
+            edge = parent_edge.forget_node_type();
+        }
+    }
+}
 
 impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge> {
     /// Moves the leaf edge handle to the next leaf edge and returns references to the
@@ -396,11 +430,8 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     /// call this method again subject to its safety conditions, or call counterpart
     /// `next_back_unchecked` subject to its safety conditions.
     pub unsafe fn next_unchecked(&mut self) -> (K, V) {
-        super::mem::replace(self, |leaf_edge| {
-            let kv = unsafe { next_kv_unchecked_dealloc(leaf_edge) };
-            let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
-            let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
-            (kv.next_leaf_edge(), (k, v))
+        super::mem::replace(self, |leaf_edge| unsafe {
+            unwrap_unchecked(leaf_edge.deallocating_next())
         })
     }
 
@@ -417,11 +448,8 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     /// call this method again subject to its safety conditions, or call counterpart
     /// `next_unchecked` subject to its safety conditions.
     pub unsafe fn next_back_unchecked(&mut self) -> (K, V) {
-        super::mem::replace(self, |leaf_edge| {
-            let kv = unsafe { next_back_kv_unchecked_dealloc(leaf_edge) };
-            let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
-            let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
-            (kv.next_back_leaf_edge(), (k, v))
+        super::mem::replace(self, |leaf_edge| unsafe {
+            unwrap_unchecked(leaf_edge.deallocating_next_back())
         })
     }
 }