From adcbe49b16dff960600b5d54171d38d6b556d228 Mon Sep 17 00:00:00 2001
From: Tyson Nottingham <tgnottingham@gmail.com>
Date: Thu, 11 Mar 2021 11:22:47 -0800
Subject: [PATCH] rustc_query_system: simplify QueryCache::iter

Minor cleanup to reduce a small amount of complexity and code bloat.
Reduces the number of mono items in rustc_query_impl by 15%.
---
 .../rustc_query_system/src/query/caches.rs    | 35 ++++++++-----------
 .../rustc_query_system/src/query/plumbing.rs  |  4 +--
 2 files changed, 17 insertions(+), 22 deletions(-)

diff --git a/compiler/rustc_query_system/src/query/caches.rs b/compiler/rustc_query_system/src/query/caches.rs
index ec71c8685804f..001bf3b216b1b 100644
--- a/compiler/rustc_query_system/src/query/caches.rs
+++ b/compiler/rustc_query_system/src/query/caches.rs
@@ -49,12 +49,11 @@ pub trait QueryCache: QueryStorage {
         index: DepNodeIndex,
     ) -> Self::Stored;
 
-    fn iter<R, L>(
+    fn iter<R>(
         &self,
-        shards: &Sharded<L>,
-        get_shard: impl Fn(&mut L) -> &mut Self::Sharded,
+        shards: &Sharded<Self::Sharded>,
         f: impl for<'a> FnOnce(
-            Box<dyn Iterator<Item = (&'a Self::Key, &'a Self::Value, DepNodeIndex)> + 'a>,
+            &'a mut dyn Iterator<Item = (&'a Self::Key, &'a Self::Value, DepNodeIndex)>,
         ) -> R,
     ) -> R;
 }
@@ -125,16 +124,14 @@ where
         value
     }
 
-    fn iter<R, L>(
+    fn iter<R>(
         &self,
-        shards: &Sharded<L>,
-        get_shard: impl Fn(&mut L) -> &mut Self::Sharded,
-        f: impl for<'a> FnOnce(Box<dyn Iterator<Item = (&'a K, &'a V, DepNodeIndex)> + 'a>) -> R,
+        shards: &Sharded<Self::Sharded>,
+        f: impl for<'a> FnOnce(&'a mut dyn Iterator<Item = (&'a K, &'a V, DepNodeIndex)>) -> R,
     ) -> R {
-        let mut shards = shards.lock_shards();
-        let mut shards: Vec<_> = shards.iter_mut().map(|shard| get_shard(shard)).collect();
-        let results = shards.iter_mut().flat_map(|shard| shard.iter()).map(|(k, v)| (k, &v.0, v.1));
-        f(Box::new(results))
+        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)
     }
 }
 
@@ -210,15 +207,13 @@ where
         &value.0
     }
 
-    fn iter<R, L>(
+    fn iter<R>(
         &self,
-        shards: &Sharded<L>,
-        get_shard: impl Fn(&mut L) -> &mut Self::Sharded,
-        f: impl for<'a> FnOnce(Box<dyn Iterator<Item = (&'a K, &'a V, DepNodeIndex)> + 'a>) -> R,
+        shards: &Sharded<Self::Sharded>,
+        f: impl for<'a> FnOnce(&'a mut dyn Iterator<Item = (&'a K, &'a V, DepNodeIndex)>) -> R,
     ) -> R {
-        let mut shards = shards.lock_shards();
-        let mut shards: Vec<_> = shards.iter_mut().map(|shard| get_shard(shard)).collect();
-        let results = shards.iter_mut().flat_map(|shard| shard.iter()).map(|(k, v)| (k, &v.0, v.1));
-        f(Box::new(results))
+        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)
     }
 }
diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs
index 0050670b338ea..c3e92b87c27d2 100644
--- a/compiler/rustc_query_system/src/query/plumbing.rs
+++ b/compiler/rustc_query_system/src/query/plumbing.rs
@@ -76,10 +76,10 @@ impl<C: QueryCache> QueryCacheStore<C> {
     pub fn iter_results<R>(
         &self,
         f: impl for<'a> FnOnce(
-            Box<dyn Iterator<Item = (&'a C::Key, &'a C::Value, DepNodeIndex)> + 'a>,
+            &'a mut dyn Iterator<Item = (&'a C::Key, &'a C::Value, DepNodeIndex)>,
         ) -> R,
     ) -> R {
-        self.cache.iter(&self.shards, |shard| &mut *shard, f)
+        self.cache.iter(&self.shards, f)
     }
 }