Move Sharded maps into each QueryCache impl

3 files changed, 44 insertions, 113 deletions

diff --git a/compiler/rustc_query_system/src/query/caches.rs b/compiler/rustc_query_system/src/query/caches.rs
index 011c2ceebb7..05cc61c83aa 100644
--- a/compiler/rustc_query_system/src/query/caches.rs
+++ b/compiler/rustc_query_system/src/query/caches.rs
@@ -1,9 +1,9 @@
 use crate::dep_graph::DepNodeIndex;
-use crate::query::plumbing::{QueryCacheStore, QueryLookup};
+use crate::query::plumbing::QueryLookup;
 
 use rustc_arena::TypedArena;
 use rustc_data_structures::fx::FxHashMap;
-use rustc_data_structures::sharded::Sharded;
+use rustc_data_structures::sharded::{self, Sharded};
 use rustc_data_structures::sync::WorkerLocal;
 use std::default::Default;
 use std::fmt::Debug;
@@ -25,15 +25,13 @@ pub trait QueryStorage {
 
 pub trait QueryCache: QueryStorage + Sized {
     type Key: Hash + Eq + Clone + Debug;
-    type Sharded: Default;
 
     /// Checks if the query is already computed and in the cache.
     /// It returns the shard index and a lock guard to the shard,
     /// which will be used if the query is not in the cache and we need
     /// to compute it.
-    fn lookup<'s, R, OnHit>(
+    fn lookup<R, OnHit>(
         &self,
-        state: &'s QueryCacheStore<Self>,
         key: &Self::Key,
         // `on_hit` can be called while holding a lock to the query state shard.
         on_hit: OnHit,
@@ -41,19 +39,9 @@ pub trait QueryCache: QueryStorage + Sized {
     where
         OnHit: FnOnce(&Self::Stored, DepNodeIndex) -> R;
 
-    fn complete(
-        &self,
-        lock_sharded_storage: &mut Self::Sharded,
-        key: Self::Key,
-        value: Self::Value,
-        index: DepNodeIndex,
-    ) -> Self::Stored;
+    fn complete(&self, key: Self::Key, value: Self::Value, index: DepNodeIndex) -> Self::Stored;
 
-    fn iter(
-        &self,
-        shards: &Sharded<Self::Sharded>,
-        f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex),
-    );
+    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex));
 }
 
 pub struct DefaultCacheSelector;
@@ -62,11 +50,13 @@ impl<K: Eq + Hash, V: Clone> CacheSelector<K, V> for DefaultCacheSelector {
     type Cache = DefaultCache<K, V>;
 }
 
-pub struct DefaultCache<K, V>(PhantomData<(K, V)>);
+pub struct DefaultCache<K, V> {
+    shards: Sharded<FxHashMap<K, (V, DepNodeIndex)>>,
+}
 
 impl<K, V> Default for DefaultCache<K, V> {
     fn default() -> Self {
-        DefaultCache(PhantomData)
+        DefaultCache { shards: Default::default() }
     }
 }
 
