Skip to content

Commit

Permalink
BTreeMap: stop tree from being owned by non-root node
Browse files Browse the repository at this point in the history
  • Loading branch information
@ssomers
ssomers committed Jan 26, 2021
1 parent 7907345 commit 417eefe
Show file tree
Hide file tree
Showing 5 changed files with 73 additions and 32 deletions.
6 changes: 3 additions & 3 deletions library/alloc/src/collections/btree/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -300,8 +300,8 @@ pub struct IterMut<'a, K: 'a, V: 'a> {
/// [`into_iter`]: IntoIterator::into_iter
#[stable(feature = "rust1", since = "1.0.0")]
pub struct IntoIter<K, V> {
front: Option<Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>>,
back: Option<Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>>,
front: Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,
back: Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,
length: usize,
}

Expand Down Expand Up @@ -1364,7 +1364,7 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
fn into_iter(self) -> IntoIter<K, V> {
let mut me = ManuallyDrop::new(self);
if let Some(root) = me.root.take() {
let (f, b) = root.full_range();
let (f, b) = root.into_dying().full_range();

IntoIter { front: Some(f), back: Some(b), length: me.length }
} else {
Expand Down
30 changes: 18 additions & 12 deletions library/alloc/src/collections/btree/navigate.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use super::unwrap_unchecked;
///
/// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is.
fn range_search<BorrowType, K, V, Q, R>(
fn range_search<BorrowType: marker::BorrowType, K, V, Q, R>(
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
range: R,
Expand Down Expand Up @@ -105,7 +105,7 @@ where
}

/// Equivalent to `range_search(k, v, ..)` but without the `Ord` bound.
fn full_range<BorrowType, K, V>(
fn full_range<BorrowType: marker::BorrowType, K, V>(
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
) -> (
Expand Down Expand Up @@ -202,15 +202,15 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::ValMut<'a>, K, V, marker::LeafOrInternal>
}
}

impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
impl<K, V> NodeRef<marker::Dying, K, V, marker::LeafOrInternal> {
/// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree.
/// The results are non-unique references allowing massively destructive mutation, so must be
/// used with the utmost care.
pub fn full_range(
self,
) -> (
Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
) {
// We duplicate the root NodeRef here -- we will never access it in a way
// that overlaps references obtained from the root.
Expand All @@ -219,7 +219,9 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
}
}

impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
impl<BorrowType: marker::BorrowType, K, V>
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>
{
/// Given a leaf edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
/// on the right side, which is either in the same leaf node or in an ancestor node.
/// If the leaf edge is the last one in the tree, returns [`Result::Err`] with the root node.
Expand Down Expand Up @@ -263,7 +265,9 @@ impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::E
}
}

impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge> {
impl<BorrowType: marker::BorrowType, K, V>
Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge>
{
/// Given an internal edge handle, returns [`Result::Ok`] with a handle to the neighboring KV
/// on the right side, which is either in the same internal node or in an ancestor node.
/// If the internal edge is the last one in the tree, returns [`Result::Err`] with the root node.
Expand Down Expand Up @@ -297,8 +301,8 @@ macro_rules! def_next_kv_uncheched_dealloc {
/// - The node carrying the next KV returned must not have been deallocated by a
/// previous call on any handle obtained for this tree.
unsafe fn $name <K, V>(
leaf_edge: Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
) -> Handle<NodeRef<marker::Owned, K, V, marker::LeafOrInternal>, marker::KV> {
leaf_edge: Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
let mut edge = leaf_edge.forget_node_type();
loop {
edge = match edge.$adjacent_kv() {
Expand Down Expand Up @@ -378,7 +382,7 @@ impl<'a, K, V> Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::E
}
}

impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
/// Moves the leaf edge handle to the next leaf edge and returns the key and value
/// in between, deallocating any node left behind while leaving the corresponding
/// edge in its parent node dangling.
Expand Down Expand Up @@ -422,7 +426,7 @@ impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
}
}

impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
/// Returns the leftmost leaf edge in or underneath a node - in other words, the edge
/// you need first when navigating forward (or last when navigating backward).
#[inline]
Expand Down Expand Up @@ -503,7 +507,9 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>
}
}

impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
impl<BorrowType: marker::BorrowType, K, V>
Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>
{
/// Returns the leaf edge closest to a KV for forward navigation.
pub fn next_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
match self.force() {
Expand Down
63 changes: 49 additions & 14 deletions library/alloc/src/collections/btree/node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,8 +93,8 @@ struct InternalNode<K, V> {
data: LeafNode<K, V>,

/// The pointers to the children of this node. `len + 1` of these are considered
/// initialized and valid. Although during the process of `into_iter` or `drop`,
/// some pointers are dangling while others still need to be traversed.
/// initialized and valid, except that near the end, while the tree is held
/// through borrow type `Dying`, some of these pointers are dangling.
edges: [MaybeUninit<BoxedNode<K, V>>; 2 * B],
}

Expand All @@ -119,7 +119,7 @@ impl<K, V> InternalNode<K, V> {
/// is not a separate type and has no destructor.
type BoxedNode<K, V> = NonNull<LeafNode<K, V>>;

/// An owned tree.
/// The root node of an owned tree.
///
/// Note that this does not have a destructor, and must be cleaned up manually.
pub type Root<K, V> = NodeRef<marker::Owned, K, V, marker::LeafOrInternal>;
Expand Down Expand Up @@ -157,18 +157,23 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::Internal> {
}

impl<K, V, Type> NodeRef<marker::Owned, K, V, Type> {
/// Mutably borrows the owned node. Unlike `reborrow_mut`, this is safe,
/// because the return value cannot be used to destroy the node itself,
/// and there cannot be other references to the tree (except during the
/// process of `into_iter` or `drop`, but that is horrific already).
/// Mutably borrows the owned root node. Unlike `reborrow_mut`, this is safe
/// because the return value cannot be used to destroy the root, and there
/// cannot be other references to the tree.
pub fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, Type> {
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
}

/// Slightly mutably borrows the owned node.
/// Slightly mutably borrows the owned root node.
pub fn borrow_valmut(&mut self) -> NodeRef<marker::ValMut<'_>, K, V, Type> {
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
}

/// Irreversibly transistions to a reference that offers traversal,
/// destructive methods and little else.
pub fn into_dying(self) -> NodeRef<marker::Dying, K, V, Type> {
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
}
}

impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
Expand Down Expand Up @@ -196,8 +201,13 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {

let top = self.node;

let internal_node = NodeRef { height: self.height, node: top, _marker: PhantomData };
*self = internal_node.first_edge().descend();
// SAFETY: we asserted to be internal.
let internal_self = unsafe { self.borrow_mut().cast_to_internal_unchecked() };
// SAFETY: we borrowed `self` exclusively and its borrow type is exclusive.
let internal_node = unsafe { &mut *NodeRef::as_internal_ptr(&internal_self) };
// SAFETY: the first edge is always initialized.
self.node = unsafe { internal_node.edges[0].assume_init_read() };
self.height -= 1;
self.clear_parent_link();

unsafe {
Expand All @@ -224,6 +234,9 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
/// although insert methods allow a mutable pointer to a value to coexist.
/// - When this is `Owned`, the `NodeRef` acts roughly like `Box<Node>`,
/// but does not have a destructor, and must be cleaned up manually.
/// - When this is `Dying`, the `NodeRef` still acts roughly like `Box<Node>`,
/// but has methods to destroy the tree bit by bit, and ordinary methods,
/// while not marked as unsafe to call, can invoke UB if called incorrectly.
/// Since any `NodeRef` allows navigating through the tree, `BorrowType`
/// effectively applies to the entire tree, not just to the node itself.
/// - `K` and `V`: These are the types of keys and values stored in the nodes.
Expand Down Expand Up @@ -280,6 +293,7 @@ unsafe impl<'a, K: Sync + 'a, V: Sync + 'a, Type> Send for NodeRef<marker::Immut
unsafe impl<'a, K: Send + 'a, V: Send + 'a, Type> Send for NodeRef<marker::Mut<'a>, K, V, Type> {}
unsafe impl<'a, K: Send + 'a, V: Send + 'a, Type> Send for NodeRef<marker::ValMut<'a>, K, V, Type> {}
unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Owned, K, V, Type> {}
unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Dying, K, V, Type> {}

impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
/// Unpack a node reference that was packed as `NodeRef::parent`.
Expand Down Expand Up @@ -343,7 +357,7 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
}
}

impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
impl<BorrowType: marker::BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
/// Finds the parent of the current node. Returns `Ok(handle)` if the current
/// node actually has a parent, where `handle` points to the edge of the parent
/// that points to the current node. Returns `Err(self)` if the current node has
Expand All @@ -356,6 +370,7 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
pub fn ascend(
self,
) -> Result<Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge>, Self> {
assert!(BorrowType::PERMITS_TRAVERSAL);
// We need to use raw pointers to nodes because, if BorrowType is marker::ValMut,
// there might be outstanding mutable references to values that we must not invalidate.
let leaf_ptr: *const _ = Self::as_leaf_ptr(&self);
Expand Down Expand Up @@ -410,13 +425,13 @@ impl<'a, K: 'a, V: 'a, Type> NodeRef<marker::Immut<'a>, K, V, Type> {
}
}

impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
impl<K, V> NodeRef<marker::Dying, K, V, marker::LeafOrInternal> {
/// Similar to `ascend`, gets a reference to a node's parent node, but also
/// deallocates the current node in the process. This is unsafe because the
/// current node will still be accessible despite being deallocated.
pub unsafe fn deallocate_and_ascend(
self,
) -> Option<Handle<NodeRef<marker::Owned, K, V, marker::Internal>, marker::Edge>> {
) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::Internal>, marker::Edge>> {
let height = self.height;
let node = self.node;
let ret = self.ascend().ok();
Expand Down Expand Up @@ -951,14 +966,17 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
}
}

impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge> {
impl<BorrowType: marker::BorrowType, K, V>
Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge>
{
/// Finds the node pointed to by this edge.
///
/// The method name assumes you picture trees with the root node on top.
///
/// `edge.descend().ascend().unwrap()` and `node.ascend().unwrap().descend()` should
/// both, upon success, do nothing.
pub fn descend(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
assert!(BorrowType::PERMITS_TRAVERSAL);
// We need to use raw pointers to nodes because, if BorrowType is
// marker::ValMut, there might be outstanding mutable references to
// values that we must not invalidate. There's no worry accessing the
Expand Down Expand Up @@ -1596,10 +1614,27 @@ pub mod marker {
pub enum LeafOrInternal {}

pub enum Owned {}
pub enum Dying {}
pub struct Immut<'a>(PhantomData<&'a ()>);
pub struct Mut<'a>(PhantomData<&'a mut ()>);
pub struct ValMut<'a>(PhantomData<&'a mut ()>);

pub trait BorrowType {
// Whether node references of this borrow type allow traversing
// to other nodes in the tree.
const PERMITS_TRAVERSAL: bool = true;
}
impl BorrowType for Owned {
// Traversal isn't needede, it happens using the result of `borrow_mut`.
// By disabling traversal, and only creating new references to roots,
// we know that every reference of the `Owned` type is to a root node.
const PERMITS_TRAVERSAL: bool = false;
}
impl BorrowType for Dying {}
impl<'a> BorrowType for Immut<'a> {}
impl<'a> BorrowType for Mut<'a> {}
impl<'a> BorrowType for ValMut<'a> {}

pub enum KV {}
pub enum Edge {}
}
Expand Down
4 changes: 2 additions & 2 deletions library/alloc/src/collections/btree/node/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,8 +95,8 @@ fn test_partial_cmp_eq() {
assert_eq!(top_edge_1.partial_cmp(&top_edge_2), None);

root1.pop_internal_level();
unsafe { root1.deallocate_and_ascend() };
unsafe { root2.deallocate_and_ascend() };
unsafe { root1.into_dying().deallocate_and_ascend() };
unsafe { root2.into_dying().deallocate_and_ascend() };
}

#[test]
Expand Down
2 changes: 1 addition & 1 deletion library/alloc/src/collections/btree/search.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ pub enum IndexResult {
Edge(usize),
}

impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
/// Looks up a given key in a (sub)tree headed by the node, recursively.
/// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
/// returns a `GoDown` with the handle of the leaf edge where the key belongs.
Expand Down

0 comments on commit 417eefe

Please sign in to comment.