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,79 +10,79 @@ 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.
-/// 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>(
-    mut node: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
-    key: &Q,
-) -> SearchResult<BorrowType, K, V, marker::LeafOrInternal, marker::Leaf>
-where
-    Q: Ord,
-    K: Borrow<Q>,
-{
-    loop {
-        match search_node(node, key) {
-            Found(handle) => return Found(handle),
-            GoDown(handle) => match handle.force() {
-                Leaf(leaf) => return GoDown(leaf),
-                Internal(internal) => {
-                    node = internal.descend();
-                    continue;
-                }
-            },
+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<Q: ?Sized>(
+        mut self,
+        key: &Q,
+    ) -> SearchResult<BorrowType, K, V, marker::LeafOrInternal, marker::Leaf>
+    where
+        Q: Ord,
+        K: Borrow<Q>,
+    {
+        loop {
+            self = match self.search_node(key) {
+                Found(handle) => return Found(handle),
+                GoDown(handle) => match handle.force() {
+                    Leaf(leaf) => return GoDown(leaf),
+                    Internal(internal) => internal.descend(),
+                },
+            }
         }
     }
 }
 
-/// Looks up a given key in a given 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>(
-    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) {
-        (idx, true) => Found(unsafe { Handle::new_kv(node, idx) }),
-        (idx, false) => GoDown(unsafe { Handle::new_edge(node, idx) }),
+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<Q: ?Sized>(self, key: &Q) -> SearchResult<BorrowType, K, V, Type, Type>
+    where
+        Q: Ord,
+        K: Borrow<Q>,
+    {
+        match self.find_index(key) {
+            IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, 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`.
-///
-/// 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>(
-    node: &NodeRef<BorrowType, K, V, Type>,
-    key: &Q,
-) -> (usize, bool)
-where
-    Q: Ord,
-    K: Borrow<Q>,
-{
-    // 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();
-    for i in 0..len {
-        let k = unsafe { node.reborrow().key_at(i) };
-        match key.cmp(k.borrow()) {
-            Ordering::Greater => {}
-            Ordering::Equal => return (i, true),
-            Ordering::Less => return (i, false),
+    /// Returns either the KV index in the node at which the key (or an equivalent)
+    /// 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 find_index<Q: ?Sized>(&self, key: &Q) -> IndexResult
+    where
+        Q: Ord,
+        K: Borrow<Q>,
+    {
+        // 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 = self.len();
+        for i in 0..len {
+            let k = unsafe { self.reborrow().key_at(i) };
+            match key.cmp(k.borrow()) {
+                Ordering::Greater => {}
+                Ordering::Equal => return IndexResult::KV(i),
+                Ordering::Less => return IndexResult::Edge(i),
+            }
         }
+        IndexResult::Edge(len)
     }
-    (len, false)
 }

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,