Add an abstraction for custom query caches

Author: Zoxc

src/librustc/ty/query/caches.rs

@@ -0,0 +1,112 @@
+use crate::dep_graph::DepNodeIndex;
+use crate::ty::query::config::QueryAccessors;
+use crate::ty::query::plumbing::{QueryLookup, QueryState, QueryStateShard};
+use crate::ty::TyCtxt;
+
+use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::sharded::Sharded;
+use std::default::Default;
+use std::hash::Hash;
+
+pub(crate) trait CacheSelector<K, V> {
+    type Cache: QueryCache<K, V>;
+}
+
+pub(crate) trait QueryCache<K, V>: Default {
+    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<'tcx, R, GetCache, OnHit, OnMiss, Q>(
+        &self,
+        state: &'tcx QueryState<'tcx, Q>,
+        get_cache: GetCache,
+        key: K,
+        // `on_hit` can be called while holding a lock to the query state shard.
+        on_hit: OnHit,
+        on_miss: OnMiss,
+    ) -> R
+    where
+        Q: QueryAccessors<'tcx>,
+        GetCache: for<'a> Fn(&'a mut QueryStateShard<'tcx, Q>) -> &'a mut Self::Sharded,
+        OnHit: FnOnce(&V, DepNodeIndex) -> R,
+        OnMiss: FnOnce(K, QueryLookup<'tcx, Q>) -> R;
+
+    fn complete(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        lock_sharded_storage: &mut Self::Sharded,
+        key: K,
+        value: V,
+        index: DepNodeIndex,
+    );
+
+    fn iter<R, L>(
+        &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,
+    ) -> R;
+}
+
+pub struct DefaultCacheSelector;
+
+impl<K: Eq + Hash, V: Clone> CacheSelector<K, V> for DefaultCacheSelector {
+    type Cache = DefaultCache;
+}
+
+#[derive(Default)]
+pub struct DefaultCache;
+
+impl<K: Eq + Hash, V: Clone> QueryCache<K, V> for DefaultCache {
+    type Sharded = FxHashMap<K, (V, DepNodeIndex)>;
+
+    #[inline(always)]
+    fn lookup<'tcx, R, GetCache, OnHit, OnMiss, Q>(
+        &self,
+        state: &'tcx QueryState<'tcx, Q>,
+        get_cache: GetCache,
+        key: K,
+        on_hit: OnHit,
+        on_miss: OnMiss,
+    ) -> R
+    where
+        Q: QueryAccessors<'tcx>,
+        GetCache: for<'a> Fn(&'a mut QueryStateShard<'tcx, Q>) -> &'a mut Self::Sharded,
+        OnHit: FnOnce(&V, DepNodeIndex) -> R,
+        OnMiss: FnOnce(K, QueryLookup<'tcx, Q>) -> R,
+    {
+        let mut lookup = state.get_lookup(&key);
+        let lock = &mut *lookup.lock;
+
+        let result = get_cache(lock).raw_entry().from_key_hashed_nocheck(lookup.key_hash, &key);
+
+        if let Some((_, value)) = result { on_hit(&value.0, value.1) } else { on_miss(key, lookup) }
+    }
+
+    #[inline]
+    fn complete(
+        &self,
+        _: TyCtxt<'tcx>,
+        lock_sharded_storage: &mut Self::Sharded,
+        key: K,
+        value: V,
+        index: DepNodeIndex,
+    ) {
+        lock_sharded_storage.insert(key, (value, index));
+    }
+
+    fn iter<R, L>(
+        &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,
+    ) -> 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))
+    }
+}

src/librustc/ty/query/config.rs

@@ -1,15 +1,15 @@
 use crate::dep_graph::SerializedDepNodeIndex;
 use crate::dep_graph::{DepKind, DepNode};
+use crate::ty::query::caches::QueryCache;
 use crate::ty::query::plumbing::CycleError;
 use crate::ty::query::queries;
