]> git.lizzy.rs Git - rust.git/blobdiff - library/alloc/src/collections/btree/map.rs
Rollup merge of #76926 - ssomers:btree_cleanup_1, r=Mark-Simulacrum
[rust.git] / library / alloc / src / collections / btree / map.rs
index 674cdaaa012d9e637d87824f285a3c77a05a35f7..b4e9929af5ff2cd693e8430fa31d403635782c1b 100644 (file)
@@ -8,6 +8,7 @@
 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;
@@ -46,7 +47,6 @@
 /// any other key, as determined by the [`Ord`] trait, changes while it is in the map. This is
 /// normally only possible through [`Cell`], [`RefCell`], global state, I/O, or unsafe code.
 ///
-/// [`Ord`]: core::cmp::Ord
 /// [`Cell`]: core::cell::Cell
 /// [`RefCell`]: core::cell::RefCell
 ///
 /// }
 /// ```
 ///
-/// `BTreeMap` also implements an [`Entry API`](#method.entry), which allows
-/// for more complex methods of getting, setting, updating and removing keys and
-/// their values:
+/// `BTreeMap` also implements an [`Entry API`], which allows for more complex
+/// methods of getting, setting, updating and removing keys and their values:
+///
+/// [`Entry API`]: BTreeMap::entry
 ///
 /// ```
 /// use std::collections::BTreeMap;
@@ -228,24 +229,23 @@ fn get(&self, key: &Q) -> Option<&K> {
     }
 
     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
             }
         }
@@ -453,13 +453,11 @@ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 
 /// A view into a vacant entry in a `BTreeMap`.
 /// It is part of the [`Entry`] enum.
-///
-/// [`Entry`]: enum.Entry.html
 #[stable(feature = "rust1", since = "1.0.0")]
 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)>,
@@ -474,13 +472,10 @@ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 
 /// A view into an occupied entry in a `BTreeMap`.
 /// It is part of the [`Entry`] enum.
-///
-/// [`Entry`]: enum.Entry.html
 #[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 +639,10 @@ pub fn first_key_value(&self) -> Option<(&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 +713,10 @@ pub fn last_key_value(&self) -> Option<(&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.
@@ -822,7 +811,7 @@ pub fn get_mut<Q: ?Sized>(&mut self, key: &Q) -> Option<&mut V>
     /// types that can be `==` without being identical. See the [module-level
     /// documentation] for more.
     ///
-    /// [module-level documentation]: index.html#insert-and-complex-keys
+    /// [module-level documentation]: crate::collections#insert-and-complex-keys
     ///
     /// # Examples
     ///
@@ -901,12 +890,12 @@ pub fn remove_entry<Q: ?Sized>(&mut self, key: &Q) -> Option<(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 +1062,12 @@ pub fn range_mut<T: ?Sized, R>(&mut self, range: R) -> RangeMut<'_, 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 +1272,17 @@ pub fn drain_filter<F>(&mut self, pred: F) -> DrainFilter<'_, K, V, F>
     }
 
     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 +1667,9 @@ pub struct DrainFilter<'a, K, V, F>
 /// 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>>,
 }
 
@@ -1716,7 +1715,7 @@ impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> {
     /// Allow Debug implementations to predict the next element.
     pub(super) fn peek(&self) -> Option<(&K, &V)> {
         let edge = self.cur_leaf_edge.as_ref()?;
-        edge.reborrow().next_kv().ok().map(|kv| kv.into_kv())
+        edge.reborrow().next_kv().ok().map(Handle::into_kv)
     }
 
     /// Implementation of a typical `DrainFilter::next` method, given the predicate.
@@ -1728,7 +1727,13 @@ pub(super) fn next<F>(&mut self, pred: &mut F) -> Option<(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 +2461,20 @@ pub fn into_key(self) -> K {
     /// ```
     #[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
             }
         };
@@ -2538,7 +2550,7 @@ pub fn get(&self) -> &V {
     /// If you need a reference to the `OccupiedEntry` that may outlive the
     /// destruction of the `Entry` value, see [`into_mut`].
     ///
-    /// [`into_mut`]: #method.into_mut
+    /// [`into_mut`]: OccupiedEntry::into_mut
     ///
     /// # Examples
     ///
@@ -2568,7 +2580,7 @@ pub fn get_mut(&mut self) -> &mut V {
     ///
     /// If you need multiple references to the `OccupiedEntry`, see [`get_mut`].
     ///
-    /// [`get_mut`]: #method.get_mut
+    /// [`get_mut`]: OccupiedEntry::get_mut
     ///
     /// # Examples
     ///
@@ -2636,9 +2648,15 @@ pub fn remove(self) -> 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
     }
 }
@@ -2646,8 +2664,9 @@ fn remove_kv(self) -> (K, V) {
 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 +2719,7 @@ fn remove_kv_tracking(
                         // 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();