BTreeMap: better distinguish the root holder from the root node

2 files changed, 51 insertions, 42 deletions

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index 7e27aeb8539..e7d243bfcb0 100644
--- a/library/alloc/src/collections/btree/map.rs
+++ b/library/alloc/src/collections/btree/map.rs
@@ -2,14 +2,14 @@
 
 use core::borrow::Borrow;
 use core::cmp::Ordering;
-use core::fmt::Debug;
+use core::fmt::{self, Debug};
 use core::hash::{Hash, Hasher};
 use core::iter::{FromIterator, FusedIterator, Peekable};
 use core::marker::PhantomData;
 use core::mem::{self, ManuallyDrop};
 use core::ops::Bound::{Excluded, Included, Unbounded};
 use core::ops::{Index, RangeBounds};
-use core::{fmt, ptr};
+use core::ptr;
 
 use super::node::{self, marker, ForceResult::*, Handle, InsertResult::*, NodeRef};
 use super::search::{self, SearchResult::*};
@@ -154,7 +154,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
 
                     {
                         let root = out_tree.root.as_mut().unwrap(); // unwrap succeeds because we just wrapped
-                        let mut out_node = match root.as_mut().force() {
+                        let mut out_node = match root.node_as_mut().force() {
                             Leaf(leaf) => leaf,
                             Internal(_) => unreachable!(),
                         };
@@ -210,7 +210,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
             // Ord` constraint, which this method lacks.
             BTreeMap { root: None, length: 0 }
         } else {
-            clone_subtree(self.root.as_ref().unwrap().as_ref()) // unwrap succeeds because not empty
+            clone_subtree(self.root.as_ref().unwrap().node_as_ref()) // unwrap succeeds because not empty
         }
     }
 }
