BTreeMap: convert search functions to methods

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

@@ -10,7 +10,7 @@ use core::ptr;
 
 use super::borrow::DormantMutRef;
 use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root};
-use super::search::{self, SearchResult::*};
+use super::search::SearchResult::*;
 use super::unwrap_unchecked;
 
 mod entry;
@@ -230,7 +230,7 @@ where
 
     fn get(&self, key: &Q) -> Option<&K> {
         let root_node = self.root.as_ref()?.reborrow();
-        match search::search_tree(root_node, key) {
+        match root_node.search_tree(key) {
             Found(handle) => Some(handle.into_kv().0),
             GoDown(_) => None,
         }
@@ -239,7 +239,7 @@ where
     fn take(&mut self, key: &Q) -> Option<K> {
         let (map, dormant_map) = DormantMutRef::new(self);
         let root_node = map.root.as_mut()?.borrow_mut();
-        match search::search_tree(root_node, key) {
+        match root_node.search_tree(key) {
             Found(handle) => {
                 Some(OccupiedEntry { handle, dormant_map, _marker: PhantomData }.remove_kv().0)
             }
@@ -250,7 +250,7 @@ where
     fn replace(&mut self, key: K) -> Option<K> {
         let (map, dormant_map) = DormantMutRef::new(self);
         let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut();
-        match search::search_tree::<marker::Mut<'_>, K, (), K>(root_node, &key) {
+        match root_node.search_tree::<K>(&key) {
             Found(mut kv) => Some(mem::replace(kv.key_mut(), key)),
             GoDown(handle) => {
                 VacantEntry { key, handle, dormant_map, _marker: PhantomData }.insert(());
@@ -526,7 +526,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
         Q: Ord,
     {
         let root_node = self.root.as_ref()?.reborrow();
-        match search::search_tree(root_node, key) {
+        match root_node.search_tree(key) {
             Found(handle) => Some(handle.into_kv().1),
             GoDown(_) => None,
         }
@@ -554,7 +554,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
         Q: Ord,
     {
         let root_node = self.root.as_ref()?.reborrow();
-        match search::search_tree(root_node, k) {
+        match root_node.search_tree(k) {
             Found(handle) => Some(handle.into_kv()),
             GoDown(_) => None,
         }
@@ -762,7 +762,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
         Q: Ord,
     {
         let root_node = self.root.as_mut()?.borrow_mut();
-        match search::search_tree(root_node, key) {
+        match root_node.search_tree(key) {
             Found(handle) => Some(handle.into_val_mut()),
             GoDown(_) => None,
         }
@@ -858,7 +858,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
     {
         let (map, dormant_map) = DormantMutRef::new(self);
         let root_node = map.root.as_mut()?.borrow_mut();
-        match search::search_tree(root_node, key) {
+        match root_node.search_tree(key) {
             Found(handle) => {
                 Some(OccupiedEntry { handle, dormant_map, _marker: PhantomData }.remove_entry())
             }
@@ -1051,7 +1051,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
         // FIXME(@porglezomp) Avoid allocating if we don't insert
         let (map, dormant_map) = DormantMutRef::new(self);
         let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut();
-        match search::search_tree(root_node, &key) {
+        match root_node.search_tree(&key) {
             Found(handle) => Occupied(OccupiedEntry { handle, dormant_map, _marker: PhantomData }),
             GoDown(handle) => {
                 Vacant(VacantEntry { key, handle, dormant_map, _marker: PhantomData })

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

@@ -5,7 +5,7 @@ use core::ops::RangeBounds;
 use core::ptr;
 
 use super::node::{marker, ForceResult::*, Handle, NodeRef};
-use super::search::{self, SearchResult};
+use super::search::SearchResult;
 use super::unwrap_unchecked;
 
 /// Finds the leaf edges delimiting a specified range in or underneath a node.
@@ -42,14 +42,14 @@ where
 
     loop {
         let front = match (min_found, range.start_bound()) {
-            (false, Included(key)) => match search::search_node(min_node, key) {
+            (false, Included(key)) => match min_node.search_node(key) {
                 SearchResult::Found(kv) => {
                     min_found = true;
                     kv.left_edge()
                 }
                 SearchResult::GoDown(edge) => edge,
             },
-            (false, Excluded(key)) => match search::search_node(min_node, key) {
+            (false, Excluded(key)) => match min_node.search_node(key) {
                 SearchResult::Found(kv) => {
                     min_found = true;
                     kv.right_edge()
@@ -62,14 +62,14 @@ where
         };
 
         let back = match (max_found, range.end_bound()) {
-            (false, Included(key)) => match search::search_node(max_node, key) {
+            (false, Included(key)) => match max_node.search_node(key) {
                 SearchResult::Found(kv) => {
                     max_found = true;
                     kv.right_edge()
                 }
                 SearchResult::GoDown(edge) => edge,
             },
-            (false, Excluded(key)) => match search::search_node(max_node, key) {
+            (false, Excluded(key)) => match max_node.search_node(key) {
                 SearchResult::Found(kv) => {
                     max_found = true;
                     kv.left_edge()

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

@@ -10,14 +10,20 @@ pub enum SearchResult<BorrowType, K, V, FoundType, GoDownType> {
     GoDown(Handle<NodeRef<BorrowType, K, V, GoDownType>, marker::Edge>),
 }
 
-/// Looks up a given key in a (sub)tree headed by the given node, recursively.
+pub enum IndexResult {
+    KV(usize),
+    Edge(usize),
+}
+
+impl<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.
     ///
     /// The result is meaningful only if the tree is ordered by key, like the tree
     /// in a `BTreeMap` is.
-pub fn search_tree<BorrowType, K, V, Q: ?Sized>(
+    pub fn search_tree<Q: ?Sized>(
-    mut node: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
+        mut self,
         key: &Q,
     ) -> SearchResult<BorrowType, K, V, marker::LeafOrInternal, marker::Leaf>
     where
@@ -25,49 +31,42 @@ where
         K: Borrow<Q>,
     {
         loop {
-        match search_node(node, key) {
+            self = match self.search_node(key) {
                 Found(handle) => return Found(handle),
                 GoDown(handle) => match handle.force() {
                     Leaf(leaf) => return GoDown(leaf),
-                Internal(internal) => {
+                    Internal(internal) => internal.descend(),
-                    node = internal.descend();
-                    continue;
-                }
                 },
             }
         }
     }
+}
 
-/// Looks up a given key in a given node, without recursion.
+impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
+    /// Looks up a given key in the node, without recursion.
     /// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
     /// returns a `GoDown` with the handle of the edge where the key might be found
     /// (if the node is internal) or where the key can be inserted.
     ///
     /// The result is meaningful only if the tree is ordered by key, like the tree
     /// in a `BTreeMap` is.
-pub fn search_node<BorrowType, K, V, Type, Q: ?Sized>(
+    pub fn search_node<Q: ?Sized>(self, key: &Q) -> SearchResult<BorrowType, K, V, Type, Type>
-    node: NodeRef<BorrowType, K, V, Type>,
-    key: &Q,
-) -> SearchResult<BorrowType, K, V, Type, Type>
     where
         Q: Ord,
         K: Borrow<Q>,
     {
-    match search_linear(&node, key) {
+        match self.find_index(key) {
-        (idx, true) => Found(unsafe { Handle::new_kv(node, idx) }),
+            IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, idx) }),
-        (idx, false) => GoDown(unsafe { Handle::new_edge(node, idx) }),
+            IndexResult::Edge(idx) => GoDown(unsafe { Handle::new_edge(self, idx) }),
         }
     }
 
     /// Returns either the KV index in the node at which the key (or an equivalent)
-/// exists and `true`, or the edge index where the key belongs and `false`.
+    /// exists, or the edge index where the key belongs.
     ///
     /// The result is meaningful only if the tree is ordered by key, like the tree
     /// in a `BTreeMap` is.
-fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
+    fn find_index<Q: ?Sized>(&self, key: &Q) -> IndexResult
-    node: &NodeRef<BorrowType, K, V, Type>,
-    key: &Q,
-) -> (usize, bool)
     where
         Q: Ord,
         K: Borrow<Q>,
@@ -75,14 +74,15 @@ where
         // This function is defined over all borrow types (immutable, mutable, owned).
         // Using `keys_at()` is fine here even if BorrowType is mutable, as all we return
         // is an index -- not a reference.
-    let len = node.len();
+        let len = self.len();
         for i in 0..len {
-        let k = unsafe { node.reborrow().key_at(i) };
+            let k = unsafe { self.reborrow().key_at(i) };
             match key.cmp(k.borrow()) {
                 Ordering::Greater => {}
-            Ordering::Equal => return (i, true),
+                Ordering::Equal => return IndexResult::KV(i),
-            Ordering::Less => return (i, false),
+                Ordering::Less => return IndexResult::Edge(i),
             }
         }
-    (len, false)
+        IndexResult::Edge(len)
+    }
 }

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

@@ -1,6 +1,6 @@
 use super::map::MIN_LEN;
 use super::node::{ForceResult::*, Root};
-use super::search::{search_node, SearchResult::*};
+use super::search::SearchResult::*;
 use core::borrow::Borrow;
 
 impl<K, V> Root<K, V> {
@@ -21,7 +21,7 @@ impl<K, V> Root<K, V> {
             let mut right_node = right_root.borrow_mut();
 
             loop {
-                let mut split_edge = match search_node(left_node, key) {
+                let mut split_edge = match left_node.search_node(key) {
                     // key is going to the right tree
                     Found(kv) => kv.left_edge(),
                     GoDown(edge) => edge,