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 NodeRef { } impl NodeRef { + fn new_internal(child: Root) -> 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>, 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 NodeRef { /// 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, 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, 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();