BTreeMap: move up reference to map's root from NodeRef

library/alloc/src/collections/btree/borrow.rs

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

library/alloc/src/collections/btree/borrow/tests.rs

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

library/alloc/src/collections/btree/map.rs

@@ -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,18 +2652,25 @@ 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
     }
 }
 
 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();

library/alloc/src/collections/btree/mod.rs

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