diff --git a/library/alloc/src/collections/btree/borrow.rs b/library/alloc/src/collections/btree/borrow.rs new file mode 100644 index 0000000000000..5c95acfbe9c20 --- /dev/null +++ b/library/alloc/src/collections/btree/borrow.rs @@ -0,0 +1,44 @@ +use core::marker::PhantomData; +use core::ptr::NonNull; + +/// Models a reborrow of some unique reference, when you know that the reborrow +/// and all its descendants (i.e., all pointers and references derived from it) +/// will not be used any more at some point, after which you want to use the +/// original unique reference again. +/// +/// The borrow checker usually handles this stacking of borrows for you, but +/// some control flows that accomplish this stacking are too complicated for +/// the compiler to follow. A `DormantMutRef` allows you to check borrowing +/// yourself, while still expressing its stacked nature, and encapsulating +/// the raw pointer code needed to do this without undefined behavior. +pub struct DormantMutRef<'a, T> { + ptr: NonNull, + _marker: PhantomData<&'a mut T>, +} + +impl<'a, T> DormantMutRef<'a, T> { + /// Capture a unique borrow, and immediately reborrow it. For the compiler, + /// the lifetime of the new reference is the same as the lifetime of the + /// original reference, but you promise to use it for a shorter period. + pub fn new(t: &'a mut T) -> (&'a mut T, Self) { + let ptr = NonNull::from(t); + // SAFETY: we hold the borrow throughout 'a via `_marker`, and we expose + // only this reference, so it is unique. + let new_ref = unsafe { &mut *ptr.as_ptr() }; + (new_ref, Self { ptr, _marker: PhantomData }) + } + + /// 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(self) -> &'a mut T { + // SAFETY: our own safety conditions imply this reference is again unique. + unsafe { &mut *self.ptr.as_ptr() } + } +} + +#[cfg(test)] +mod tests; diff --git a/library/alloc/src/collections/btree/borrow/tests.rs b/library/alloc/src/collections/btree/borrow/tests.rs new file mode 100644 index 0000000000000..56a8434fc71e6 --- /dev/null +++ b/library/alloc/src/collections/btree/borrow/tests.rs @@ -0,0 +1,19 @@ +use super::DormantMutRef; + +#[test] +fn test_borrow() { + let mut data = 1; + let mut stack = vec![]; + let mut rr = &mut data; + for factor in [2, 3, 7].iter() { + let (r, dormant_r) = DormantMutRef::new(rr); + rr = r; + assert_eq!(*rr, 1); + stack.push((factor, dormant_r)); + } + while let Some((factor, dormant_r)) = stack.pop() { + let r = unsafe { dormant_r.awaken() }; + *r *= factor; + } + assert_eq!(data, 42); +} diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 674cdaaa012d9..aed898be08fb6 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -8,6 +8,7 @@ use core::mem::{self, ManuallyDrop}; use core::ops::{Index, RangeBounds}; use core::ptr; +use super::borrow::DormantMutRef; use super::node::{self, marker, ForceResult::*, Handle, InsertResult::*, NodeRef}; use super::search::{self, SearchResult::*}; use super::unwrap_unchecked; @@ -228,24 +229,23 @@ where } fn take(&mut self, key: &Q) -> Option { - let root_node = self.root.as_mut()?.node_as_mut(); + let (map, dormant_map) = DormantMutRef::new(self); + let root_node = map.root.as_mut()?.node_as_mut(); match search::search_tree(root_node, key) { - Found(handle) => Some( - OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData } - .remove_kv() - .0, - ), + Found(handle) => { + Some(OccupiedEntry { handle, dormant_map, _marker: PhantomData }.remove_kv().0) + } GoDown(_) => None, } } fn replace(&mut self, key: K) -> Option { - let root = Self::ensure_is_owned(&mut self.root); - match search::search_tree::, K, (), K>(root.node_as_mut(), &key) { + let (map, dormant_map) = DormantMutRef::new(self); + let root_node = Self::ensure_is_owned(&mut map.root).node_as_mut(); + match search::search_tree::, K, (), K>(root_node, &key) { Found(handle) => Some(mem::replace(handle.into_key_mut(), key)), GoDown(handle) => { - VacantEntry { key, handle, length: &mut self.length, _marker: PhantomData } - .insert(()); + VacantEntry { key, handle, dormant_map, _marker: PhantomData }.insert(()); None } } @@ -459,7 +459,7 @@ impl Debug for Entry<'_, K, V> { pub struct VacantEntry<'a, K: 'a, V: 'a> { key: K, handle: Handle, K, V, marker::Leaf>, marker::Edge>, - length: &'a mut usize, + dormant_map: DormantMutRef<'a, BTreeMap>, // Be invariant in `K` and `V` _marker: PhantomData<&'a mut (K, V)>, @@ -479,8 +479,7 @@ impl Debug for VacantEntry<'_, K, V> { #[stable(feature = "rust1", since = "1.0.0")] pub struct OccupiedEntry<'a, K: 'a, V: 'a> { handle: Handle, K, V, marker::LeafOrInternal>, marker::KV>, - - length: &'a mut usize, + dormant_map: DormantMutRef<'a, BTreeMap>, // Be invariant in `K` and `V` _marker: PhantomData<&'a mut (K, V)>, @@ -644,13 +643,10 @@ impl BTreeMap { /// ``` #[unstable(feature = "map_first_last", issue = "62924")] pub fn first_entry(&mut self) -> Option> { - let root_node = self.root.as_mut()?.node_as_mut(); + let (map, dormant_map) = DormantMutRef::new(self); + let root_node = map.root.as_mut()?.node_as_mut(); let kv = root_node.first_leaf_edge().right_kv().ok()?; - Some(OccupiedEntry { - handle: kv.forget_node_type(), - length: &mut self.length, - _marker: PhantomData, - }) + Some(OccupiedEntry { handle: kv.forget_node_type(), dormant_map, _marker: PhantomData }) } /// Removes and returns the first element in the map. @@ -721,13 +717,10 @@ impl BTreeMap { /// ``` #[unstable(feature = "map_first_last", issue = "62924")] pub fn last_entry(&mut self) -> Option> { - let root_node = self.root.as_mut()?.node_as_mut(); + let (map, dormant_map) = DormantMutRef::new(self); + let root_node = map.root.as_mut()?.node_as_mut(); let kv = root_node.last_leaf_edge().left_kv().ok()?; - Some(OccupiedEntry { - handle: kv.forget_node_type(), - length: &mut self.length, - _marker: PhantomData, - }) + Some(OccupiedEntry { handle: kv.forget_node_type(), dormant_map, _marker: PhantomData }) } /// Removes and returns the last element in the map. @@ -901,12 +894,12 @@ impl BTreeMap { K: Borrow, Q: Ord, { - let root_node = self.root.as_mut()?.node_as_mut(); + let (map, dormant_map) = DormantMutRef::new(self); + let root_node = map.root.as_mut()?.node_as_mut(); match search::search_tree(root_node, key) { - Found(handle) => Some( - OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData } - .remove_entry(), - ), + Found(handle) => { + Some(OccupiedEntry { handle, dormant_map, _marker: PhantomData }.remove_entry()) + } GoDown(_) => None, } } @@ -1073,13 +1066,12 @@ impl BTreeMap { #[stable(feature = "rust1", since = "1.0.0")] pub fn entry(&mut self, key: K) -> Entry<'_, K, V> { // FIXME(@porglezomp) Avoid allocating if we don't insert - let root = Self::ensure_is_owned(&mut self.root); - match search::search_tree(root.node_as_mut(), &key) { - Found(handle) => { - Occupied(OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData }) - } + let (map, dormant_map) = DormantMutRef::new(self); + let root_node = Self::ensure_is_owned(&mut map.root).node_as_mut(); + match search::search_tree(root_node, &key) { + Found(handle) => Occupied(OccupiedEntry { handle, dormant_map, _marker: PhantomData }), GoDown(handle) => { - Vacant(VacantEntry { key, handle, length: &mut self.length, _marker: PhantomData }) + Vacant(VacantEntry { key, handle, dormant_map, _marker: PhantomData }) } } } @@ -1284,9 +1276,17 @@ impl BTreeMap { } pub(super) fn drain_filter_inner(&mut self) -> DrainFilterInner<'_, K, V> { - let root_node = self.root.as_mut().map(|r| r.node_as_mut()); - let front = root_node.map(|rn| rn.first_leaf_edge()); - DrainFilterInner { length: &mut self.length, cur_leaf_edge: front } + if let Some(root) = self.root.as_mut() { + let (root, dormant_root) = DormantMutRef::new(root); + let front = root.node_as_mut().first_leaf_edge(); + DrainFilterInner { + length: &mut self.length, + dormant_root: Some(dormant_root), + cur_leaf_edge: Some(front), + } + } else { + DrainFilterInner { length: &mut self.length, dormant_root: None, cur_leaf_edge: None } + } } /// Creates a consuming iterator visiting all the keys, in sorted order. @@ -1671,6 +1671,9 @@ where /// of the predicate, thus also serving for BTreeSet::DrainFilter. pub(super) struct DrainFilterInner<'a, K: 'a, V: 'a> { length: &'a mut usize, + // dormant_root is wrapped in an Option to be able to `take` it. + dormant_root: Option>>, + // cur_leaf_edge is wrapped in an Option because maps without root lack a leaf edge. cur_leaf_edge: Option, K, V, marker::Leaf>, marker::Edge>>, } @@ -1728,7 +1731,13 @@ impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> { let (k, v) = kv.kv_mut(); if pred(k, v) { *self.length -= 1; - let (kv, pos) = kv.remove_kv_tracking(); + let (kv, pos) = kv.remove_kv_tracking(|| { + // SAFETY: we will touch the root in a way that will not + // invalidate the position returned. + let root = unsafe { self.dormant_root.take().unwrap().awaken() }; + root.pop_internal_level(); + self.dormant_root = Some(DormantMutRef::new(root).1); + }); self.cur_leaf_edge = Some(pos); return Some(kv); } @@ -2456,13 +2465,20 @@ impl<'a, K: Ord, V> VacantEntry<'a, K, V> { /// ``` #[stable(feature = "rust1", since = "1.0.0")] pub fn insert(self, value: V) -> &'a mut V { - *self.length += 1; - let out_ptr = match self.handle.insert_recursing(self.key, value) { - (Fit(_), val_ptr) => val_ptr, + (Fit(_), val_ptr) => { + // Safety: We have consumed self.handle and the handle returned. + let map = unsafe { self.dormant_map.awaken() }; + map.length += 1; + val_ptr + } (Split(ins), val_ptr) => { - let root = ins.left.into_root_mut(); + drop(ins.left); + // Safety: We have consumed self.handle and the reference returned. + let map = unsafe { self.dormant_map.awaken() }; + let root = map.root.as_mut().unwrap(); root.push_internal_level().push(ins.k, ins.v, ins.right); + map.length += 1; val_ptr } }; @@ -2636,9 +2652,15 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> { // Body of `remove_entry`, separate to keep the above implementations short. fn remove_kv(self) -> (K, V) { - *self.length -= 1; - - let (old_kv, _) = self.handle.remove_kv_tracking(); + let mut emptied_internal_root = false; + let (old_kv, _) = self.handle.remove_kv_tracking(|| emptied_internal_root = true); + // SAFETY: we consumed the intermediate root borrow, `self.handle`. + let map = unsafe { self.dormant_map.awaken() }; + map.length -= 1; + if emptied_internal_root { + let root = map.root.as_mut().unwrap(); + root.pop_internal_level(); + } old_kv } } @@ -2646,8 +2668,9 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> { impl<'a, K: 'a, V: 'a> Handle, K, V, marker::LeafOrInternal>, marker::KV> { /// Removes a key/value-pair from the map, and returns that pair, as well as /// the leaf edge corresponding to that former pair. - fn remove_kv_tracking( + fn remove_kv_tracking( self, + handle_emptied_internal_root: F, ) -> ((K, V), Handle, K, V, marker::Leaf>, marker::Edge>) { let (old_kv, mut pos, was_internal) = match self.force() { Leaf(leaf) => { @@ -2700,7 +2723,7 @@ impl<'a, K: 'a, V: 'a> Handle, K, V, marker::LeafOrInter // The parent that was just emptied must be the root, // because nodes on a lower level would not have been // left with a single child. - parent.into_root_mut().pop_internal_level(); + handle_emptied_internal_root(); break; } else { cur_node = parent.forget_type(); diff --git a/library/alloc/src/collections/btree/mod.rs b/library/alloc/src/collections/btree/mod.rs index 6c8a588eb58f3..1f85c44b185ed 100644 --- a/library/alloc/src/collections/btree/mod.rs +++ b/library/alloc/src/collections/btree/mod.rs @@ -1,3 +1,4 @@ +mod borrow; pub mod map; mod navigate; mod node; diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index f9de88e4c13c8..8832619a404cb 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -168,40 +168,20 @@ impl Root { /// Borrows and returns an immutable reference to the node owned by the root. pub fn node_as_ref(&self) -> NodeRef, K, V, marker::LeafOrInternal> { - NodeRef { - height: self.height, - node: self.node.as_ptr(), - root: ptr::null(), - _marker: PhantomData, - } + NodeRef { height: self.height, node: self.node.as_ptr(), _marker: PhantomData } } /// Borrows and returns a mutable reference to the node owned by the root. pub fn node_as_mut(&mut self) -> NodeRef, K, V, marker::LeafOrInternal> { - NodeRef { - height: self.height, - node: self.node.as_ptr(), - root: self as *mut _, - _marker: PhantomData, - } + NodeRef { height: self.height, node: self.node.as_ptr(), _marker: PhantomData } } pub fn node_as_valmut(&mut self) -> NodeRef, K, V, marker::LeafOrInternal> { - NodeRef { - height: self.height, - node: self.node.as_ptr(), - root: ptr::null(), - _marker: PhantomData, - } + NodeRef { height: self.height, node: self.node.as_ptr(), _marker: PhantomData } } pub fn into_ref(self) -> NodeRef { - NodeRef { - height: self.height, - node: self.node.as_ptr(), - root: ptr::null(), - _marker: PhantomData, - } + NodeRef { height: self.height, node: self.node.as_ptr(), _marker: PhantomData } } /// Adds a new internal node with a single edge, pointing to the previous root, and make that @@ -214,12 +194,8 @@ impl Root { self.node = BoxedNode::from_internal(new_node); self.height += 1; - let mut ret = NodeRef { - height: self.height, - node: self.node.as_ptr(), - root: self as *mut _, - _marker: PhantomData, - }; + let mut ret = + NodeRef { height: self.height, node: self.node.as_ptr(), _marker: PhantomData }; unsafe { ret.reborrow_mut().first_edge().correct_parent_link(); @@ -228,11 +204,15 @@ impl Root { ret } - /// Removes the internal root node, using its first child as the new root. - /// As it is intended only to be called when the root has only one child, - /// no cleanup is done on any of the other children of the root. + /// Removes the internal root node, using its first child as the new root node. + /// As it is intended only to be called when the root node has only one child, + /// no cleanup is done on any of the other children. /// This decreases the height by 1 and is the opposite of `push_internal_level`. - /// Panics if there is no internal level, i.e. if the root is a leaf. + /// + /// Requires exclusive access to the `Root` object but not to the root node; + /// it will not invalidate existing handles or references to the root node. + /// + /// Panics if there is no internal level, i.e., if the root node is a leaf. pub fn pop_internal_level(&mut self) { assert!(self.height > 0); @@ -276,8 +256,6 @@ pub struct NodeRef { /// The number of levels below the node. height: usize, node: NonNull>, - // `root` is null unless the borrow type is `Mut` - root: *const Root, _marker: PhantomData<(BorrowType, Type)>, } @@ -342,7 +320,7 @@ impl NodeRef { /// Temporarily takes out another, immutable reference to the same node. fn reborrow(&self) -> NodeRef, K, V, Type> { - NodeRef { height: self.height, node: self.node, root: self.root, _marker: PhantomData } + NodeRef { height: self.height, node: self.node, _marker: PhantomData } } /// Exposes the leaf "portion" of any leaf or internal node. @@ -390,12 +368,7 @@ impl NodeRef { let parent_as_leaf = unsafe { (*self.as_leaf_ptr()).parent as *const LeafNode }; if let Some(non_zero) = NonNull::new(parent_as_leaf as *mut _) { Ok(Handle { - node: NodeRef { - height: self.height + 1, - node: non_zero, - root: self.root, - _marker: PhantomData, - }, + node: NodeRef { height: self.height + 1, node: non_zero, _marker: PhantomData }, idx: unsafe { usize::from(*(*self.as_leaf_ptr()).parent_idx.as_ptr()) }, _marker: PhantomData, }) @@ -465,21 +438,21 @@ impl<'a, K, V, Type> NodeRef, K, V, Type> { /// Unsafely asserts to the compiler some static information about whether this /// node is a `Leaf` or an `Internal`. unsafe fn cast_unchecked(self) -> NodeRef, K, V, NewType> { - NodeRef { height: self.height, node: self.node, root: self.root, _marker: PhantomData } + NodeRef { height: self.height, node: self.node, _marker: PhantomData } } /// Temporarily takes out another, mutable reference to the same node. Beware, as /// this method is very dangerous, doubly so since it may not immediately appear /// dangerous. /// - /// Because mutable pointers can roam anywhere around the tree and can even (through - /// `into_root_mut`) mess with the root of the tree, the result of `reborrow_mut` - /// can easily be used to make the original mutable pointer dangling, or, in the case - /// of a reborrowed handle, out of bounds. - // FIXME(@gereeter) consider adding yet another type parameter to `NodeRef` that restricts - // the use of `ascend` and `into_root_mut` on reborrowed pointers, preventing this unsafety. + /// Because mutable pointers can roam anywhere around the tree, the returned + /// pointer can easily be used to make the original pointer dangling, out of + /// bounds, or invalid under stacked borrow rules. + // FIXME(@gereeter) consider adding yet another type parameter to `NodeRef` + // that restricts the use of navigation methods on reborrowed pointers, + // preventing this unsafety. unsafe fn reborrow_mut(&mut self) -> NodeRef, K, V, Type> { - NodeRef { height: self.height, node: self.node, root: self.root, _marker: PhantomData } + NodeRef { height: self.height, node: self.node, _marker: PhantomData } } /// Exposes the leaf "portion" of any leaf or internal node for writing. @@ -522,12 +495,6 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef, K, V, Type> { } impl<'a, K: 'a, V: 'a, Type> NodeRef, K, V, Type> { - /// Gets a mutable reference to the root itself. This is useful primarily when the - /// height of the tree needs to be adjusted. Never call this on a reborrowed pointer. - pub fn into_root_mut(self) -> &'a mut Root { - unsafe { &mut *(self.root as *mut Root) } - } - fn into_key_slice_mut(mut self) -> &'a mut [K] { // SAFETY: The keys of a node must always be initialized up to length. unsafe { @@ -757,14 +724,12 @@ impl NodeRef { ForceResult::Leaf(NodeRef { height: self.height, node: self.node, - root: self.root, _marker: PhantomData, }) } else { ForceResult::Internal(NodeRef { height: self.height, node: self.node, - root: self.root, _marker: PhantomData, }) } @@ -855,12 +820,7 @@ impl<'a, K, V, NodeType, HandleType> Handle, K, V, NodeT /// this method is very dangerous, doubly so since it may not immediately appear /// dangerous. /// - /// Because mutable pointers can roam anywhere around the tree and can even (through - /// `into_root_mut`) mess with the root of the tree, the result of `reborrow_mut` - /// can easily be used to make the original mutable pointer dangling, or, in the case - /// of a reborrowed handle, out of bounds. - // FIXME(@gereeter) consider adding yet another type parameter to `NodeRef` that restricts - // the use of `ascend` and `into_root_mut` on reborrowed pointers, preventing this unsafety. + /// For details, see `NodeRef::reborrow_mut`. pub unsafe fn reborrow_mut( &mut self, ) -> Handle, K, V, NodeType>, HandleType> { @@ -1106,7 +1066,6 @@ impl Handle, marke NodeRef { height: self.node.height - 1, node: unsafe { (&*(*internal_node).edges.get_unchecked(self.idx).as_ptr()).as_ptr() }, - root: self.node.root, _marker: PhantomData, } } @@ -1499,14 +1458,14 @@ unsafe fn move_edges( impl NodeRef { /// Removes any static information asserting that this node is a `Leaf` node. pub fn forget_type(self) -> NodeRef { - NodeRef { height: self.height, node: self.node, root: self.root, _marker: PhantomData } + NodeRef { height: self.height, node: self.node, _marker: PhantomData } } } impl NodeRef { /// Removes any static information asserting that this node is an `Internal` node. pub fn forget_type(self) -> NodeRef { - NodeRef { height: self.height, node: self.node, root: self.root, _marker: PhantomData } + NodeRef { height: self.height, node: self.node, _marker: PhantomData } } }