Move iter_results to dyn FnMut rather than a generic

This means that we're no longer generating the iteration/locking code for each
invocation site of iter_results, rather just once per query.

This is a 15% win in instruction counts when compiling the rustc_query_impl crate.

5 files changed, 55 insertions, 49 deletions

diff --git a/compiler/rustc_middle/src/ty/query/on_disk_cache.rs b/compiler/rustc_middle/src/ty/query/on_disk_cache.rs
index 416199b3840..43b775f8083 100644
--- a/compiler/rustc_middle/src/ty/query/on_disk_cache.rs
+++ b/compiler/rustc_middle/src/ty/query/on_disk_cache.rs
@@ -1185,20 +1185,27 @@ where
 
     assert!(Q::query_state(tcx).all_inactive());
     let cache = Q::query_cache(tcx);
-    cache.iter_results(|results| {
-        for (key, value, dep_node) in results {
-            if Q::cache_on_disk(tcx, &key, Some(value)) {
-                let dep_node = SerializedDepNodeIndex::new(dep_node.index());
-
-                // Record position of the cache entry.
-                query_result_index
-                    .push((dep_node, AbsoluteBytePos::new(encoder.encoder.position())));
-
-                // Encode the type check tables with the `SerializedDepNodeIndex`
-                // as tag.
-                encoder.encode_tagged(dep_node, value)?;
+    let mut res = Ok(());
+    cache.iter_results(&mut |key, value, dep_node| {
+        if res.is_err() {
+            return;
+        }
+        if Q::cache_on_disk(tcx, &key, Some(value)) {
+            let dep_node = SerializedDepNodeIndex::new(dep_node.index());
+
+            // Record position of the cache entry.
+            query_result_index.push((dep_node, AbsoluteBytePos::new(encoder.encoder.position())));
+
+            // Encode the type check tables with the `SerializedDepNodeIndex`
+            // as tag.
+            match encoder.encode_tagged(dep_node, value) {
+                Ok(()) => {}
+                Err(e) => {
+                    res = Err(e);
+                }
             }
         }
-        Ok(())
-    })
+    });
+
+    res
 }
diff --git a/compiler/rustc_query_impl/src/profiling_support.rs b/compiler/rustc_query_impl/src/profiling_support.rs
index 24485889731..2517793ecea 100644
--- a/compiler/rustc_query_impl/src/profiling_support.rs
+++ b/compiler/rustc_query_impl/src/profiling_support.rs
@@ -250,8 +250,8 @@ fn alloc_self_profile_query_strings_for_query_cache<'tcx, C>(
             // need to invoke queries itself, we cannot keep the query caches
             // locked while doing so. Instead we copy out the
             // `(query_key, dep_node_index)` pairs and release the lock again.
-            let query_keys_and_indices: Vec<_> = query_cache
-                .iter_results(|results| results.map(|(k, _, i)| (k.clone(), i)).collect());
+            let mut query_keys_and_indices = Vec::new();
+            query_cache.iter_results(&mut |k, _, i| query_keys_and_indices.push((k.clone(), i)));
 
             // Now actually allocate the strings. If allocating the strings
             // generates new entries in the query cache, we'll miss them but
@@ -275,14 +275,15 @@ fn alloc_self_profile_query_strings_for_query_cache<'tcx, C>(
             let query_name = profiler.get_or_alloc_cached_string(query_name);
             let event_id = event_id_builder.from_label(query_name).to_string_id();
 
-            query_cache.iter_results(|results| {
-                let query_invocation_ids: Vec<_> = results.map(|v| v.2.into()).collect();
-
-                profiler.bulk_map_query_invocation_id_to_single_string(
-                    query_invocation_ids.into_iter(),
-                    event_id,
-                );
+            let mut query_invocation_ids = Vec::new();
+            query_cache.iter_results(&mut |_, _, i| {
+                query_invocation_ids.push(i.into());
             });
+
+            profiler.bulk_map_query_invocation_id_to_single_string(
+                query_invocation_ids.into_iter(),
+                event_id,
+            );
         }
     });
 }
