Lift range_search up one level of abstraction

3 files changed, 45 insertions, 42 deletions

diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
index 8eabc177304..0b3f603686d 100644
--- a/src/liballoc/collections/btree/map.rs
+++ b/src/liballoc/collections/btree/map.rs
@@ -1861,65 +1861,51 @@ where
     let mut max_node = root2;
     let mut min_found = false;
     let mut max_found = false;
-    let mut diverged = false;
 
     loop {
-        let min_edge = match (min_found, range.start_bound()) {
-            (false, Included(key)) => match search::search_linear(&min_node, key) {
-                (i, true) => {
+        let front = match (min_found, range.start_bound()) {
+            (false, Included(key)) => match search::search_node(min_node, key) {
+                Found(kv) => {
                     min_found = true;
-                    i
+                    kv.left_edge()
                 }
-                (i, false) => i,
+                GoDown(edge) => edge,
             },
-            (false, Excluded(key)) => match search::search_linear(&min_node, key) {
-                (i, true) => {
+            (false, Excluded(key)) => match search::search_node(min_node, key) {
+                Found(kv) => {
                     min_found = true;
-                    i + 1
+                    kv.right_edge()
                 }
-                (i, false) => i,
+                GoDown(edge) => edge,
             },
-            (_, Unbounded) => 0,
-            (true, Included(_)) => min_node.len(),
-            (true, Excluded(_)) => 0,
+            (true, Included(_)) => min_node.last_edge(),
+            (true, Excluded(_)) => min_node.first_edge(),
+            (_, Unbounded) => min_node.first_edge(),
         };
 
-        let max_edge = match (max_found, range.end_bound()) {
-            (false, Included(key)) => match search::search_linear(&max_node, key) {
-                (i, true) => {
+        let back = match (max_found, range.end_bound()) {
+            (false, Included(key)) => match search::search_node(max_node, key) {
+                Found(kv) => {
                     max_found = true;
-                    i + 1
+                    kv.right_edge()
                 }
-                (i, false) => i,
+                GoDown(edge) => edge,
             },
-            (false, Excluded(key)) => match search::search_linear(&max_node, key) {
-                (i, true) => {
+            (false, Excluded(key)) => match search::search_node(max_node, key) {
+                Found(kv) => {
                     max_found = true;
-                    i
+                    kv.left_edge()
                 }
-                (i, false) => i,
+                GoDown(edge) => edge,
             },
-            (_, Unbounded) => max_node.len(),
-            (true, Included(_)) => 0,
-            (true, Excluded(_)) => max_node.len(),
+            (true, Included(_)) => max_node.first_edge(),
+            (true, Excluded(_)) => max_node.last_edge(),
+            (_, Unbounded) => max_node.last_edge(),
         };
 
-        if !diverged {
-            if max_edge < min_edge {
-                panic!("Ord is ill-defined in BTreeMap range")
-            }
-            if min_edge != max_edge {
-                diverged = true;
-            }
+        if front.partial_cmp(&back) == Some(Ordering::Greater) {
+            panic!("Ord is ill-defined in BTreeMap range");
         }
-
-        // Safety guarantee: `min_edge` is always in range for `min_node`, because
-        // `min_edge` is unconditionally calculated for each iteration's value of `min_node`,
-        // either (if not found) as the edge index returned by `search_linear`,
-        // or (if found) as the KV index returned by `search_linear`, possibly + 1.
-        // Likewise for `max_node` versus `max_edge`.
-        let front = unsafe { Handle::new_edge(min_node, min_edge) };
-        let back = unsafe { Handle::new_edge(max_node, max_edge) };
         match (front.force(), back.force()) {
             (Leaf(f), Leaf(b)) => {
                 return (f, b);
diff --git a/src/liballoc/collections/btree/node.rs b/src/liballoc/collections/btree/node.rs
index e4123c9a20b..abf926186e8 100644
--- a/src/liballoc/collections/btree/node.rs
+++ b/src/liballoc/collections/btree/node.rs
@@ -31,6 +31,7 @@
 // - A node of length `n` has `n` keys, `n` values, and (in an internal node) `n + 1` edges.
 //   This implies that even an empty internal node has at least one edge.
 
+use core::cmp::Ordering;
 use core::marker::PhantomData;
 use core::mem::{self, MaybeUninit};
 use core::ptr::{self, NonNull, Unique};
@@ -832,6 +833,14 @@ impl<BorrowType, K, V, NodeType, HandleType> PartialEq
     }
 }
 
+impl<BorrowType, K, V, NodeType, HandleType> PartialOrd
+    for Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType>
+{
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        if self.node.node == other.node.node { Some(self.idx.cmp(&other.idx)) } else { None }
+    }
+}
+
 impl<BorrowType, K, V, NodeType, HandleType>
     Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType>
 {
diff --git a/src/liballoc/collections/btree/search.rs b/src/liballoc/collections/btree/search.rs
index e680e364147..2ba5cebbdee 100644
--- a/src/liballoc/collections/btree/search.rs
+++ b/src/liballoc/collections/btree/search.rs
@@ -10,6 +10,10 @@ 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 possible leaf edge where the key
+/// belongs.
 pub fn search_tree<BorrowType, K, V, Q: ?Sized>(
     mut node: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
     key: &Q,
@@ -32,6 +36,10 @@ where
     }
 }
 
+/// 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 a leaf, a `GoDown` edge is not an actual edge but a possible edge.
 pub fn search_node<BorrowType, K, V, Type, Q: ?Sized>(
     node: NodeRef<BorrowType, K, V, Type>,
     key: &Q,
@@ -50,8 +58,8 @@ where
 /// or could exist, and whether it exists in the node itself. If it doesn't
 /// exist in the node itself, it may exist in the subtree with that index
 /// (if the node has subtrees). If the key doesn't exist in node or subtree,
-/// the returned index is the position or subtree to insert at.
-pub fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
+/// the returned index is the position or subtree where the key belongs.
+fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
     node: &NodeRef<BorrowType, K, V, Type>,
     key: &Q,
 ) -> (usize, bool)