Separate the query cache from the query state.

Author: cjgillot

compiler/rustc_data_structures/src/sharded.rs

@@ -63,23 +63,9 @@ impl<T> Sharded<T> {
         if SHARDS == 1 { &self.shards[0].0 } else { self.get_shard_by_hash(make_hash(val)) }
     }
 
-    /// Get a shard with a pre-computed hash value. If `get_shard_by_value` is
-    /// ever used in combination with `get_shard_by_hash` on a single `Sharded`
-    /// instance, then `hash` must be computed with `FxHasher`. Otherwise,
-    /// `hash` can be computed with any hasher, so long as that hasher is used
-    /// consistently for each `Sharded` instance.
-    #[inline]
-    pub fn get_shard_index_by_hash(&self, hash: u64) -> usize {
-        let hash_len = mem::size_of::<usize>();
-        // Ignore the top 7 bits as hashbrown uses these and get the next SHARD_BITS highest bits.
-        // hashbrown also uses the lowest bits, so we can't use those
-        let bits = (hash >> (hash_len * 8 - 7 - SHARD_BITS)) as usize;
-        bits % SHARDS
-    }
-
     #[inline]
     pub fn get_shard_by_hash(&self, hash: u64) -> &Lock<T> {
-        &self.shards[self.get_shard_index_by_hash(hash)].0
+        &self.shards[get_shard_index_by_hash(hash)].0
     }
 
     #[inline]
@@ -166,3 +152,17 @@ fn make_hash<K: Hash + ?Sized>(val: &K) -> u64 {
     val.hash(&mut state);
     state.finish()
 }
+
+/// Get a shard with a pre-computed hash value. If `get_shard_by_value` is
+/// ever used in combination with `get_shard_by_hash` on a single `Sharded`
+/// instance, then `hash` must be computed with `FxHasher`. Otherwise,
+/// `hash` can be computed with any hasher, so long as that hasher is used
+/// consistently for each `Sharded` instance.
+#[inline]
+pub fn get_shard_index_by_hash(hash: u64) -> usize {
+    let hash_len = mem::size_of::<usize>();
+    // Ignore the top 7 bits as hashbrown uses these and get the next SHARD_BITS highest bits.
+    // hashbrown also uses the lowest bits, so we can't use those
+    let bits = (hash >> (hash_len * 8 - 7 - SHARD_BITS)) as usize;
+    bits % SHARDS
+}

compiler/rustc_middle/src/ty/context.rs

@@ -963,6 +963,7 @@ pub struct GlobalCtxt<'tcx> {
     pub(crate) definitions: &'tcx Definitions,
 
     pub queries: query::Queries<'tcx>,
+    pub query_caches: query::QueryCaches<'tcx>,
 
     maybe_unused_trait_imports: FxHashSet<LocalDefId>,
     maybe_unused_extern_crates: Vec<(LocalDefId, Span)>,
