Use `map_while` instead of `take_while` + `map`

2 files changed, 6 insertions, 4 deletions

diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index 4467980054f..376513d3586 100644
--- a/compiler/rustc_data_structures/src/lib.rs
+++ b/compiler/rustc_data_structures/src/lib.rs
@@ -27,6 +27,8 @@
 #![feature(min_type_alias_impl_trait)]
 #![allow(rustc::default_hash_types)]
 #![deny(unaligned_references)]
+#![feature(iter_map_while)]
+#![feature(bool_to_option)]
 
 #[macro_use]
 extern crate tracing;
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 2054f5eebf5..e92db9ea128 100644
--- a/compiler/rustc_data_structures/src/sorted_map/index_map.rs
+++ b/compiler/rustc_data_structures/src/sorted_map/index_map.rs
@@ -86,10 +86,10 @@ impl<I: Idx, K: Ord, V> SortedIndexMultiMap<I, K, V> {
     /// insertion order.
     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()
-            .take_while(move |&&i| self.items[i].0.eq(&key))
-            .map(move |&idx| (idx, &self.items[idx].1))
+        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))
+        })
     }
 }