-use crate::ty::query::{Query, QueryCache};
+use crate::ty::query::{Query, QueryState};
 use crate::ty::TyCtxt;
 use rustc_data_structures::profiling::ProfileCategory;
 use rustc_hir::def_id::{CrateNum, DefId};
 
 use crate::ich::StableHashingContext;
 use rustc_data_structures::fingerprint::Fingerprint;
-use rustc_data_structures::sharded::Sharded;
 use std::borrow::Cow;
 use std::fmt::Debug;
 use std::hash::Hash;
@@ -30,10 +30,12 @@ pub(crate) trait QueryAccessors<'tcx>: QueryConfig<'tcx> {
     const ANON: bool;
     const EVAL_ALWAYS: bool;
 
+    type Cache: QueryCache<Self::Key, Self::Value>;
+
     fn query(key: Self::Key) -> Query<'tcx>;
 
     // Don't use this method to access query results, instead use the methods on TyCtxt
-    fn query_cache<'a>(tcx: TyCtxt<'tcx>) -> &'a Sharded<QueryCache<'tcx, Self>>;
+    fn query_state<'a>(tcx: TyCtxt<'tcx>) -> &'a QueryState<'tcx, Self>;
 
     fn to_dep_node(tcx: TyCtxt<'tcx>, key: &Self::Key) -> DepNode;
 
@@ -61,7 +63,10 @@ pub(crate) trait QueryDescription<'tcx>: QueryAccessors<'tcx> {
     }
 }
 
