Auto merge of #86392 - JohnTitor:use-partition-point, r=petrochenkov

Prefer `partition_point` to look up assoc items

Since we now have `partition_point` (instead of `equal_range`), I think it's worth trying to use it instead of manually finding it.
`partition_point` uses `binary_search_by` internally (#85406) and its performance has been improved (#74024), so I guess this will make a performance difference.

1 files changed, 5 insertions, 47 deletions

diff --git a/compiler/rustc_data_structures/src/sorted_map/index_map.rs b/compiler/rustc_data_structures/src/sorted_map/index_map.rs
index 01cd1cec924..65689ab769c 100644
--- a/compiler/rustc_data_structures/src/sorted_map/index_map.rs
+++ b/compiler/rustc_data_structures/src/sorted_map/index_map.rs
@@ -94,13 +94,15 @@ impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
         Q: Ord + ?Sized,
         K: Borrow<Q>,
     {
-        // FIXME: This should be in the standard library as `equal_range`. See rust-lang/rfcs#2184.
         match self.binary_search_idx(key) {
             Err(_) => self.idxs_to_items_enumerated(&[]),
 
             Ok(idx) => {
-                let start = self.find_lower_bound(key, idx);
-                let end = self.find_upper_bound(key, idx);
+                let start = self.idx_sorted_by_item_key[..idx]
+                    .partition_point(|&i| self.items[i].0.borrow() != key);
+                let end = idx
+                    + self.idx_sorted_by_item_key[idx..]
+                        .partition_point(|&i| self.items[i].0.borrow() == key);
                 self.idxs_to_items_enumerated(&self.idx_sorted_by_item_key[start..end])
             }
         }
@@ -114,50 +116,6 @@ impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
         self.idx_sorted_by_item_key.binary_search_by(|&idx| self.items[idx].0.borrow().cmp(key))
     }
 
-    /// Returns the index into the `idx_sorted_by_item_key` array of the first item equal to
-    /// `key`.
-    ///
-    /// `initial` must be an index into that same array for an item that is equal to `key`.
-    fn find_lower_bound<Q>(&self, key: &Q, initial: usize) -> usize
-    where
-        Q: Ord + ?Sized,
-        K: Borrow<Q>,
-    {
-        debug_assert!(self.items[self.idx_sorted_by_item_key[initial]].0.borrow() == key);
-
-        // FIXME: At present, this uses linear search, meaning lookup is only `O(log n)` if duplicate
-        // entries are rare. It would be better to start with a linear search for the common case but
-        // fall back to an exponential search if many duplicates are found. This applies to
-        // `upper_bound` as well.
-        let mut start = initial;
-        while start != 0 && self.items[self.idx_sorted_by_item_key[start - 1]].0.borrow() == key {
-            start -= 1;
-        }
-
-        start
-    }
-
-    /// Returns the index into the `idx_sorted_by_item_key` array of the first item greater than
-    /// `key`, or `self.len()` if no such item exists.
-    ///
-    /// `initial` must be an index into that same array for an item that is equal to `key`.
-    fn find_upper_bound<Q>(&self, key: &Q, initial: usize) -> usize
-    where
-        Q: Ord + ?Sized,
-        K: Borrow<Q>,
-    {
-        debug_assert!(self.items[self.idx_sorted_by_item_key[initial]].0.borrow() == key);
-
-        // See the FIXME for `find_lower_bound`.
-        let mut end = initial + 1;
-        let len = self.items.len();
-        while end < len && self.items[self.idx_sorted_by_item_key[end]].0.borrow() == key {
-            end += 1;
-        }
-
-        end
-    }
-
     fn idxs_to_items_enumerated(&'a self, idxs: &'a [I]) -> impl 'a + Iterator<Item = (I, &'a V)> {
         idxs.iter().map(move |&idx| (idx, &self.items[idx].1))
     }