diff --git a/compiler/rustc_query_impl/src/stats.rs b/compiler/rustc_query_impl/src/stats.rs
index 4d52483c3b8..e877034bd7b 100644
--- a/compiler/rustc_query_impl/src/stats.rs
+++ b/compiler/rustc_query_impl/src/stats.rs
@@ -50,13 +50,12 @@ where
         key_type: type_name::<C::Key>(),
         value_size: mem::size_of::<C::Value>(),
         value_type: type_name::<C::Value>(),
-        entry_count: map.iter_results(|results| results.count()),
+        entry_count: 0,
         local_def_id_keys: None,
     };
-    map.iter_results(|results| {
-        for (key, _, _) in results {
-            key.key_stats(&mut stats)
-        }
+    map.iter_results(&mut |key, _, _| {
+        stats.entry_count += 1;
+        key.key_stats(&mut stats)
     });
     stats
 }
diff --git a/compiler/rustc_query_system/src/query/caches.rs b/compiler/rustc_query_system/src/query/caches.rs
index 001bf3b216b..2e4c8d0565a 100644
--- a/compiler/rustc_query_system/src/query/caches.rs
+++ b/compiler/rustc_query_system/src/query/caches.rs
@@ -49,13 +49,11 @@ pub trait QueryCache: QueryStorage {
         index: DepNodeIndex,
     ) -> Self::Stored;
 
-    fn iter<R>(
+    fn iter(
         &self,
         shards: &Sharded<Self::Sharded>,
-        f: impl for<'a> FnOnce(
-            &'a mut dyn Iterator<Item = (&'a Self::Key, &'a Self::Value, DepNodeIndex)>,
-        ) -> R,
-    ) -> R;
+        f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex),
+    );
 }
 
 pub struct DefaultCacheSelector;
@@ -124,14 +122,17 @@ where
         value
     }
 
-    fn iter<R>(
+    fn iter(
         &self,
         shards: &Sharded<Self::Sharded>,
-        f: impl for<'a> FnOnce(&'a mut dyn Iterator<Item = (&'a K, &'a V, DepNodeIndex)>) -> R,
-    ) -> R {
+        f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex),
+    ) {
         let shards = shards.lock_shards();
-        let mut results = shards.iter().flat_map(|shard| shard.iter()).map(|(k, v)| (k, &v.0, v.1));
-        f(&mut results)
+        for shard in shards.iter() {
+            for (k, v) in shard.iter() {
+                f(k, &v.0, v.1);
+            }
+        }
     }
 }
 
@@ -207,13 +208,16 @@ where
         &value.0
     }
 
-    fn iter<R>(
+    fn iter(
         &self,
         shards: &Sharded<Self::Sharded>,
-        f: impl for<'a> FnOnce(&'a mut dyn Iterator<Item = (&'a K, &'a V, DepNodeIndex)>) -> R,
-    ) -> R {
+        f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex),
+    ) {
         let shards = shards.lock_shards();
-        let mut results = shards.iter().flat_map(|shard| shard.iter()).map(|(k, v)| (k, &v.0, v.1));
-        f(&mut results)
+        for shard in shards.iter() {
+            for (k, v) in shard.iter() {
+                f(k, &v.0, v.1);
+            }
+        }
     }
 }
diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs
index fb8a53048fa..f7b83812e89 100644
--- a/compiler/rustc_query_system/src/query/plumbing.rs
+++ b/compiler/rustc_query_system/src/query/plumbing.rs
@@ -73,12 +73,7 @@ impl<C: QueryCache> QueryCacheStore<C> {
         (QueryLookup { key_hash, shard }, lock)
     }
 
-    pub fn iter_results<R>(
-        &self,
-        f: impl for<'a> FnOnce(
-            &'a mut dyn Iterator<Item = (&'a C::Key, &'a C::Value, DepNodeIndex)>,
-        ) -> R,
-    ) -> R {
+    pub fn iter_results(&self, f: &mut dyn FnMut(&C::Key, &C::Value, DepNodeIndex)) {
         self.cache.iter(&self.shards, f)
     }
 }