-impl<'tcx, M: QueryAccessors<'tcx, Key = DefId>> QueryDescription<'tcx> for M {
+impl<'tcx, M: QueryAccessors<'tcx, Key = DefId>> QueryDescription<'tcx> for M
+where
+    <M as QueryAccessors<'tcx>>::Cache: QueryCache<DefId, <M as QueryConfig<'tcx>>::Value>,
+{
     default fn describe(tcx: TyCtxt<'_>, def_id: DefId) -> Cow<'static, str> {
         if !tcx.sess.verbose() {
             format!("processing `{}`", tcx.def_path_str(def_id)).into()

src/librustc/ty/query/keys.rs

@@ -4,6 +4,7 @@ use crate::infer::canonical::Canonical;
 use crate::mir;
 use crate::traits;
 use crate::ty::fast_reject::SimplifiedType;
+use crate::ty::query::caches::DefaultCacheSelector;
 use crate::ty::subst::SubstsRef;
 use crate::ty::{self, Ty, TyCtxt};
 use rustc_hir::def_id::{CrateNum, DefId, DefIndex, LOCAL_CRATE};
@@ -12,7 +13,9 @@ use rustc_span::{Span, DUMMY_SP};
 
 /// The `Key` trait controls what types can legally be used as the key
 /// for a query.
-pub(super) trait Key {
+pub trait Key {
+    type CacheSelector;
+
     /// Given an instance of this key, what crate is it referring to?
     /// This is used to find the provider.
     fn query_crate(&self) -> CrateNum;
@@ -23,6 +26,8 @@ pub(super) trait Key {
 }
 
 impl<'tcx> Key for ty::InstanceDef<'tcx> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -33,6 +38,8 @@ impl<'tcx> Key for ty::InstanceDef<'tcx> {
 }
 
 impl<'tcx> Key for ty::Instance<'tcx> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -43,6 +50,8 @@ impl<'tcx> Key for ty::Instance<'tcx> {
 }
 
 impl<'tcx> Key for mir::interpret::GlobalId<'tcx> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         self.instance.query_crate()
     }
@@ -53,6 +62,8 @@ impl<'tcx> Key for mir::interpret::GlobalId<'tcx> {
 }
 
 impl<'tcx> Key for mir::interpret::LitToConstInput<'tcx> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -63,6 +74,8 @@ impl<'tcx> Key for mir::interpret::LitToConstInput<'tcx> {
 }
 
 impl Key for CrateNum {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         *self
     }
@@ -72,6 +85,8 @@ impl Key for CrateNum {
 }
 
 impl Key for DefIndex {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -81,6 +96,8 @@ impl Key for DefIndex {
 }
 
 impl Key for DefId {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         self.krate
     }
@@ -90,6 +107,8 @@ impl Key for DefId {
 }
 
 impl Key for (DefId, DefId) {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         self.0.krate
     }
@@ -99,6 +118,8 @@ impl Key for (DefId, DefId) {
 }
 
 impl Key for (CrateNum, DefId) {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         self.0
     }
@@ -108,6 +129,8 @@ impl Key for (CrateNum, DefId) {
 }
 
 impl Key for (DefId, SimplifiedType) {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         self.0.krate
     }
@@ -117,6 +140,8 @@ impl Key for (DefId, SimplifiedType) {
 }
 
 impl<'tcx> Key for SubstsRef<'tcx> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -126,6 +151,8 @@ impl<'tcx> Key for SubstsRef<'tcx> {
 }
 
 impl<'tcx> Key for (DefId, SubstsRef<'tcx>) {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         self.0.krate
     }
@@ -135,6 +162,8 @@ impl<'tcx> Key for (DefId, SubstsRef<'tcx>) {
 }
 
 impl<'tcx> Key for (ty::ParamEnv<'tcx>, ty::PolyTraitRef<'tcx>) {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         self.1.def_id().krate
     }
@@ -144,6 +173,8 @@ impl<'tcx> Key for (ty::ParamEnv<'tcx>, ty::PolyTraitRef<'tcx>) {
 }
 
 impl<'tcx> Key for (&'tcx ty::Const<'tcx>, mir::Field) {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -153,6 +184,8 @@ impl<'tcx> Key for (&'tcx ty::Const<'tcx>, mir::Field) {
 }
 
 impl<'tcx> Key for ty::PolyTraitRef<'tcx> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         self.def_id().krate
     }
@@ -162,6 +195,8 @@ impl<'tcx> Key for ty::PolyTraitRef<'tcx> {
 }
 
 impl<'tcx> Key for &'tcx ty::Const<'tcx> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -171,6 +206,8 @@ impl<'tcx> Key for &'tcx ty::Const<'tcx> {
 }
 
 impl<'tcx> Key for Ty<'tcx> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -180,6 +217,8 @@ impl<'tcx> Key for Ty<'tcx> {
 }
 
 impl<'tcx> Key for ty::ParamEnv<'tcx> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -189,6 +228,8 @@ impl<'tcx> Key for ty::ParamEnv<'tcx> {
 }
 
 impl<'tcx, T: Key> Key for ty::ParamEnvAnd<'tcx, T> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         self.value.query_crate()
     }
@@ -198,6 +239,8 @@ impl<'tcx, T: Key> Key for ty::ParamEnvAnd<'tcx, T> {
 }
 
 impl<'tcx> Key for traits::Environment<'tcx> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -207,6 +250,8 @@ impl<'tcx> Key for traits::Environment<'tcx> {
 }
 
 impl Key for Symbol {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -218,6 +263,8 @@ impl Key for Symbol {
 /// Canonical query goals correspond to abstract trait operations that
 /// are not tied to any crate in particular.
 impl<'tcx, T> Key for Canonical<'tcx, T> {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }
@@ -228,6 +275,8 @@ impl<'tcx, T> Key for Canonical<'tcx, T> {
 }
 
 impl Key for (Symbol, u32, u32) {
+    type CacheSelector = DefaultCacheSelector;
+
     fn query_crate(&self) -> CrateNum {
         LOCAL_CRATE
     }

src/librustc/ty/query/mod.rs

@@ -76,6 +76,9 @@ use self::keys::Key;
 mod values;
 use self::values::Value;
 
+mod caches;
+use self::caches::CacheSelector;
+
 mod config;
 use self::config::QueryAccessors;
 pub use self::config::QueryConfig;