@@ -87,19 +77,16 @@ where
     V: Clone + Debug,
 {
     type Key = K;
-    type Sharded = FxHashMap<K, (V, DepNodeIndex)>;
 
     #[inline(always)]
-    fn lookup<'s, R, OnHit>(
-        &self,
-        state: &'s QueryCacheStore<Self>,
-        key: &K,
-        on_hit: OnHit,
-    ) -> Result<R, QueryLookup>
+    fn lookup<R, OnHit>(&self, key: &K, on_hit: OnHit) -> Result<R, QueryLookup>
     where
         OnHit: FnOnce(&V, DepNodeIndex) -> R,
     {
-        let (lookup, lock) = state.get_lookup(key);
+        let key_hash = sharded::make_hash(key);
+        let shard = sharded::get_shard_index_by_hash(key_hash);
+        let lock = self.shards.get_shard_by_index(shard).lock();
+        let lookup = QueryLookup { key_hash, shard };
         let result = lock.raw_entry().from_key_hashed_nocheck(lookup.key_hash, key);
 
         if let Some((_, value)) = result {
@@ -111,23 +98,13 @@ where
     }
 
     #[inline]
-    fn complete(
-        &self,
-        lock_sharded_storage: &mut Self::Sharded,
-        key: K,
-        value: V,
-        index: DepNodeIndex,
-    ) -> Self::Stored {
-        lock_sharded_storage.insert(key, (value.clone(), index));
+    fn complete(&self, key: K, value: V, index: DepNodeIndex) -> Self::Stored {
+        self.shards.get_shard_by_value(&key).lock().insert(key, (value.clone(), index));
         value
     }
 
-    fn iter(
-        &self,
-        shards: &Sharded<Self::Sharded>,
-        f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex),
-    ) {
-        let shards = shards.lock_shards();
+    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
+        let shards = self.shards.lock_shards();
         for shard in shards.iter() {
             for (k, v) in shard.iter() {
                 f(k, &v.0, v.1);
@@ -144,12 +121,15 @@ impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector<K, V> for ArenaCacheSelector<'tc
 
 pub struct ArenaCache<'tcx, K, V> {
     arena: WorkerLocal<TypedArena<(V, DepNodeIndex)>>,
-    phantom: PhantomData<(K, &'tcx V)>,
+    shards: Sharded<FxHashMap<K, &'tcx (V, DepNodeIndex)>>,
 }
 
 impl<'tcx, K, V> Default for ArenaCache<'tcx, K, V> {
     fn default() -> Self {
-        ArenaCache { arena: WorkerLocal::new(|_| TypedArena::default()), phantom: PhantomData }
+        ArenaCache {
+            arena: WorkerLocal::new(|_| TypedArena::default()),
+            shards: Default::default(),
+        }
     }
 }
 
@@ -171,19 +151,16 @@ where
     V: Debug,
 {
     type Key = K;
-    type Sharded = FxHashMap<K, &'tcx (V, DepNodeIndex)>;
 
     #[inline(always)]
-    fn lookup<'s, R, OnHit>(
-        &self,
-        state: &'s QueryCacheStore<Self>,
-        key: &K,
-        on_hit: OnHit,
-    ) -> Result<R, QueryLookup>
+    fn lookup<R, OnHit>(&self, key: &K, on_hit: OnHit) -> Result<R, QueryLookup>
     where
         OnHit: FnOnce(&&'tcx V, DepNodeIndex) -> R,
     {
-        let (lookup, lock) = state.get_lookup(key);
+        let key_hash = sharded::make_hash(key);
+        let shard = sharded::get_shard_index_by_hash(key_hash);
+        let lock = self.shards.get_shard_by_index(shard).lock();
+        let lookup = QueryLookup { key_hash, shard };
         let result = lock.raw_entry().from_key_hashed_nocheck(lookup.key_hash, key);
 
         if let Some((_, value)) = result {
@@ -195,25 +172,15 @@ where
     }
 
     #[inline]
-    fn complete(
-        &self,
-        lock_sharded_storage: &mut Self::Sharded,
-        key: K,
-        value: V,
-        index: DepNodeIndex,
-    ) -> Self::Stored {
+    fn complete(&self, key: K, value: V, index: DepNodeIndex) -> Self::Stored {
         let value = self.arena.alloc((value, index));
         let value = unsafe { &*(value as *const _) };
-        lock_sharded_storage.insert(key, value);
+        self.shards.get_shard_by_value(&key).lock().insert(key, value);
         &value.0
     }
 
-    fn iter(
-        &self,
-        shards: &Sharded<Self::Sharded>,
-        f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex),
-    ) {
-        let shards = shards.lock_shards();
+    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
+        let shards = self.shards.lock_shards();
         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/config.rs b/compiler/rustc_query_system/src/query/config.rs
index b1ff1e15a9d..1ffcdbce6fc 100644
--- a/compiler/rustc_query_system/src/query/config.rs
+++ b/compiler/rustc_query_system/src/query/config.rs
@@ -4,7 +4,7 @@ use crate::dep_graph::DepNode;
 use crate::dep_graph::SerializedDepNodeIndex;
 use crate::ich::StableHashingContext;
 use crate::query::caches::QueryCache;
-use crate::query::{QueryCacheStore, QueryContext, QueryState};
+use crate::query::{QueryContext, QueryState};
 
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_errors::DiagnosticBuilder;
@@ -64,7 +64,7 @@ pub trait QueryDescription<CTX: QueryContext>: QueryConfig {
         CTX: 'a;
 
     // Don't use this method to access query results, instead use the methods on TyCtxt
-    fn query_cache<'a>(tcx: CTX) -> &'a QueryCacheStore<Self::Cache>
+    fn query_cache<'a>(tcx: CTX) -> &'a Self::Cache
     where
         CTX: 'a;
 
diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs
index 77e1fd3f2cc..9278bb602e1 100644
--- a/compiler/rustc_query_system/src/query/plumbing.rs
+++ b/compiler/rustc_query_system/src/query/plumbing.rs
@@ -12,7 +12,7 @@ use rustc_data_structures::fx::{FxHashMap, FxHasher};
 #[cfg(parallel_compiler)]
 use rustc_data_structures::profiling::TimingGuard;
 use rustc_data_structures::sharded::{get_shard_index_by_hash, Sharded};
-use rustc_data_structures::sync::{Lock, LockGuard};
+use rustc_data_structures::sync::Lock;
 use rustc_data_structures::thin_vec::ThinVec;
 use rustc_errors::{DiagnosticBuilder, FatalError};
 use rustc_session::Session;
@@ -24,21 +24,10 @@ use std::hash::{Hash, Hasher};
 use std::mem;
 use std::ptr;
 
-pub struct QueryCacheStore<C: QueryCache> {
-    cache: C,
-    shards: Sharded<C::Sharded>,
-}
-
-impl<C: QueryCache + Default> Default for QueryCacheStore<C> {
-    fn default() -> Self {
-        Self { cache: C::default(), shards: Default::default() }
-    }
-}
-
 /// Values used when checking a query cache which can be reused on a cache-miss to execute the query.
 pub struct QueryLookup {
     pub(super) key_hash: u64,
-    shard: usize,
+    pub(super) shard: usize,
 }
 
 // We compute the key's hash once and then use it for both the
@@ -50,22 +39,6 @@ fn hash_for_shard<K: Hash>(key: &K) -> u64 {
     hasher.finish()
 }
 
-impl<C: QueryCache> QueryCacheStore<C> {
-    pub(super) fn get_lookup<'tcx>(
-        &'tcx self,
-        key: &C::Key,
-    ) -> (QueryLookup, LockGuard<'tcx, C::Sharded>) {
-        let key_hash = hash_for_shard(key);
-        let shard = get_shard_index_by_hash(key_hash);
-        let lock = self.shards.get_shard_by_index(shard).lock();
-        (QueryLookup { key_hash, shard }, lock)
-    }
-
-    pub fn iter_results(&self, f: &mut dyn FnMut(&C::Key, &C::Value, DepNodeIndex)) {
-        self.cache.iter(&self.shards, f)
-    }
-}
-
 struct QueryStateShard<K> {
     active: FxHashMap<K, QueryResult>,
 }
@@ -239,12 +212,7 @@ where
 
     /// Completes the query by updating the query cache with the `result`,
     /// signals the waiter and forgets the JobOwner, so it won't poison the query
-    fn complete<C>(
-        self,
-        cache: &QueryCacheStore<C>,
-        result: C::Value,
-        dep_node_index: DepNodeIndex,
-    ) -> C::Stored
+    fn complete<C>(self, cache: &C, result: C::Value, dep_node_index: DepNodeIndex) -> C::Stored
     where
         C: QueryCache<Key = K>,
     {
@@ -265,10 +233,7 @@ where
                     QueryResult::Poisoned => panic!(),
                 }
             };
-            let result = {
-                let mut lock = cache.shards.get_shard_by_index(shard).lock();
-                cache.cache.complete(&mut lock, key, result, dep_node_index)
-            };
+            let result = cache.complete(key, result, dep_node_index);
             (job, result)
         };
 
@@ -334,7 +299,7 @@ where
 #[inline]
 pub fn try_get_cached<'a, CTX, C, R, OnHit>(
     tcx: CTX,
-    cache: &'a QueryCacheStore<C>,
+    cache: &'a C,
     key: &C::Key,
     // `on_hit` can be called while holding a lock to the query cache
     on_hit: OnHit,
@@ -344,7 +309,7 @@ where
     CTX: DepContext,
     OnHit: FnOnce(&C::Stored) -> R,
 {
-    cache.cache.lookup(cache, &key, |value, index| {
+    cache.lookup(&key, |value, index| {
         if unlikely!(tcx.profiler().enabled()) {
             tcx.profiler().query_cache_hit(index.into());
         }
@@ -356,7 +321,7 @@ where
 fn try_execute_query<CTX, C>(
     tcx: CTX,
     state: &QueryState<C::Key>,
-    cache: &QueryCacheStore<C>,
+    cache: &C,
     span: Span,
     key: C::Key,
     lookup: QueryLookup,
@@ -375,14 +340,13 @@ where
             (result, Some(dep_node_index))
         }
         TryGetJob::Cycle(error) => {
-            let result = mk_cycle(tcx, error, query.handle_cycle_error, &cache.cache);
+            let result = mk_cycle(tcx, error, query.handle_cycle_error, cache);
             (result, None)
         }
         #[cfg(parallel_compiler)]
         TryGetJob::JobCompleted(query_blocked_prof_timer) => {
             let (v, index) = cache
-                .cache
-                .lookup(cache, &key, |value, index| (value.clone(), index))
+                .lookup(&key, |value, index| (value.clone(), index))
                 .unwrap_or_else(|_| panic!("value must be in cache after waiting"));
 
             if unlikely!(tcx.dep_context().profiler().enabled()) {
@@ -760,7 +724,7 @@ where
     // We may be concurrently trying both execute and force a query.
     // Ensure that only one of them runs the query.
     let cache = Q::query_cache(tcx);
-    let cached = cache.cache.lookup(cache, &key, |_, index| {
+    let cached = cache.lookup(&key, |_, index| {
         if unlikely!(tcx.dep_context().profiler().enabled()) {
             tcx.dep_context().profiler().query_cache_hit(index.into());
         }