OpenE2K/rust
2
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

BTreeMap: pull the map's root out of NodeRef

Browse Source
This commit is contained in:
Stein Somers 2020-08-16 22:32:50 +02:00
parent d92155bf6a
commit 2b54ab880c
5 changed files with 163 additions and 117 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
library/alloc/src/collections/btree/borrow.rs Normal file
View File

@ -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<T>,
_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;

19
library/alloc/src/collections/btree/borrow/tests.rs Normal file
View File

@ -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);
}

121
library/alloc/src/collections/btree/map.rs
View File

@ -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<K> {
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<K> {
let root = Self::ensure_is_owned(&mut self.root);
match search::search_tree::<marker::Mut<'_>, 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::<marker::Mut<'_>, 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<K: Debug + Ord, V: Debug> Debug for Entry<'_, K, V> {
pub struct VacantEntry<'a, K: 'a, V: 'a> {
key: K,
handle: Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>,
length: &'a mut usize,
dormant_map: DormantMutRef<'a, BTreeMap<K, V>>,
// Be invariant in `K` and `V`
_marker: PhantomData<&'a mut (K, V)>,
@ -479,8 +479,7 @@ impl<K: Debug + Ord, V> Debug for VacantEntry<'_, K, V> {
#[stable(feature = "rust1", since = "1.0.0")]
pub struct OccupiedEntry<'a, K: 'a, V: 'a> {
handle: Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV>,
length: &'a mut usize,
dormant_map: DormantMutRef<'a, BTreeMap<K, V>>,
// Be invariant in `K` and `V`
_marker: PhantomData<&'a mut (K, V)>,
@ -644,13 +643,10 @@ impl<K: Ord, V> BTreeMap<K, V> {
/// ```
#[unstable(feature = "map_first_last", issue = "62924")]
pub fn first_entry(&mut self) -> Option<OccupiedEntry<'_, K, V>> {
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<K: Ord, V> BTreeMap<K, V> {
/// ```
#[unstable(feature = "map_first_last", issue = "62924")]
pub fn last_entry(&mut self) -> Option<OccupiedEntry<'_, K, V>> {
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<K: Ord, V> BTreeMap<K, V> {
K: Borrow<Q>,
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<K: Ord, V> BTreeMap<K, V> {
#[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<K: Ord, V> BTreeMap<K, V> {
}
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<DormantMutRef<'a, node::Root<K, V>>>,
// cur_leaf_edge is wrapped in an Option because maps without root lack a leaf edge.
cur_leaf_edge: Option<Handle<NodeRef<marker::Mut<'a>, 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<NodeRef<marker::Mut<'a>, 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<F: FnOnce()>(
self,
handle_emptied_internal_root: F,
) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, 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<NodeRef<marker::Mut<'a>, 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();

1
library/alloc/src/collections/btree/mod.rs
View File

@ -1,3 +1,4 @@
mod borrow;
pub mod map;
mod navigate;
mod node;

95
library/alloc/src/collections/btree/node.rs
View File

@ -168,40 +168,20 @@ impl<K, V> Root<K, V> {
/// Borrows and returns an immutable reference to the node owned by the root.
pub fn node_as_ref(&self) -> NodeRef<marker::Immut<'_>, 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<marker::Mut<'_>, 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<marker::ValMut<'_>, 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<marker::Owned, 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 }
}
/// Adds a new internal node with a single edge, pointing to the previous root, and make that
@ -214,12 +194,8 @@ impl<K, V> Root<K, V> {
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<K, V> Root<K, V> {
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<BorrowType, K, V, Type> {
/// The number of levels below the node.
height: usize,
node: NonNull<LeafNode<K, V>>,
// `root` is null unless the borrow type is `Mut`
root: *const Root<K, V>,
_marker: PhantomData<(BorrowType, Type)>,
}
@ -342,7 +320,7 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
/// Temporarily takes out another, immutable reference to the same node.
fn reborrow(&self) -> NodeRef<marker::Immut<'_>, 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<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
let parent_as_leaf = unsafe { (*self.as_leaf_ptr()).parent as *const LeafNode<K, V> };
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<marker::Mut<'a>, 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<NewType>(self) -> NodeRef<marker::Mut<'a>, 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<marker::Mut<'_>, 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<marker::Immut<'a>, K, V, Type> {
}
impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Mut<'a>, 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<K, V> {
unsafe { &mut *(self.root as *mut Root<K, V>) }
}
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<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
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<NodeRef<marker::Mut<'a>, 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<NodeRef<marker::Mut<'_>, K, V, NodeType>, HandleType> {
@ -1106,7 +1066,6 @@ impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, 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<K, V>(
impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Leaf> {
/// Removes any static information asserting that this node is a `Leaf` node.
pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
NodeRef { height: self.height, node: self.node, root: self.root, _marker: PhantomData }
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
}
}
impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
/// Removes any static information asserting that this node is an `Internal` node.
pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
NodeRef { height: self.height, node: self.node, root: self.root, _marker: PhantomData }
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
}
}