@@ -223,14 +223,16 @@ where
     type Key = K;
 
     fn get(&self, key: &Q) -> Option<&K> {
-        match search::search_tree(self.root.as_ref()?.as_ref(), key) {
+        let root_node = self.root.as_ref()?.node_as_ref();
+        match search::search_tree(root_node, key) {
             Found(handle) => Some(handle.into_kv().0),
             GoDown(_) => None,
         }
     }
 
     fn take(&mut self, key: &Q) -> Option<K> {
-        match search::search_tree(self.root.as_mut()?.as_mut(), key) {
+        let root_node = self.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()
@@ -242,7 +244,7 @@ where
 
     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.as_mut(), &key) {
+        match search::search_tree::<marker::Mut<'_>, K, (), K>(root.node_as_mut(), &key) {
             Found(handle) => Some(mem::replace(handle.into_kv_mut().0, key)),
             GoDown(handle) => {
                 VacantEntry { key, handle, length: &mut self.length, _marker: PhantomData }
@@ -565,7 +567,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
         K: Borrow<Q>,
         Q: Ord,
     {
-        match search::search_tree(self.root.as_ref()?.as_ref(), key) {
+        let root_node = self.root.as_ref()?.node_as_ref();
+        match search::search_tree(root_node, key) {
             Found(handle) => Some(handle.into_kv().1),
             GoDown(_) => None,
         }
@@ -592,7 +595,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
         K: Borrow<Q>,
         Q: Ord,
     {
-        match search::search_tree(self.root.as_ref()?.as_ref(), k) {
+        let root_node = self.root.as_ref()?.node_as_ref();
+        match search::search_tree(root_node, k) {
             Found(handle) => Some(handle.into_kv()),
             GoDown(_) => None,
         }
@@ -617,8 +621,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// ```
     #[unstable(feature = "map_first_last", issue = "62924")]
     pub fn first_key_value(&self) -> Option<(&K, &V)> {
-        let front = self.root.as_ref()?.as_ref().first_leaf_edge();
-        front.right_kv().ok().map(Handle::into_kv)
+        let root_node = self.root.as_ref()?.node_as_ref();
+        root_node.first_leaf_edge().right_kv().ok().map(Handle::into_kv)
     }
 
     /// Returns the first entry in the map for in-place manipulation.
@@ -643,8 +647,8 @@ 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 front = self.root.as_mut()?.as_mut().first_leaf_edge();
-        let kv = front.right_kv().ok()?;
+        let root_node = self.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,
@@ -694,8 +698,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
     /// ```
     #[unstable(feature = "map_first_last", issue = "62924")]
     pub fn last_key_value(&self) -> Option<(&K, &V)> {
-        let back = self.root.as_ref()?.as_ref().last_leaf_edge();
-        back.left_kv().ok().map(Handle::into_kv)
+        let root_node = self.root.as_ref()?.node_as_ref();
+        root_node.last_leaf_edge().left_kv().ok().map(Handle::into_kv)
     }
 
     /// Returns the last entry in the map for in-place manipulation.
@@ -720,8 +724,8 @@ 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 back = self.root.as_mut()?.as_mut().last_leaf_edge();
-        let kv = back.left_kv().ok()?;
+        let root_node = self.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,
@@ -805,7 +809,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
         K: Borrow<Q>,
         Q: Ord,
     {
-        match search::search_tree(self.root.as_mut()?.as_mut(), key) {
+        let root_node = self.root.as_mut()?.node_as_mut();
+        match search::search_tree(root_node, key) {
             Found(handle) => Some(handle.into_kv_mut().1),
             GoDown(_) => None,
         }
@@ -899,7 +904,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
         K: Borrow<Q>,
         Q: Ord,
     {
-        match search::search_tree(self.root.as_mut()?.as_mut(), key) {
+        let root_node = self.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(),
@@ -995,7 +1001,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
         R: RangeBounds<T>,
     {
         if let Some(root) = &self.root {
-            let (f, b) = range_search(root.as_ref(), range);
+            let (f, b) = range_search(root.node_as_ref(), range);
 
             Range { front: Some(f), back: Some(b) }
         } else {
@@ -1041,7 +1047,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
         R: RangeBounds<T>,
     {
         if let Some(root) = &mut self.root {
-            let (f, b) = range_search(root.as_mut(), range);
+            let (f, b) = range_search(root.node_as_mut(), range);
 
             RangeMut { front: Some(f), back: Some(b), _marker: PhantomData }
         } else {
@@ -1071,7 +1077,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
     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.as_mut(), &key) {
+        match search::search_tree(root.node_as_mut(), &key) {
             Found(handle) => {
                 Occupied(OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData })
             }
@@ -1083,7 +1089,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
 
     fn from_sorted_iter<I: Iterator<Item = (K, V)>>(&mut self, iter: I) {
         let root = Self::ensure_is_owned(&mut self.root);
-        let mut cur_node = root.as_mut().last_leaf_edge().into_node();
+        let mut cur_node = root.node_as_mut().last_leaf_edge().into_node();
         // Iterate through all key-value pairs, pushing them into nodes at the right level.
         for (key, value) in iter {
             // Try to push key-value pair into the current leaf node.
@@ -1133,7 +1139,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
 
     fn fix_right_edge(root: &mut node::Root<K, V>) {
         // Handle underfull nodes, start from the top.
-        let mut cur_node = root.as_mut();
+        let mut cur_node = root.node_as_mut();
         while let Internal(internal) = cur_node.force() {
             // Check if right-most child is underfull.
             let mut last_edge = internal.last_edge();
@@ -1201,8 +1207,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
         }
 
         {
-            let mut left_node = left_root.as_mut();
-            let mut right_node = right_root.as_mut();
+            let mut left_node = left_root.node_as_mut();
+            let mut right_node = right_root.node_as_mut();
 
             loop {
                 let mut split_edge = match search::search_node(left_node, key) {
@@ -1280,7 +1286,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
         DrainFilter { pred, inner: self.drain_filter_inner() }
     }
     pub(super) fn drain_filter_inner(&mut self) -> DrainFilterInner<'_, K, V> {
-        let front = self.root.as_mut().map(|r| r.as_mut().first_leaf_edge());
+        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,
@@ -1315,7 +1322,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
             res
         }
 
-        self.length = dfs(self.root.as_ref().unwrap().as_ref());
+        self.length = dfs(self.root.as_ref().unwrap().node_as_ref());
     }
 
     /// Creates a consuming iterator visiting all the keys, in sorted order.
@@ -2251,7 +2258,7 @@ impl<K, V> BTreeMap<K, V> {
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn iter(&self) -> Iter<'_, K, V> {
         if let Some(root) = &self.root {
-            let (f, b) = full_range_search(root.as_ref());
+            let (f, b) = full_range_search(root.node_as_ref());
 
             Iter { range: Range { front: Some(f), back: Some(b) }, length: self.length }
         } else {
@@ -2283,7 +2290,7 @@ impl<K, V> BTreeMap<K, V> {
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
         if let Some(root) = &mut self.root {
-            let (f, b) = full_range_search(root.as_mut());
+            let (f, b) = full_range_search(root.node_as_mut());
 
             IterMut {
                 range: RangeMut { front: Some(f), back: Some(b), _marker: PhantomData },
@@ -2895,7 +2902,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInter
 impl<K, V> node::Root<K, V> {
     /// Removes empty levels on the top, but keep an empty leaf if the entire tree is empty.
     fn fix_top(&mut self) {
-        while self.height() > 0 && self.as_ref().len() == 0 {
+        while self.height() > 0 && self.node_as_ref().len() == 0 {
             self.pop_internal_level();
         }
     }
@@ -2904,7 +2911,7 @@ impl<K, V> node::Root<K, V> {
         self.fix_top();
 
         {
-            let mut cur_node = self.as_mut();
+            let mut cur_node = self.node_as_mut();
 
             while let Internal(node) = cur_node.force() {
                 let mut last_kv = node.last_kv();
@@ -2930,7 +2937,7 @@ impl<K, V> node::Root<K, V> {
         self.fix_top();
 
         {
-            let mut cur_node = self.as_mut();
+            let mut cur_node = self.node_as_mut();
 
             while let Internal(node) = cur_node.force() {
                 let mut first_kv = node.first_kv();
diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
index 4e52c16d20d..f70869148d5 100644
--- a/library/alloc/src/collections/btree/node.rs
+++ b/library/alloc/src/collections/btree/node.rs
@@ -163,7 +163,8 @@ impl<K, V> Root<K, V> {
         Root { node: BoxedNode::from_leaf(Box::new(unsafe { LeafNode::new() })), height: 0 }
     }
 
-    pub fn as_ref(&self) -> NodeRef<marker::Immut<'_>, K, V, marker::LeafOrInternal> {
+    /// Borrows and returns an immutable reference to the node owned by the root.
+    pub fn node_as_ref(&self) -> NodeRef<marker::Immut<'_>, K, V, marker::LeafOrInternal> {
         NodeRef {
             height: self.height,
             node: self.node.as_ptr(),
@@ -172,7 +173,8 @@ impl<K, V> Root<K, V> {
         }
     }
 
-    pub fn as_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, marker::LeafOrInternal> {
+    /// Borrows and returns a mutable reference to the node owned by the root.
+    pub fn node_as_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, marker::LeafOrInternal> {
         NodeRef {
             height: self.height,
             node: self.node.as_ptr(),
@@ -226,12 +228,12 @@ impl<K, V> Root<K, V> {
 
         self.node = unsafe {
             BoxedNode::from_ptr(
-                self.as_mut().cast_unchecked::<marker::Internal>().first_edge().descend().node,
+                self.node_as_mut().cast_unchecked::<marker::Internal>().first_edge().descend().node,
             )
         };
         self.height -= 1;
         unsafe {
-            (*self.as_mut().as_leaf_mut()).parent = ptr::null();
+            (*self.node_as_mut().as_leaf_mut()).parent = ptr::null();
         }
 
         unsafe {
@@ -616,7 +618,7 @@ impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
                     let edge =
                         ptr::read(internal.as_internal().edges.get_unchecked(idx + 1).as_ptr());
                     let mut new_root = Root { node: edge, height: internal.height - 1 };
-                    (*new_root.as_mut().as_leaf_mut()).parent = ptr::null();
+                    (*new_root.node_as_mut().as_leaf_mut()).parent = ptr::null();
                     Some(new_root)
                 }
             };
@@ -648,7 +650,7 @@ impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
                     );
 
                     let mut new_root = Root { node: edge, height: internal.height - 1 };
-                    (*new_root.as_mut().as_leaf_mut()).parent = ptr::null();
+                    (*new_root.node_as_mut().as_leaf_mut()).parent = ptr::null();
 
                     for i in 0..old_len {
                         Handle::new_edge(internal.reborrow_mut(), i).correct_parent_link();
@@ -868,7 +870,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
             } else {
                 unsafe {
                     Handle::new_edge(
-                        right.as_mut().cast_unchecked::<marker::Leaf>(),
+                        right.node_as_mut().cast_unchecked::<marker::Leaf>(),
                         self.idx - (B + 1),
                     )
                     .insert_fit(key, val)
@@ -943,7 +945,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::
             } else {
                 unsafe {
                     Handle::new_edge(
-                        right.as_mut().cast_unchecked::<marker::Internal>(),
+                        right.node_as_mut().cast_unchecked::<marker::Internal>(),
                         self.idx - (B + 1),
                     )
                     .insert_fit(key, val, edge);
@@ -1117,7 +1119,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::
             let mut new_root = Root { node: BoxedNode::from_internal(new_node), height };
 
             for i in 0..(new_len + 1) {
-                Handle::new_edge(new_root.as_mut().cast_unchecked(), i).correct_parent_link();
+                Handle::new_edge(new_root.node_as_mut().cast_unchecked(), i).correct_parent_link();
             }
 
             (self.node, k, v, new_root)