BTree: refactor clone, mostly for readability

Author: ssomers

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

@@ -0,0 +1,54 @@
+use super::node::{marker, NodeRef};
+
+/// Simple tree around a raw `NodeRef`, that provides a destructor.
+/// Unlike smarter sister `BTreeMap`, it does not keep track of element count,
+/// does not allow replacing the root, therefore can be statically tyoed by the
+/// node type of the root, and cannot represent a tree without root node
+/// (its `Option` exists only to disable the `drop` method when needed).
+pub struct BareTree<K, V, RootType>(Option<NodeRef<marker::Owned, K, V, RootType>>);
+
+impl<K, V, RootType> Drop for BareTree<K, V, RootType> {
+    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() } {
+                unsafe { kv.drop_key_val() };
+                cur_edge = next_edge;
+            }
+        }
+    }
+}
+
+impl<K, V> BareTree<K, V, marker::Leaf> {
+    /// Returns a new tree consisting of a single leaf that is initially empty.
+    pub fn new_leaf() -> Self {
+        Self(Some(NodeRef::new_leaf()))
+    }
+}
+
+impl<K, V> BareTree<K, V, marker::Internal> {
+    /// Returns a new tree with an internal root, that initially has no elements
+    /// and one child.
+    pub fn new_internal(child: NodeRef<marker::Owned, K, V, marker::LeafOrInternal>) -> Self {
+        Self(Some(NodeRef::new_internal(child)))
+    }
+}
+
+impl<K, V, RootType> BareTree<K, V, RootType> {
+    /// Mutably borrows the owned root node.
+    pub fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, RootType> {
+        self.0.as_mut().unwrap().borrow_mut()
+    }
+
+    /// Removes any static information asserting that the root is a `Leaf` or
+    /// `Internal` node.
+    pub fn forget_root_type(mut self) -> BareTree<K, V, marker::LeafOrInternal> {
+        BareTree(self.0.take().map(NodeRef::forget_type))
+    }
+
+    /// Consumes the tree, returning a `NodeRef` that still enforces ownership
+    /// but lacks a destructor.
+    pub fn into_inner(mut self) -> NodeRef<marker::Owned, K, V, RootType> {
+        self.0.take().unwrap()
+    }
+}

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

@@ -9,6 +9,7 @@ use core::mem::{self, ManuallyDrop};
 use core::ops::{Index, RangeBounds};
 use core::ptr;
 
+use super::bare_tree::BareTree;
 use super::borrow::DormantMutRef;
 use super::dedup_sorted_iter::DedupSortedIter;
 use super::navigate::{LazyLeafRange, LeafRange};
@@ -171,75 +172,47 @@ 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> {
-        fn clone_subtree<'a, K: Clone, V: Clone>(
+        fn clone_subtree<'a, K: Clone + 'a, V: Clone + 'a>(
             node: NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>,
-        ) -> BTreeMap<K, V>
-        where
-            K: 'a,
-            V: 'a,
-        {
+        ) -> BareTree<K, V, marker::LeafOrInternal> {
             match node.force() {
                 Leaf(leaf) => {
-                    let mut out_tree = BTreeMap { root: Some(Root::new()), length: 0 };
-
+                    let mut out_tree = BareTree::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.forget_root_type()
                 }
                 Internal(internal) => {
-                    let mut out_tree = clone_subtree(internal.first_edge().descend());
-
+                    let mut in_edge = internal.first_edge();
+                    let first_child = clone_subtree(in_edge.descend());
+                    let mut out_tree = BareTree::new_internal(first_child.into_inner());
                     {
-                        let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root);
-                        let mut out_node = out_root.push_internal_level();
-                        let mut in_edge = internal.first_edge();
+                        let mut out_node = out_tree.borrow_mut();
                         while let Ok(kv) = in_edge.right_kv() {
                             let (k, v) = kv.into_kv();
                             in_edge = kv.right_edge();
 
                             let k = (*k).clone();
                             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.into_inner());
                         }
                     }
-
-                    out_tree
+                    out_tree.forget_root_type()
                 }
             }
         }
 
-        if self.is_empty() {
-            BTreeMap::new()
-        } else {
-            clone_subtree(self.root.as_ref().unwrap().reborrow()) // unwrap succeeds because not empty
-        }
+        let cloned_root = self.root.as_ref().map(|r| clone_subtree(r.reborrow()).into_inner());
+        BTreeMap { root: cloned_root, length: self.length }
     }
 }
 

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

@@ -1,4 +1,5 @@
 mod append;
+mod bare_tree;
 mod borrow;
 mod dedup_sorted_iter;
 mod fix;

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

@@ -441,7 +441,7 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     ///   `deallocating_next_back`.
     /// - The returned KV handle is only valid to access the key and value,
     ///   and only valid until the next call to a `deallocating_` method.
-    unsafe fn deallocating_next(
+    pub unsafe fn deallocating_next(
         self,
     ) -> Option<(Self, Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>)>
     {

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

@@ -213,7 +213,8 @@ unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Owned, K, V, Type>
 unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Dying, K, V, Type> {}
 
 impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
-    fn new_leaf() -> Self {
+    /// Returns a new owned leaf node, that is initially empty.
+    pub fn new_leaf() -> Self {
         Self::from_new_leaf(LeafNode::new())
     }
 
@@ -223,7 +224,8 @@ 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 {
+    /// Returns a new owned internal node, that initially has no elements and one child.
+    pub fn new_internal(child: Root<K, V>) -> Self {
         let mut new_node = unsafe { InternalNode::new() };
         new_node.edges[0].write(child.node);
         unsafe { NodeRef::from_new_internal(new_node, child.height + 1) }
@@ -651,15 +653,8 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
     }
 }
 
-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, NodeType> NodeRef<BorrowType, K, V, NodeType> {
+    /// 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 }
     }
@@ -1505,31 +1500,15 @@ impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
     }
 }
 
-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> {