BTreeMap: clean up implementation of clone

ssomers · ssomers · commit 35c7567c7d69 · 2021-02-02T17:25:13.000+01:00
diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
@@ -147,42 +147,59 @@ unsafe impl<#[may_dangle] K, #[may_dangle] V> Drop for BTreeMap<K, V> {
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
     fn clone(&self) -> BTreeMap<K, V> {
+        struct GuardedTree<K, V, Type>(Option<NodeRef<marker::Owned, K, V, Type>>);
+        impl<K, V, Type> GuardedTree<K, V, Type> {
+            fn new(root: NodeRef<marker::Owned, K, V, Type>) -> Self {
+                Self(Some(root))
+            }
+            fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, Type> {
+                self.0.as_mut().unwrap().borrow_mut()
+            }
+            fn leak(mut self) -> Root<K, V> {
+                self.0.take().unwrap().forget_type()
+            }
+        }
+        impl<K, V, Type> Drop for GuardedTree<K, V, Type> {
+            fn drop(&mut self) {
+                if let Some(root) = self.0.take() {
+                    let mut cur_edge = root.forget_type().into_dying().first_leaf_edge();
+                    while let Some((next_edge, _kv)) = unsafe { cur_edge.deallocating_next() } {
+                        cur_edge = next_edge;
+                    }
+                }
+            }
+        }
+
         fn clone_subtree<'a, K: Clone, V: Clone>(
             node: NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>,
-        ) -> BTreeMap<K, V>
+        ) -> Root<K, V>
         where
             K: 'a,
             V: 'a,
         {
             match node.force() {
                 Leaf(leaf) => {
-                    let mut out_tree = BTreeMap { root: Some(Root::new()), length: 0 };
+                    let mut out_tree = GuardedTree::new(NodeRef::new_leaf());
 
                     {
-                        let root = out_tree.root.as_mut().unwrap(); // unwrap succeeds because we just wrapped
-                        let mut out_node = match root.borrow_mut().force() {
-                            Leaf(leaf) => leaf,
-                            Internal(_) => unreachable!(),
-                        };
-
+                        let mut out_node = out_tree.borrow_mut();
                         let mut in_edge = leaf.first_edge();
                         while let Ok(kv) = in_edge.right_kv() {
                             let (k, v) = kv.into_kv();
                             in_edge = kv.right_edge();
 
                             out_node.push(k.clone(), v.clone());
-                            out_tree.length += 1;
                         }
                     }
 
-                    out_tree
+                    out_tree.leak()
                 }
                 Internal(internal) => {
-                    let mut out_tree = clone_subtree(internal.first_edge().descend());
+                    let first_child = clone_subtree(internal.first_edge().descend());
+                    let mut out_tree = GuardedTree::new(NodeRef::new_internal(first_child));
 
                     {
-                        let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root);
-                        let mut out_node = out_root.push_internal_level();
+                        let mut out_node = out_tree.borrow_mut();
                         let mut in_edge = internal.first_edge();
                         while let Ok(kv) = in_edge.right_kv() {
                             let (k, v) = kv.into_kv();
@@ -192,32 +209,17 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
                             let v = (*v).clone();
                             let subtree = clone_subtree(in_edge.descend());
 
-                            // We can't destructure subtree directly
-                            // because BTreeMap implements Drop
-                            let (subroot, sublength) = unsafe {
-                                let subtree = ManuallyDrop::new(subtree);
-                                let root = ptr::read(&subtree.root);
-                                let length = subtree.length;
-                                (root, length)
-                            };
-
-                            out_node.push(k, v, subroot.unwrap_or_else(Root::new));
-                            out_tree.length += 1 + sublength;
+                            out_node.push(k, v, subtree);
                         }
                     }
 
-                    out_tree
+                    out_tree.leak()
                 }
             }
         }
 
-        if self.is_empty() {
-            // Ideally we'd call `BTreeMap::new` here, but that has the `K:
-            // Ord` constraint, which this method lacks.
-            BTreeMap { root: None, length: 0 }
-        } else {
-            clone_subtree(self.root.as_ref().unwrap().reborrow()) // unwrap succeeds because not empty
-        }
+        let cloned_root = self.root.as_ref().map(|root| clone_subtree(root.reborrow()));
+        BTreeMap { root: cloned_root, length: self.length }
     }
 }
 
diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs
@@ -301,7 +301,7 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     ///
     /// # Safety
     /// The next KV must not have been previously returned by counterpart `deallocating_next_back`.
-    unsafe fn deallocating_next(self) -> Option<(Self, (K, V))> {
+    pub unsafe fn deallocating_next(self) -> Option<(Self, (K, V))> {
         let mut edge = self.forget_node_type();
         loop {
             edge = match edge.right_kv() {
diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
@@ -132,7 +132,7 @@ impl<K, V> Root<K, V> {
 }
 
 impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
-    fn new_leaf() -> Self {
+    pub fn new_leaf() -> Self {
         Self::from_new_leaf(Box::new(unsafe { LeafNode::new() }))
     }
 
@@ -142,7 +142,7 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
 }
 
 impl<K, V> NodeRef<marker::Owned, K, V, marker::Internal> {
-    fn new_internal(child: Root<K, V>) -> Self {
+    pub fn new_internal(child: Root<K, V>) -> Self {
         let mut new_node = Box::new(unsafe { InternalNode::new() });
         new_node.edges[0].write(child.node);
         NodeRef::from_new_internal(new_node, child.height + 1)
@@ -1461,45 +1461,22 @@ impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
     }
 }
 
-impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Leaf> {
-    /// Removes any static information asserting that this node is a `Leaf` node.
-    pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
-        NodeRef { height: self.height, node: self.node, _marker: PhantomData }
-    }
-}
-
-impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
-    /// Removes any static information asserting that this node is an `Internal` node.
+impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
+    /// Removes any static information asserting that this node is a `Leaf` or `Internal` node.
     pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
         NodeRef { height: self.height, node: self.node, _marker: PhantomData }
     }
 }
 
-impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
-    pub fn forget_node_type(
-        self,
-    ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::Edge> {
-        unsafe { Handle::new_edge(self.node.forget_type(), self.idx) }
-    }
-}
-
-impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge> {
+impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
     pub fn forget_node_type(
         self,
     ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::Edge> {
         unsafe { Handle::new_edge(self.node.forget_type(), self.idx) }
     }
 }
 
-impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::KV> {
-    pub fn forget_node_type(
-        self,
-    ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
-        unsafe { Handle::new_kv(self.node.forget_type(), self.idx) }
-    }
-}
-
-impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV> {
+impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV> {
     pub fn forget_node_type(
         self,
     ) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
