BTreeMap: Change internal insert function to return a handle

Amanieu · Amanieu · commit eb70c82d297c · 2023-02-01T00:08:59.000Z
This is a prerequisite for cursor support for `BTreeMap`.
diff --git a/library/alloc/src/collections/btree/borrow.rs b/library/alloc/src/collections/btree/borrow.rs
@@ -41,6 +41,28 @@ impl<'a, T> DormantMutRef<'a, T> {
         // SAFETY: our own safety conditions imply this reference is again unique.
         unsafe { &mut *self.ptr.as_ptr() }
     }
+
+    /// Borrows a new mutable reference from the unique borrow initially captured.
+    ///
+    /// # Safety
+    ///
+    /// The reborrow must have ended, i.e., the reference returned by `new` and
+    /// all pointers and references derived from it, must not be used anymore.
+    pub unsafe fn reborrow(&mut self) -> &'a mut T {
+        // SAFETY: our own safety conditions imply this reference is again unique.
+        unsafe { &mut *self.ptr.as_ptr() }
+    }
+
+    /// Borrows a new shared reference from the unique borrow initially captured.
+    ///
+    /// # Safety
+    ///
+    /// The reborrow must have ended, i.e., the reference returned by `new` and
+    /// all pointers and references derived from it, must not be used anymore.
+    pub unsafe fn reborrow_shared(&self) -> &'a T {
+        // SAFETY: our own safety conditions imply this reference is again unique.
+        unsafe { &*self.ptr.as_ptr() }
+    }
 }
 
 #[cfg(test)]
diff --git a/library/alloc/src/collections/btree/map/entry.rs b/library/alloc/src/collections/btree/map/entry.rs
@@ -347,7 +347,7 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> {
     /// assert_eq!(map["poneyland"], 37);
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
-    pub fn insert(self, value: V) -> &'a mut V {
+    pub fn insert(mut self, value: V) -> &'a mut V {
         let out_ptr = match self.handle {
             None => {
                 // SAFETY: There is no tree yet so no reference to it exists.
@@ -358,25 +358,27 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> {
                 map.length = 1;
                 val_ptr
             }
-            Some(handle) => match handle.insert_recursing(self.key, value, self.alloc.clone()) {
-                (None, val_ptr) => {
-                    // SAFETY: We have consumed self.handle.
-                    let map = unsafe { self.dormant_map.awaken() };
-                    map.length += 1;
-                    val_ptr
-                }
-                (Some(ins), val_ptr) => {
-                    drop(ins.left);
-                    // SAFETY: We have consumed self.handle and dropped the
-                    // remaining reference to the tree, ins.left.
-                    let map = unsafe { self.dormant_map.awaken() };
-                    let root = map.root.as_mut().unwrap(); // same as ins.left
-                    root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right);
-                    map.length += 1;
-                    val_ptr
-                }
-            },
+            Some(handle) => {
+                let new_handle =
+                    handle.insert_recursing(self.key, value, self.alloc.clone(), |ins| {
+                        drop(ins.left);
+                        // SAFETY: Pushing a new root node doesn't invalidate
+                        // handles to existing nodes.
+                        let map = unsafe { self.dormant_map.reborrow() };
+                        let root = map.root.as_mut().unwrap(); // same as ins.left
+                        root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right)
+                    });
+
+                // Get the pointer to the value
+                let val_ptr = new_handle.into_val_mut();
+
+                // SAFETY: We have consumed self.handle.
+                let map = unsafe { self.dormant_map.awaken() };
+                map.length += 1;
+                val_ptr
+            }
         };
+
         // Now that we have finished growing the tree using borrowed references,
         // dereference the pointer to a part of it, that we picked up along the way.
         unsafe { &mut *out_ptr }
diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
@@ -442,6 +442,24 @@ impl<'a, K, V, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
         // SAFETY: we have exclusive access to the entire node.
         unsafe { &mut *ptr }
     }
+
+    /// Returns a dormant copy of this node with its lifetime erased which can
+    /// be reawakened later.
+    pub fn dormant(&self) -> NodeRef<marker::DormantMut, K, V, Type> {
+        NodeRef { height: self.height, node: self.node, _marker: PhantomData }
+    }
+}
+
+impl<K, V, Type> NodeRef<marker::DormantMut, K, V, Type> {
+    /// Revert to the unique borrow initially captured.
+    ///
+    /// # Safety
+    ///
+    /// The reborrow must have ended, i.e., the reference returned by `new` and
+    /// all pointers and references derived from it, must not be used anymore.
+    pub unsafe fn awaken<'a>(self) -> NodeRef<marker::Mut<'a>, K, V, Type> {
+        NodeRef { height: self.height, node: self.node, _marker: PhantomData }
+    }
 }
 
 impl<K, V, Type> NodeRef<marker::Dying, K, V, Type> {
@@ -798,6 +816,25 @@ impl<'a, K, V, NodeType, HandleType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeT
         // We can't use Handle::new_kv or Handle::new_edge because we don't know our type
         Handle { node: unsafe { self.node.reborrow_mut() }, idx: self.idx, _marker: PhantomData }
     }
+
+    /// Returns a dormant copy of this handle which can be reawakened later.
+    ///
+    /// See [`DormantMutRef`] for more details.
+    pub fn dormant(&self) -> Handle<NodeRef<marker::DormantMut, K, V, NodeType>, HandleType> {
+        Handle { node: self.node.dormant(), idx: self.idx, _marker: PhantomData }
+    }
+}
+
+impl<K, V, NodeType, HandleType> Handle<NodeRef<marker::DormantMut, K, V, NodeType>, HandleType> {
+    /// Revert to the unique borrow initially captured.
+    ///
+    /// # Safety
+    ///
+    /// The reborrow must have ended, i.e., the reference returned by `new` and
+    /// all pointers and references derived from it, must not be used anymore.
+    pub unsafe fn awaken<'a>(self) -> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, HandleType> {
+        Handle { node: unsafe { self.node.awaken() }, idx: self.idx, _marker: PhantomData }
+    }
 }
 
 impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