@@ -1154,6 +1155,7 @@ impl<'tcx> TyCtxt<'tcx> {
             untracked_crate: krate,
             definitions,
             queries: query::Queries::new(providers, extern_providers, on_disk_query_result_cache),
+            query_caches: query::QueryCaches::default(),
             ty_rcache: Default::default(),
             pred_rcache: Default::default(),
             selection_cache: Default::default(),

compiler/rustc_middle/src/ty/query/on_disk_cache.rs

@@ -1244,10 +1244,9 @@ where
         .prof
         .extra_verbose_generic_activity("encode_query_results_for", std::any::type_name::<Q>());
 
-    let state = Q::query_state(tcx);
-    assert!(state.all_inactive());
-
-    state.iter_results(|results| {
+    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());

compiler/rustc_middle/src/ty/query/plumbing.rs

@@ -355,6 +355,11 @@ macro_rules! define_queries {
             $(pub type $name<$tcx> = <query_storage::$name<$tcx> as QueryStorage>::Stored;)*
         }
 
+        #[derive(Default)]
+        pub struct QueryCaches<$tcx> {
+            $($(#[$attr])* $name: QueryCacheStore<query_storage::$name<$tcx>>,)*
+        }
+
         $(impl<$tcx> QueryConfig for queries::$name<$tcx> {
             type Key = $($K)*;
             type Value = $V;
@@ -370,10 +375,17 @@ macro_rules! define_queries {
             type Cache = query_storage::$name<$tcx>;
 
             #[inline(always)]
-            fn query_state<'a>(tcx: TyCtxt<$tcx>) -> &'a QueryState<crate::dep_graph::DepKind, <TyCtxt<$tcx> as QueryContext>::Query, Self::Cache> {
+            fn query_state<'a>(tcx: TyCtxt<$tcx>) -> &'a QueryState<crate::dep_graph::DepKind, Query<$tcx>, Self::Key> {
                 &tcx.queries.$name
             }
 
+            #[inline(always)]
+            fn query_cache<'a>(tcx: TyCtxt<$tcx>) -> &'a QueryCacheStore<Self::Cache>
+                where 'tcx:'a
+            {
+                &tcx.query_caches.$name
+            }
+
             #[inline]
             fn compute(tcx: TyCtxt<'tcx>, key: Self::Key) -> Self::Value {
                 let provider = tcx.queries.providers.get(key.query_crate())
@@ -479,7 +491,7 @@ macro_rules! define_queries {
                     alloc_self_profile_query_strings_for_query_cache(
                         self,
                         stringify!($name),
-                        &self.queries.$name,
+                        &self.query_caches.$name,
                         &mut string_cache,
                     );
                 })*
@@ -525,8 +537,8 @@ macro_rules! define_queries_struct {
 
             $($(#[$attr])*  $name: QueryState<
                 crate::dep_graph::DepKind,
-                <TyCtxt<$tcx> as QueryContext>::Query,
-                query_storage::$name<$tcx>,
+                Query<$tcx>,
+                query_keys::$name<$tcx>,
             >,)*
         }
 

compiler/rustc_middle/src/ty/query/profiling_support.rs

@@ -5,7 +5,7 @@ use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::profiling::SelfProfiler;
 use rustc_hir::def_id::{CrateNum, DefId, DefIndex, LocalDefId, CRATE_DEF_INDEX, LOCAL_CRATE};
 use rustc_hir::definitions::DefPathData;
-use rustc_query_system::query::{QueryCache, QueryContext, QueryState};
+use rustc_query_system::query::{QueryCache, QueryCacheStore};
 use std::fmt::Debug;
 use std::io::Write;
 
@@ -230,7 +230,7 @@ where
 pub(super) fn alloc_self_profile_query_strings_for_query_cache<'tcx, C>(
     tcx: TyCtxt<'tcx>,
     query_name: &'static str,
-    query_state: &QueryState<crate::dep_graph::DepKind, <TyCtxt<'tcx> as QueryContext>::Query, C>,
+    query_cache: &QueryCacheStore<C>,
     string_cache: &mut QueryKeyStringCache,
 ) where
     C: QueryCache,
@@ -251,7 +251,7 @@ pub(super) 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_state
+            let query_keys_and_indices: Vec<_> = query_cache
                 .iter_results(|results| results.map(|(k, _, i)| (k.clone(), i)).collect());
 
             // Now actually allocate the strings. If allocating the strings
@@ -276,7 +276,7 @@ pub(super) 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_state.iter_results(|results| {
+            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(

compiler/rustc_middle/src/ty/query/stats.rs

@@ -1,10 +1,9 @@
 use crate::ty::query::queries;
 use crate::ty::TyCtxt;
 use rustc_hir::def_id::{DefId, LOCAL_CRATE};
-use rustc_query_system::query::{QueryAccessors, QueryCache, QueryContext, QueryState};
+use rustc_query_system::query::{QueryAccessors, QueryCache, QueryCacheStore};
 
 use std::any::type_name;
-use std::hash::Hash;
 use std::mem;
 #[cfg(debug_assertions)]
 use std::sync::atomic::Ordering;
@@ -37,10 +36,8 @@ struct QueryStats {
     local_def_id_keys: Option<usize>,
 }
 
-fn stats<D, Q, C>(name: &'static str, map: &QueryState<D, Q, C>) -> QueryStats
+fn stats<C>(name: &'static str, map: &QueryCacheStore<C>) -> QueryStats
 where
-    D: Copy + Clone + Eq + Hash,
-    Q: Clone,
     C: QueryCache,
 {
     let mut stats = QueryStats {
@@ -128,12 +125,10 @@ macro_rules! print_stats {
 
             $(
                 queries.push(stats::<
-                    crate::dep_graph::DepKind,
-                    <TyCtxt<'_> as QueryContext>::Query,
                     <queries::$name<'_> as QueryAccessors<TyCtxt<'_>>>::Cache,
                 >(
                     stringify!($name),
-                    &tcx.queries.$name,
+                    &tcx.query_caches.$name,
                 ));
             )*
 

compiler/rustc_query_system/src/query/caches.rs

@@ -1,5 +1,5 @@
 use crate::dep_graph::DepNodeIndex;
-use crate::query::plumbing::{QueryLookup, QueryState};
+use crate::query::plumbing::{QueryCacheStore, QueryLookup};
 
 use rustc_arena::TypedArena;
 use rustc_data_structures::fx::FxHashMap;
@@ -31,13 +31,13 @@ pub trait QueryCache: QueryStorage {
     /// 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, D, Q, R, OnHit>(
+    fn lookup<'s, R, OnHit>(
         &self,
-        state: &'s QueryState<D, Q, 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,
-    ) -> Result<R, QueryLookup<'s, D, Q, Self::Key, Self::Sharded>>
+    ) -> Result<R, QueryLookup<'s, Self::Sharded>>
     where
         OnHit: FnOnce(&Self::Stored, DepNodeIndex) -> R;
 
@@ -93,17 +93,17 @@ where
     type Sharded = FxHashMap<K, (V, DepNodeIndex)>;
 
     #[inline(always)]
-    fn lookup<'s, D, Q, R, OnHit>(
+    fn lookup<'s, R, OnHit>(
         &self,
-        state: &'s QueryState<D, Q, Self>,
+        state: &'s QueryCacheStore<Self>,
         key: &K,
         on_hit: OnHit,
-    ) -> Result<R, QueryLookup<'s, D, Q, K, Self::Sharded>>
+    ) -> Result<R, QueryLookup<'s, Self::Sharded>>
     where
         OnHit: FnOnce(&V, DepNodeIndex) -> R,
     {
         let lookup = state.get_lookup(key);
-        let result = lookup.lock.cache.raw_entry().from_key_hashed_nocheck(lookup.key_hash, key);
+        let result = lookup.lock.raw_entry().from_key_hashed_nocheck(lookup.key_hash, key);
 
         if let Some((_, value)) = result {
             let hit_result = on_hit(&value.0, value.1);
@@ -176,17 +176,17 @@ where
     type Sharded = FxHashMap<K, &'tcx (V, DepNodeIndex)>;
 
     #[inline(always)]
-    fn lookup<'s, D, Q, R, OnHit>(
+    fn lookup<'s, R, OnHit>(
         &self,
-        state: &'s QueryState<D, Q, Self>,
+        state: &'s QueryCacheStore<Self>,
         key: &K,
         on_hit: OnHit,
-    ) -> Result<R, QueryLookup<'s, D, Q, K, Self::Sharded>>
+    ) -> Result<R, QueryLookup<'s, Self::Sharded>>
     where
         OnHit: FnOnce(&&'tcx V, DepNodeIndex) -> R,
     {
         let lookup = state.get_lookup(key);
-        let result = lookup.lock.cache.raw_entry().from_key_hashed_nocheck(lookup.key_hash, key);
+        let result = lookup.lock.raw_entry().from_key_hashed_nocheck(lookup.key_hash, key);
 
         if let Some((_, value)) = result {
             let hit_result = on_hit(&&value.0, value.1);

compiler/rustc_query_system/src/query/config.rs

@@ -4,7 +4,7 @@ use crate::dep_graph::DepNode;
 use crate::dep_graph::SerializedDepNodeIndex;
 use crate::query::caches::QueryCache;
 use crate::query::plumbing::CycleError;
-use crate::query::{QueryContext, QueryState};
+use crate::query::{QueryCacheStore, QueryContext, QueryState};
 
 use rustc_data_structures::fingerprint::Fingerprint;
 use std::fmt::Debug;
@@ -73,7 +73,12 @@ pub trait QueryAccessors<CTX: QueryContext>: QueryConfig {
     type Cache: QueryCache<Key = Self::Key, Stored = Self::Stored, Value = Self::Value>;
 
     // Don't use this method to access query results, instead use the methods on TyCtxt
-    fn query_state<'a>(tcx: CTX) -> &'a QueryState<CTX::DepKind, CTX::Query, Self::Cache>;
+    fn query_state<'a>(tcx: CTX) -> &'a QueryState<CTX::DepKind, CTX::Query, Self::Key>;
+
+    // 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>
+    where
+        CTX: 'a;
 
     fn to_dep_node(tcx: CTX, key: &Self::Key) -> DepNode<CTX::DepKind>
     where