Auto merge of #83267 - ssomers:btree_prune_range_search_overlap, r=Mark-Simulacrum

BTree: no longer search arrays twice to check Ord

A possible addition to / partial replacement of #83147: no longer linearly search the upper bound of a range in the initial portion of the keys we already know are below the lower bound.
- Should be faster: fewer key comparisons at the cost of some instructions dealing with offsets
- Makes code a little more complicated.
- No longer detects ill-defined `Ord` implementations, but that wasn't a publicised feature, and was quite incomplete, and was only done in the `range` and `range_mut` methods.
r? `@Mark-Simulacrum`

3 files changed, 43 insertions, 27 deletions

diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs
index e636e490e1b..3a74b6a6fa8 100644
--- a/library/alloc/src/collections/btree/map/tests.rs
+++ b/library/alloc/src/collections/btree/map/tests.rs
@@ -776,7 +776,6 @@ fn test_range_backwards_4() {
 }
 
 #[test]
-#[should_panic]
 fn test_range_finding_ill_order_in_map() {
     let mut map = BTreeMap::new();
     map.insert(Cyclic3::B, ());
@@ -789,7 +788,6 @@ fn test_range_finding_ill_order_in_map() {
 }
 
 #[test]
-#[should_panic]
 fn test_range_finding_ill_order_in_range_ord() {
     // Has proper order the first time asked, then flips around.
     struct EvilTwin(i32);
diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs
index c2c99a9360c..4399feaccc9 100644
--- a/library/alloc/src/collections/btree/navigate.rs
+++ b/library/alloc/src/collections/btree/navigate.rs
@@ -29,11 +29,18 @@ impl<BorrowType, K, V> LeafRange<BorrowType, K, V> {
 
 impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
     /// Finds the distinct leaf edges delimiting a specified range in a tree.
-    /// Returns either a pair of different handles into the same tree or a pair
-    /// of empty options.
+    ///
+    /// If such distinct edges exist, returns them in ascending order, meaning
+    /// that a non-zero number of calls to `next_unchecked` on the `front` of
+    /// the result and/or calls to `next_back_unchecked` on the `back` of the
+    /// result will eventually reach the same edge.
+    ///
+    /// If there are no such edges, i.e., if the tree contains no key within
+    /// the range, returns a pair of empty options.
+    ///
     /// # Safety
-    /// Unless `BorrowType` is `Immut`, do not use the duplicate handles to
-    /// visit the same KV twice.
+    /// Unless `BorrowType` is `Immut`, do not use the handles to visit the same
+    /// KV twice.
     unsafe fn find_leaf_edges_spanning_range<Q: ?Sized, R>(
         self,
         range: R,
diff --git a/library/alloc/src/collections/btree/search.rs b/library/alloc/src/collections/btree/search.rs
index 5dc62d4ec70..5651a03c47a 100644
--- a/library/alloc/src/collections/btree/search.rs
+++ b/library/alloc/src/collections/btree/search.rs
@@ -68,13 +68,16 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
     /// of the range is different from the edge matching the upper bound, i.e.,
     /// the nearest node that has at least one key contained in the range.
     ///
-    /// If found, returns an `Ok` with that node, the pair of edge indices in it
-    /// delimiting the range, and the corresponding pair of bounds for
-    /// continuing the search in the child nodes, in case the node is internal.
+    /// If found, returns an `Ok` with that node, the strictly ascending pair of
+    /// edge indices in the node delimiting the range, and the corresponding
+    /// pair of bounds for continuing the search in the child nodes, in case
+    /// the node is internal.
     ///
     /// If not found, returns an `Err` with the leaf edge matching the entire
     /// range.
     ///
+    /// As a diagnostic service, panics if the range specifies impossible bounds.
+    ///
     /// The result is meaningful only if the tree is ordered by key.
     pub fn search_tree_for_bifurcation<'r, Q: ?Sized, R>(
         mut self,
@@ -112,10 +115,8 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
         let mut upper_bound = SearchBound::from_range(end);
         loop {
             let (lower_edge_idx, lower_child_bound) = self.find_lower_bound_index(lower_bound);
-            let (upper_edge_idx, upper_child_bound) = self.find_upper_bound_index(upper_bound);
-            if lower_edge_idx > upper_edge_idx {
-                panic!("Ord is ill-defined in BTreeMap range")
-            }
+            let (upper_edge_idx, upper_child_bound) =
+                unsafe { self.find_upper_bound_index(upper_bound, lower_edge_idx) };
             if lower_edge_idx < upper_edge_idx {
                 return Ok((
                     self,
@@ -125,6 +126,7 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
                     upper_child_bound,
                 ));
             }
+            debug_assert_eq!(lower_edge_idx, upper_edge_idx);
             let common_edge = unsafe { Handle::new_edge(self, lower_edge_idx) };
             match common_edge.force() {
                 Leaf(common_edge) => return Err(common_edge),
@@ -164,7 +166,7 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
         Q: ?Sized + Ord,
         K: Borrow<Q>,
     {
-        let (edge_idx, bound) = self.find_upper_bound_index(bound);
+        let (edge_idx, bound) = unsafe { self.find_upper_bound_index(bound, 0) };
         let edge = unsafe { Handle::new_edge(self, edge_idx) };
         (edge, bound)
     }
@@ -183,29 +185,33 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
         Q: Ord,
         K: Borrow<Q>,
     {
-        match self.find_key_index(key) {
+        match unsafe { self.find_key_index(key, 0) } {
             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, or the edge index where the key belongs.
+    /// exists, or the edge index where the key belongs, starting from a particular index.
     ///
     /// The result is meaningful only if the tree is ordered by key, like the tree
     /// in a `BTreeMap` is.
-    fn find_key_index<Q: ?Sized>(&self, key: &Q) -> IndexResult
+    ///
+    /// # Safety
+    /// `start_index` must be a valid edge index for the node.
+    unsafe fn find_key_index<Q: ?Sized>(&self, key: &Q, start_index: usize) -> IndexResult
     where
         Q: Ord,
         K: Borrow<Q>,
     {
         let node = self.reborrow();
         let keys = node.keys();
-        for (i, k) in keys.iter().enumerate() {
+        debug_assert!(start_index <= keys.len());
+        for (offset, k) in unsafe { keys.get_unchecked(start_index..) }.iter().enumerate() {
             match key.cmp(k.borrow()) {
                 Ordering::Greater => {}
-                Ordering::Equal => return IndexResult::KV(i),
-                Ordering::Less => return IndexResult::Edge(i),
+                Ordering::Equal => return IndexResult::KV(start_index + offset),
+                Ordering::Less => return IndexResult::Edge(start_index + offset),
             }
         }
         IndexResult::Edge(keys.len())
@@ -225,11 +231,11 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
         K: Borrow<Q>,
     {
         match bound {
-            Included(key) => match self.find_key_index(key) {
+            Included(key) => match unsafe { self.find_key_index(key, 0) } {
                 IndexResult::KV(idx) => (idx, AllExcluded),
                 IndexResult::Edge(idx) => (idx, bound),
             },
-            Excluded(key) => match self.find_key_index(key) {
+            Excluded(key) => match unsafe { self.find_key_index(key, 0) } {
                 IndexResult::KV(idx) => (idx + 1, AllIncluded),
                 IndexResult::Edge(idx) => (idx, bound),
             },
@@ -238,26 +244,31 @@ impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
         }
     }
 
-    /// Clone of `find_lower_bound_index` for the upper bound.
-    fn find_upper_bound_index<'r, Q>(
+    /// Mirror image of `find_lower_bound_index` for the upper bound,
+    /// with an additional parameter to skip part of the key array.
+    ///
+    /// # Safety
+    /// `start_index` must be a valid edge index for the node.
+    unsafe fn find_upper_bound_index<'r, Q>(
         &self,
         bound: SearchBound<&'r Q>,
+        start_index: usize,
     ) -> (usize, SearchBound<&'r Q>)
     where
         Q: ?Sized + Ord,
         K: Borrow<Q>,
     {
         match bound {
-            Included(key) => match self.find_key_index(key) {
+            Included(key) => match unsafe { self.find_key_index(key, start_index) } {
                 IndexResult::KV(idx) => (idx + 1, AllExcluded),
                 IndexResult::Edge(idx) => (idx, bound),
             },
-            Excluded(key) => match self.find_key_index(key) {
+            Excluded(key) => match unsafe { self.find_key_index(key, start_index) } {
                 IndexResult::KV(idx) => (idx, AllIncluded),
                 IndexResult::Edge(idx) => (idx, bound),
             },
             AllIncluded => (self.len(), AllIncluded),
-            AllExcluded => (0, AllExcluded),
+            AllExcluded => (start_index, AllExcluded),
         }
     }
 }