BTreeMap: stop tree from being owned by non-root node

5 files changed, 73 insertions, 32 deletions

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index ecc28731874..79dc694e6be 100644
--- a/library/alloc/src/collections/btree/map.rs
+++ b/library/alloc/src/collections/btree/map.rs
@@ -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,
 }
 
@@ -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 {
diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs
index 17b4689e4a0..2773b427fb1 100644
--- a/library/alloc/src/collections/btree/navigate.rs
+++ b/library/alloc/src/collections/btree/navigate.rs
@@ -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,
@@ -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>,
 ) -> (
@@ -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.
@@ -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.
@@ -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.
@@ -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() {
@@ -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.
@@ -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]
@@ -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() {
diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
index 57fad4217b6..67c4f09306d 100644
--- a/library/alloc/src/collections/btree/node.rs
+++ b/library/alloc/src/collections/btree/node.rs
@@ -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],
 }
 
@@ -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>;
@@ -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> {
@@ -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 {
@@ -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.
@@ -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`.
@@ -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
@@ -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);
@@ -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();
@@ -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.
     ///
     /// 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
     /// 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
@@ -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 {}
 }
diff --git a/library/alloc/src/collections/btree/node/tests.rs b/library/alloc/src/collections/btree/node/tests.rs
index 11433cd845b..acb7210ca7c 100644
--- a/library/alloc/src/collections/btree/node/tests.rs
+++ b/library/alloc/src/collections/btree/node/tests.rs
@@ -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]
diff --git a/library/alloc/src/collections/btree/search.rs b/library/alloc/src/collections/btree/search.rs
index d8bebed758f..f87444b7cd3 100644
--- a/library/alloc/src/collections/btree/search.rs
+++ b/library/alloc/src/collections/btree/search.rs
@@ -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.