Auto merge of #81073 - ssomers:btree_owned_root_vs_dying, r=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.
This commit is contained in:
bors 2021-01-29 04:06:38 +00:00
commit c6bc46227a
5 changed files with 73 additions and 32 deletions

View File

@ -300,8 +300,8 @@ pub struct IterMut<'a, K: 'a, V: 'a> {
/// [`into_iter`]: IntoIterator::into_iter /// [`into_iter`]: IntoIterator::into_iter
#[stable(feature = "rust1", since = "1.0.0")] #[stable(feature = "rust1", since = "1.0.0")]
pub struct IntoIter<K, V> { pub struct IntoIter<K, V> {
front: 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::Owned, K, V, marker::Leaf>, marker::Edge>>, back: Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,
length: usize, length: usize,
} }
@ -1364,7 +1364,7 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
fn into_iter(self) -> IntoIter<K, V> { fn into_iter(self) -> IntoIter<K, V> {
let mut me = ManuallyDrop::new(self); let mut me = ManuallyDrop::new(self);
if let Some(root) = me.root.take() { 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 } IntoIter { front: Some(f), back: Some(b), length: me.length }
} else { } else {

View File

@ -12,7 +12,7 @@ use super::unwrap_unchecked;
/// ///
/// The result is meaningful only if the tree is ordered by key, like the tree /// The result is meaningful only if the tree is ordered by key, like the tree
/// in a `BTreeMap` is. /// 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>, root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>, root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
range: R, range: R,
@ -105,7 +105,7 @@ where
} }
/// Equivalent to `range_search(k, v, ..)` but without the `Ord` bound. /// 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>, root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>, root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
) -> ( ) -> (
@ -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. /// 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 /// The results are non-unique references allowing massively destructive mutation, so must be
/// used with the utmost care. /// used with the utmost care.
pub fn full_range( pub fn full_range(
self, self,
) -> ( ) -> (
Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>, Handle<NodeRef<marker::Dying, 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>,
) { ) {
// We duplicate the root NodeRef here -- we will never access it in a way // We duplicate the root NodeRef here -- we will never access it in a way
// that overlaps references obtained from the root. // that overlaps references obtained from the root.
@ -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 /// 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. /// 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. /// If the leaf edge is the last one in the tree, returns [`Result::Err`] with the root node.
@ -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 /// 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. /// 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. /// If the internal edge is the last one in the tree, returns [`Result::Err`] with the root node.
@ -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 /// - The node carrying the next KV returned must not have been deallocated by a
/// previous call on any handle obtained for this tree. /// previous call on any handle obtained for this tree.
unsafe fn $name <K, V>( unsafe fn $name <K, V>(
leaf_edge: Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>, leaf_edge: Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
) -> Handle<NodeRef<marker::Owned, K, V, marker::LeafOrInternal>, marker::KV> { ) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
let mut edge = leaf_edge.forget_node_type(); let mut edge = leaf_edge.forget_node_type();
loop { loop {
edge = match edge.$adjacent_kv() { edge = match edge.$adjacent_kv() {
@ -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 /// 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 /// in between, deallocating any node left behind while leaving the corresponding
/// edge in its parent node dangling. /// edge in its parent node dangling.
@ -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 /// 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). /// you need first when navigating forward (or last when navigating backward).
#[inline] #[inline]
@ -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. /// 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> { pub fn next_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
match self.force() { match self.force() {

View File

@ -93,8 +93,8 @@ struct InternalNode<K, V> {
data: LeafNode<K, V>, data: LeafNode<K, V>,
/// The pointers to the children of this node. `len + 1` of these are considered /// 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`, /// initialized and valid, except that near the end, while the tree is held
/// some pointers are dangling while others still need to be traversed. /// through borrow type `Dying`, some of these pointers are dangling.
edges: [MaybeUninit<BoxedNode<K, V>>; 2 * B], edges: [MaybeUninit<BoxedNode<K, V>>; 2 * B],
} }
@ -119,7 +119,7 @@ impl<K, V> InternalNode<K, V> {
/// is not a separate type and has no destructor. /// is not a separate type and has no destructor.
type BoxedNode<K, V> = NonNull<LeafNode<K, V>>; 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. /// 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>; pub type Root<K, V> = NodeRef<marker::Owned, K, V, marker::LeafOrInternal>;
@ -157,18 +157,23 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::Internal> {
} }
impl<K, V, Type> NodeRef<marker::Owned, K, V, Type> { impl<K, V, Type> NodeRef<marker::Owned, K, V, Type> {
/// Mutably borrows the owned node. Unlike `reborrow_mut`, this is safe, /// Mutably borrows the owned root node. Unlike `reborrow_mut`, this is safe
/// because the return value cannot be used to destroy the node itself, /// because the return value cannot be used to destroy the root, and there
/// and there cannot be other references to the tree (except during the /// cannot be other references to the tree.
/// process of `into_iter` or `drop`, but that is horrific already).
pub fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, Type> { pub fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, Type> {
NodeRef { height: self.height, node: self.node, _marker: PhantomData } 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> { pub fn borrow_valmut(&mut self) -> NodeRef<marker::ValMut<'_>, K, V, Type> {
NodeRef { height: self.height, node: self.node, _marker: PhantomData } 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> { impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
@ -196,8 +201,13 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
let top = self.node; let top = self.node;
let internal_node = NodeRef { height: self.height, node: top, _marker: PhantomData }; // SAFETY: we asserted to be internal.
*self = internal_node.first_edge().descend(); 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(); self.clear_parent_link();
unsafe { unsafe {
@ -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. /// although insert methods allow a mutable pointer to a value to coexist.
/// - When this is `Owned`, the `NodeRef` acts roughly like `Box<Node>`, /// - When this is `Owned`, the `NodeRef` acts roughly like `Box<Node>`,
/// but does not have a destructor, and must be cleaned up manually. /// 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` /// Since any `NodeRef` allows navigating through the tree, `BorrowType`
/// effectively applies to the entire tree, not just to the node itself. /// 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. /// - `K` and `V`: These are the types of keys and values stored in the nodes.
@ -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::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<'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::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> { impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
/// Unpack a node reference that was packed as `NodeRef::parent`. /// Unpack a node reference that was packed as `NodeRef::parent`.
@ -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 /// 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 /// 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 /// that points to the current node. Returns `Err(self)` if the current node has
@ -356,6 +370,7 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
pub fn ascend( pub fn ascend(
self, self,
) -> Result<Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge>, 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, // 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 might be outstanding mutable references to values that we must not invalidate.
let leaf_ptr: *const _ = Self::as_leaf_ptr(&self); let leaf_ptr: *const _ = Self::as_leaf_ptr(&self);
@ -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 /// 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 /// deallocates the current node in the process. This is unsafe because the
/// current node will still be accessible despite being deallocated. /// current node will still be accessible despite being deallocated.
pub unsafe fn deallocate_and_ascend( pub unsafe fn deallocate_and_ascend(
self, 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 height = self.height;
let node = self.node; let node = self.node;
let ret = self.ascend().ok(); let ret = self.ascend().ok();
@ -951,7 +966,9 @@ 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. /// Finds the node pointed to by this edge.
/// ///
/// The method name assumes you picture trees with the root node on top. /// The method name assumes you picture trees with the root node on top.
@ -959,6 +976,7 @@ impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marke
/// `edge.descend().ascend().unwrap()` and `node.ascend().unwrap().descend()` should /// `edge.descend().ascend().unwrap()` and `node.ascend().unwrap().descend()` should
/// both, upon success, do nothing. /// both, upon success, do nothing.
pub fn descend(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> { 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 // We need to use raw pointers to nodes because, if BorrowType is
// marker::ValMut, there might be outstanding mutable references to // marker::ValMut, there might be outstanding mutable references to
// values that we must not invalidate. There's no worry accessing the // values that we must not invalidate. There's no worry accessing the
@ -1596,10 +1614,27 @@ pub mod marker {
pub enum LeafOrInternal {} pub enum LeafOrInternal {}
pub enum Owned {} pub enum Owned {}
pub enum Dying {}
pub struct Immut<'a>(PhantomData<&'a ()>); pub struct Immut<'a>(PhantomData<&'a ()>);
pub struct Mut<'a>(PhantomData<&'a mut ()>); pub struct Mut<'a>(PhantomData<&'a mut ()>);
pub struct ValMut<'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 KV {}
pub enum Edge {} pub enum Edge {}
} }

View File

@ -95,8 +95,8 @@ fn test_partial_cmp_eq() {
assert_eq!(top_edge_1.partial_cmp(&top_edge_2), None); assert_eq!(top_edge_1.partial_cmp(&top_edge_2), None);
root1.pop_internal_level(); root1.pop_internal_level();
unsafe { root1.deallocate_and_ascend() }; unsafe { root1.into_dying().deallocate_and_ascend() };
unsafe { root2.deallocate_and_ascend() }; unsafe { root2.into_dying().deallocate_and_ascend() };
} }
#[test] #[test]

View File

@ -15,7 +15,7 @@ pub enum IndexResult {
Edge(usize), 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. /// 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 `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. /// returns a `GoDown` with the handle of the leaf edge where the key belongs.