Auto merge of #86429 - JohnTitor:get-by-key-enum-part-2, r=oli-obk

Improve `get_by_key_enumerated` more

Follow-up of #86392, this applies the suggestions by `@m-ou-se.`

r? `@m-ou-se`

4 files changed, 24 insertions, 50 deletions

diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index 4467980054f..041d52aa20a 100644
--- a/compiler/rustc_data_structures/src/lib.rs
+++ b/compiler/rustc_data_structures/src/lib.rs
@@ -7,24 +7,26 @@
 //! This API is completely unstable and subject to change.
 
 #![doc(html_root_url = "https://doc.rust-lang.org/nightly/nightly-rustc/")]
+#![feature(allow_internal_unstable)]
 #![feature(array_windows)]
+#![feature(associated_type_bounds)]
+#![feature(auto_traits)]
+#![feature(bool_to_option)]
+#![feature(const_panic)]
 #![feature(control_flow_enum)]
+#![feature(core_intrinsics)]
+#![feature(extend_one)]
+#![feature(hash_raw_entry)]
 #![feature(in_band_lifetimes)]
+#![feature(iter_map_while)]
+#![feature(maybe_uninit_uninit_array)]
 #![feature(min_specialization)]
-#![feature(auto_traits)]
+#![feature(min_type_alias_impl_trait)]
+#![feature(new_uninit)]
 #![feature(nll)]
-#![feature(allow_internal_unstable)]
-#![feature(hash_raw_entry)]
-#![feature(core_intrinsics)]
+#![feature(once_cell)]
 #![feature(test)]
-#![feature(associated_type_bounds)]
 #![feature(thread_id_value)]
-#![feature(extend_one)]
-#![feature(const_panic)]
-#![feature(new_uninit)]
-#![feature(once_cell)]
-#![feature(maybe_uninit_uninit_array)]
-#![feature(min_type_alias_impl_trait)]
 #![allow(rustc::default_hash_types)]
 #![deny(unaligned_references)]
 
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 65689ab769c..e92db9ea128 100644
--- a/compiler/rustc_data_structures/src/sorted_map/index_map.rs
+++ b/compiler/rustc_data_structures/src/sorted_map/index_map.rs
@@ -1,6 +1,5 @@
 //! A variant of `SortedMap` that preserves insertion order.
 
-use std::borrow::Borrow;
 use std::hash::{Hash, Hasher};
 use std::iter::FromIterator;
 
@@ -76,11 +75,7 @@ impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
     ///
     /// If there are multiple items that are equivalent to `key`, they will be yielded in
     /// insertion order.
-    pub fn get_by_key<Q: 'a>(&'a self, key: &Q) -> impl 'a + Iterator<Item = &'a V>
-    where
-        Q: Ord + ?Sized,
-        K: Borrow<Q>,
-    {
+    pub fn get_by_key(&'a self, key: K) -> impl 'a + Iterator<Item = &'a V> {
         self.get_by_key_enumerated(key).map(|(_, v)| v)
     }
 
@@ -89,35 +84,12 @@ impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
     ///
     /// If there are multiple items that are equivalent to `key`, they will be yielded in
     /// insertion order.
-    pub fn get_by_key_enumerated<Q>(&self, key: &Q) -> impl '_ + Iterator<Item = (I, &V)>
-    where
-        Q: Ord + ?Sized,
-        K: Borrow<Q>,
-    {
-        match self.binary_search_idx(key) {
-            Err(_) => self.idxs_to_items_enumerated(&[]),
-
-            Ok(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])
-            }
-        }
-    }
-
-    fn binary_search_idx<Q>(&self, key: &Q) -> Result<usize, usize>
-    where
-        Q: Ord + ?Sized,
-        K: Borrow<Q>,
-    {
-        self.idx_sorted_by_item_key.binary_search_by(|&idx| self.items[idx].0.borrow().cmp(key))
-    }
-
-    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))
+    pub fn get_by_key_enumerated(&'a self, key: K) -> impl '_ + Iterator<Item = (I, &V)> {
+        let lower_bound = self.idx_sorted_by_item_key.partition_point(|&i| self.items[i].0 < key);
+        self.idx_sorted_by_item_key[lower_bound..].iter().map_while(move |&i| {
+            let (k, v) = &self.items[i];
+            (k == &key).then_some((i, v))
+        })
     }
 }
 
diff --git a/compiler/rustc_data_structures/src/sorted_map/tests.rs b/compiler/rustc_data_structures/src/sorted_map/tests.rs
index 7d91e1fdcef..1e977d709f1 100644
--- a/compiler/rustc_data_structures/src/sorted_map/tests.rs
+++ b/compiler/rustc_data_structures/src/sorted_map/tests.rs
@@ -14,11 +14,11 @@ fn test_sorted_index_multi_map() {
     }
 
     // `get_by_key` works.
-    assert_eq!(set.get_by_key(&3).copied().collect::<Vec<_>>(), vec![0]);
-    assert!(set.get_by_key(&4).next().is_none());
+    assert_eq!(set.get_by_key(3).copied().collect::<Vec<_>>(), vec![0]);
+    assert!(set.get_by_key(4).next().is_none());
 
     // `get_by_key` returns items in insertion order.
-    let twos: Vec<_> = set.get_by_key_enumerated(&2).collect();
+    let twos: Vec<_> = set.get_by_key_enumerated(2).collect();
     let idxs: Vec<usize> = twos.iter().map(|(i, _)| *i).collect();
     let values: Vec<usize> = twos.iter().map(|(_, &v)| v).collect();
 
diff --git a/compiler/rustc_middle/src/ty/assoc.rs b/compiler/rustc_middle/src/ty/assoc.rs
index d005f63ed43..2d177551664 100644
--- a/compiler/rustc_middle/src/ty/assoc.rs
+++ b/compiler/rustc_middle/src/ty/assoc.rs
@@ -124,7 +124,7 @@ impl<'tcx> AssocItems<'tcx> {
         &self,
         name: Symbol,
     ) -> impl '_ + Iterator<Item = &ty::AssocItem> {
-        self.items.get_by_key(&name).copied()
+        self.items.get_by_key(name).copied()
     }
 
     /// Returns an iterator over all associated items with the given name.