From d199c5b0206b964abadc8539ad4e959ec8e9413b Mon Sep 17 00:00:00 2001
From: Stein Somers <git@steinsomers.be>
Date: Wed, 18 Nov 2020 18:19:38 +0100
Subject: [PATCH] BTreeMap: expose new_internal function and sanitize
 from_new_internal

---
 library/alloc/src/collections/btree/node.rs   | 21 +++++++++++--------
 .../alloc/src/collections/btree/node/tests.rs |  6 ++----
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
index e7d66a9a6a83e..220c98b294cef 100644
--- a/library/alloc/src/collections/btree/node.rs
+++ b/library/alloc/src/collections/btree/node.rs
@@ -142,8 +142,17 @@ 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 {
+        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)
+    }
+
     fn from_new_internal(internal: Box<InternalNode<K, V>>, height: usize) -> Self {
-        NodeRef { height, node: NonNull::from(Box::leak(internal)).cast(), _marker: PhantomData }
+        let node = NonNull::from(Box::leak(internal)).cast();
+        let mut this = NodeRef { height, node, _marker: PhantomData };
+        this.borrow_mut().correct_all_childrens_parent_links();
+        this
     }
 }
 
@@ -167,11 +176,7 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
     /// make that new node the root node, and return it. This increases the height by 1
     /// and is the opposite of `pop_internal_level`.
     pub fn push_internal_level(&mut self) -> NodeRef<marker::Mut<'_>, K, V, marker::Internal> {
-        let mut new_node = Box::new(unsafe { InternalNode::new() });
-        new_node.edges[0].write(self.node);
-        let mut new_root = NodeRef::from_new_internal(new_node, self.height + 1);
-        new_root.borrow_mut().first_edge().correct_parent_link();
-        *self = new_root.forget_type();
+        super::mem::take_mut(self, |old_root| NodeRef::new_internal(old_root).forget_type());
 
         // `self.borrow_mut()`, except that we just forgot we're internal now:
         NodeRef { height: self.height, node: self.node, _marker: PhantomData }
@@ -1193,9 +1198,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>,
             );
 
             let height = self.node.height;
-            let mut right = NodeRef::from_new_internal(new_node, height);
-
-            right.borrow_mut().correct_childrens_parent_links(0..new_len + 1);
+            let right = NodeRef::from_new_internal(new_node, height);
 
             SplitResult { left: self.node, kv, right }
         }
diff --git a/library/alloc/src/collections/btree/node/tests.rs b/library/alloc/src/collections/btree/node/tests.rs
index 7fe8ff743c040..48ce9f2bd89c8 100644
--- a/library/alloc/src/collections/btree/node/tests.rs
+++ b/library/alloc/src/collections/btree/node/tests.rs
@@ -79,10 +79,8 @@ fn test_splitpoint() {
 #[test]
 fn test_partial_cmp_eq() {
     let mut root1 = NodeRef::new_leaf();
-    let mut leaf1 = root1.borrow_mut();
-    leaf1.push(1, ());
-    let mut root1 = root1.forget_type();
-    root1.push_internal_level();
+    root1.borrow_mut().push(1, ());
+    let mut root1 = NodeRef::new_internal(root1.forget_type()).forget_type();
     let root2 = Root::new();
     root1.reborrow().assert_back_pointers();
     root2.reborrow().assert_back_pointers();