Auto merge of rust-lang#68499 - ssomers:btree_search_tidying, r=Mark-Simulacrum

BtreeMap range_search spruced up

rust-lang#39457 created a lower level entry point for `range_search` to operate on, but it's really not hard to move it up a level of abstraction, making it somewhat shorter and reusing existing unsafe code (`new_edge` is unsafe although it is currently not tagged as such).

Benchmark added. Comparison says there's no real difference:
```
>cargo benchcmp old3.txt new3.txt --threshold 5
 name                                           old3.txt ns/iter  new3.txt ns/iter  diff ns/iter   diff %  speedup
 btree::map::find_seq_100                       19                21                           2   10.53%   x 0.90
 btree::map::range_excluded_unbounded           3,117             2,838                     -279   -8.95%   x 1.10
 btree::map::range_included_unbounded           1,768             1,871                      103    5.83%   x 0.94
 btree::set::intersection_10k_neg_vs_10k_pos    35                37                           2    5.71%   x 0.95
 btree::set::intersection_staggered_100_vs_10k  2,488             2,314                     -174   -6.99%   x 1.08
 btree::set::is_subset_10k_vs_100               3                 2                           -1  -33.33%   x 1.50
```

r? @Mark-Simulacrum

Author: bors

src/liballoc/benches/btree/map.rs

@@ -1,5 +1,6 @@
 use std::collections::BTreeMap;
 use std::iter::Iterator;
+use std::ops::Bound::{Excluded, Unbounded};
 use std::vec::Vec;
 
 use rand::{seq::SliceRandom, thread_rng, Rng};
@@ -200,3 +201,58 @@ pub fn first_and_last_100(b: &mut Bencher) {
 pub fn first_and_last_10k(b: &mut Bencher) {
     bench_first_and_last(b, 10_000);
 }
+
+#[bench]
+pub fn range_excluded_excluded(b: &mut Bencher) {
+    let size = 144;
+    let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
+    b.iter(|| {
+        for first in 0..size {
+            for last in first + 1..size {
+                black_box(map.range((Excluded(first), Excluded(last))));
+            }
+        }
+    });
+}
+
+#[bench]
+pub fn range_excluded_unbounded(b: &mut Bencher) {
+    let size = 144;
+    let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
+    b.iter(|| {
+        for first in 0..size {
+            black_box(map.range((Excluded(first), Unbounded)));
+        }
+    });
+}
+
+#[bench]
+pub fn range_included_included(b: &mut Bencher) {
+    let size = 144;
+    let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
+    b.iter(|| {
+        for first in 0..size {
+            for last in first..size {
+                black_box(map.range(first..=last));
+            }
+        }
+    });
+}
+
+#[bench]
+pub fn range_included_unbounded(b: &mut Bencher) {
+    let size = 144;
+    let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
+    b.iter(|| {
+        for first in 0..size {
+            black_box(map.range(first..));
+        }
+    });
+}
+
+#[bench]
+pub fn range_unbounded_unbounded(b: &mut Bencher) {
+    let size = 144;
+    let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
+    b.iter(|| map.range(..));
+}

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);

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>
 {

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)