@@ -851,9 +888,11 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
     /// Inserts a new key-value pair between the key-value pairs to the right and left of
     /// this edge. This method assumes that there is enough space in the node for the new
     /// pair to fit.
-    ///
-    /// The returned pointer points to the inserted value.
-    fn insert_fit(&mut self, key: K, val: V) -> *mut V {
+    unsafe fn insert_fit(
+        mut self,
+        key: K,
+        val: V,
+    ) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV> {
         debug_assert!(self.node.len() < CAPACITY);
         let new_len = self.node.len() + 1;
 
@@ -862,7 +901,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
             slice_insert(self.node.val_area_mut(..new_len), self.idx, val);
             *self.node.len_mut() = new_len as u16;
 
-            self.node.val_area_mut(self.idx).assume_init_mut()
+            Handle::new_kv(self.node, self.idx)
         }
     }
 }
@@ -871,30 +910,37 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
     /// Inserts a new key-value pair between the key-value pairs to the right and left of
     /// this edge. This method splits the node if there isn't enough room.
     ///
-    /// The returned pointer points to the inserted value.
+    /// Returns a dormant handle to the inserted node which can be reawakened
+    /// once splitting is complete.
     fn insert<A: Allocator + Clone>(
-        mut self,
+        self,
         key: K,
         val: V,
         alloc: A,
-    ) -> (Option<SplitResult<'a, K, V, marker::Leaf>>, *mut V) {
+    ) -> (
+        Option<SplitResult<'a, K, V, marker::Leaf>>,
+        Handle<NodeRef<marker::DormantMut, K, V, marker::Leaf>, marker::KV>,
+    ) {
         if self.node.len() < CAPACITY {
-            let val_ptr = self.insert_fit(key, val);
-            (None, val_ptr)
+            // SAFETY: There is enough space in the node for insertion.
+            let handle = unsafe { self.insert_fit(key, val) };
+            (None, handle.dormant())
         } else {
             let (middle_kv_idx, insertion) = splitpoint(self.idx);
             let middle = unsafe { Handle::new_kv(self.node, middle_kv_idx) };
             let mut result = middle.split(alloc);
-            let mut insertion_edge = match insertion {
+            let insertion_edge = match insertion {
                 LeftOrRight::Left(insert_idx) => unsafe {
                     Handle::new_edge(result.left.reborrow_mut(), insert_idx)
                 },
                 LeftOrRight::Right(insert_idx) => unsafe {
                     Handle::new_edge(result.right.borrow_mut(), insert_idx)
                 },
             };
-            let val_ptr = insertion_edge.insert_fit(key, val);
-            (Some(result), val_ptr)
+            // SAFETY: We just split the node, so there is enough space for
+            // insertion.
+            let handle = unsafe { insertion_edge.insert_fit(key, val).dormant() };
+            (Some(result), handle)
         }
     }
 }
@@ -976,21 +1022,31 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
         key: K,
         value: V,
         alloc: A,
-    ) -> (Option<SplitResult<'a, K, V, marker::LeafOrInternal>>, *mut V) {
-        let (mut split, val_ptr) = match self.insert(key, value, alloc.clone()) {
-            (None, val_ptr) => return (None, val_ptr),
-            (Some(split), val_ptr) => (split.forget_node_type(), val_ptr),
+        split_root: impl FnOnce(SplitResult<'a, K, V, marker::LeafOrInternal>),
+    ) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV> {
+        let (mut split, handle) = match self.insert(key, value, alloc.clone()) {
+            // SAFETY: we have finished splitting and can now re-awaken the
+            // handle to the inserted element.
+            (None, handle) => return unsafe { handle.awaken() },
+            (Some(split), handle) => (split.forget_node_type(), handle),
         };
 
         loop {
             split = match split.left.ascend() {
                 Ok(parent) => {
                     match parent.insert(split.kv.0, split.kv.1, split.right, alloc.clone()) {
-                        None => return (None, val_ptr),
+                        // SAFETY: we have finished splitting and can now re-awaken the
+                        // handle to the inserted element.
+                        None => return unsafe { handle.awaken() },
                         Some(split) => split.forget_node_type(),
                     }
                 }
-                Err(root) => return (Some(SplitResult { left: root, ..split }), val_ptr),
+                Err(root) => {
+                    split_root(SplitResult { left: root, ..split });
+                    // SAFETY: we have finished splitting and can now re-awaken the
+                    // handle to the inserted element.
+                    return unsafe { handle.awaken() };
+                }
             };
         }
     }
@@ -1667,6 +1723,7 @@ pub mod marker {
 
     pub enum Owned {}
     pub enum Dying {}
+    pub enum DormantMut {}
     pub struct Immut<'a>(PhantomData<&'a ()>);
     pub struct Mut<'a>(PhantomData<&'a mut ()>);
     pub struct ValMut<'a>(PhantomData<&'a mut ()>);
@@ -1688,6 +1745,7 @@ pub mod marker {
     impl<'a> BorrowType for Immut<'a> {}
     impl<'a> BorrowType for Mut<'a> {}
     impl<'a> BorrowType for ValMut<'a> {}
+    impl BorrowType for DormantMut {}
 
     pub enum KV {}
     pub enum Edge {}
