diff --git a/library/alloc/src/collections/btree/borrow.rs b/library/alloc/src/collections/btree/borrow.rs index 016f139a501..000b9bd0fab 100644 --- 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 index 370b58864af..e9366eec9ce 100644 --- 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 index 6912466448f..76b304e756f 100644 --- 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, 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 { + NodeRef { height: self.height, node: self.node, _marker: PhantomData } + } +} + +impl NodeRef { + /// 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, K, V, Type> { + NodeRef { height: self.height, node: self.node, _marker: PhantomData } + } } impl NodeRef { @@ -798,6 +816,25 @@ impl<'a, K, V, NodeType, HandleType> Handle, 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, HandleType> { + Handle { node: self.node.dormant(), idx: self.idx, _marker: PhantomData } + } +} + +impl Handle, 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, K, V, NodeType>, HandleType> { + Handle { node: unsafe { self.node.awaken() }, idx: self.idx, _marker: PhantomData } + } } impl Handle, marker::Edge> { @@ -851,9 +888,11 @@ impl<'a, K: 'a, V: 'a> Handle, 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, 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, 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,21 +910,26 @@ impl<'a, K: 'a, V: 'a> Handle, 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( - mut self, + self, key: K, val: V, alloc: A, - ) -> (Option>, *mut V) { + ) -> ( + Option>, + Handle, 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) }, @@ -893,8 +937,10 @@ impl<'a, K: 'a, V: 'a> Handle, K, V, marker::Leaf>, mark 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, K, V, marker::Leaf>, mark key: K, value: V, alloc: A, - ) -> (Option>, *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, 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 {}