Skip to content

Commit

Permalink
Auto merge of rust-lang#81073 - ssomers:btree_owned_root_vs_dying, r=…
Browse files Browse the repository at this point in the history 
…Mark-Simulacrum

BTreeMap: prevent tree from ever being owned by non-root node

This introduces a new marker type, `Dying`, which is used to note trees which are in the process of deallocation. On such trees, some fields may be in an inconsistent state as we are deallocating the tree. Unfortunately, there's not a great way to express conditional unsafety, so the methods for traversal can cause UB if not invoked correctly, but not marked as such. This is not a regression from the previous state, but rather isolates the destructive methods to solely being called on the dying state.
  • Loading branch information
@bors
bors committed Jan 29, 2021
2 parents 74500b9 + 417eefe commit c6bc462
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 c6bc462

Please sign in to comment.