Make QueryCache generic on the context.

5 files changed, 50 insertions, 46 deletions

diff --git a/src/librustc/ty/query/caches.rs b/src/librustc/ty/query/caches.rs
index 7dd858142da..f740fada1e5 100644
--- a/src/librustc/ty/query/caches.rs
+++ b/src/librustc/ty/query/caches.rs
@@ -1,6 +1,6 @@
 use crate::dep_graph::DepNodeIndex;
+use crate::ty::query::config::QueryContext;
 use crate::ty::query::plumbing::{QueryLookup, QueryState, QueryStateShard};
-use crate::ty::TyCtxt;
 
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sharded::Sharded;
@@ -8,11 +8,11 @@ use std::default::Default;
 use std::hash::Hash;
 use std::marker::PhantomData;
 
-pub(crate) trait CacheSelector<K, V> {
-    type Cache: QueryCache<Key = K, Value = V>;
+pub(crate) trait CacheSelector<CTX: QueryContext, K, V> {
+    type Cache: QueryCache<CTX, Key = K, Value = V>;
 }
 
-pub(crate) trait QueryCache: Default {
+pub(crate) trait QueryCache<CTX: QueryContext>: Default {
     type Key;
     type Value;
     type Sharded: Default;
@@ -21,9 +21,9 @@ pub(crate) trait QueryCache: Default {
     /// 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>(
+    fn lookup<R, GetCache, OnHit, OnMiss>(
         &self,
-        state: &'tcx QueryState<TyCtxt<'tcx>, Self>,
+        state: &QueryState<CTX, Self>,
         get_cache: GetCache,
         key: Self::Key,
         // `on_hit` can be called while holding a lock to the query state shard.
@@ -32,14 +32,14 @@ pub(crate) trait QueryCache: Default {
     ) -> R
     where
         GetCache: for<'a> Fn(
-            &'a mut QueryStateShard<TyCtxt<'tcx>, Self::Key, Self::Sharded>,
+            &'a mut QueryStateShard<CTX, Self::Key, Self::Sharded>,
         ) -> &'a mut Self::Sharded,
         OnHit: FnOnce(&Self::Value, DepNodeIndex) -> R,
-        OnMiss: FnOnce(Self::Key, QueryLookup<'tcx, TyCtxt<'tcx>, Self::Key, Self::Sharded>) -> R;
+        OnMiss: FnOnce(Self::Key, QueryLookup<'_, CTX, Self::Key, Self::Sharded>) -> R;
 
     fn complete(
         &self,
-        tcx: TyCtxt<'tcx>,
+        tcx: CTX,
         lock_sharded_storage: &mut Self::Sharded,
         key: Self::Key,
         value: Self::Value,
@@ -58,7 +58,7 @@ pub(crate) trait QueryCache: Default {
 
 pub struct DefaultCacheSelector;
 
-impl<K: Eq + Hash, V: Clone> CacheSelector<K, V> for DefaultCacheSelector {
+impl<CTX: QueryContext, K: Eq + Hash, V: Clone> CacheSelector<CTX, K, V> for DefaultCacheSelector {
     type Cache = DefaultCache<K, V>;
 }
 
@@ -70,26 +70,25 @@ impl<K, V> Default for DefaultCache<K, V> {
     }
 }
 
-impl<K: Eq + Hash, V: Clone> QueryCache for DefaultCache<K, V> {
+impl<CTX: QueryContext, K: Eq + Hash, V: Clone> QueryCache<CTX> for DefaultCache<K, V> {
     type Key = K;
     type Value = V;
     type Sharded = FxHashMap<K, (V, DepNodeIndex)>;
 
     #[inline(always)]
-    fn lookup<'tcx, R, GetCache, OnHit, OnMiss>(
+    fn lookup<R, GetCache, OnHit, OnMiss>(
         &self,
-        state: &'tcx QueryState<TyCtxt<'tcx>, Self>,
+        state: &QueryState<CTX, Self>,
         get_cache: GetCache,
         key: K,
         on_hit: OnHit,
         on_miss: OnMiss,
     ) -> R
     where
-        GetCache: for<'a> Fn(
-            &'a mut QueryStateShard<TyCtxt<'tcx>, K, Self::Sharded>,
-        ) -> &'a mut Self::Sharded,
+        GetCache:
+            for<'a> Fn(&'a mut QueryStateShard<CTX, K, Self::Sharded>) -> &'a mut Self::Sharded,
         OnHit: FnOnce(&V, DepNodeIndex) -> R,
-        OnMiss: FnOnce(K, QueryLookup<'tcx, TyCtxt<'tcx>, K, Self::Sharded>) -> R,
+        OnMiss: FnOnce(K, QueryLookup<'_, CTX, K, Self::Sharded>) -> R,
     {
         let mut lookup = state.get_lookup(&key);
         let lock = &mut *lookup.lock;
@@ -102,7 +101,7 @@ impl<K: Eq + Hash, V: Clone> QueryCache for DefaultCache<K, V> {
     #[inline]
     fn complete(
         &self,
-        _: TyCtxt<'tcx>,
+        _: CTX,
         lock_sharded_storage: &mut Self::Sharded,
         key: K,
         value: V,
diff --git a/src/librustc/ty/query/config.rs b/src/librustc/ty/query/config.rs
index e65f3094820..fd2f855d5b1 100644
--- a/src/librustc/ty/query/config.rs
+++ b/src/librustc/ty/query/config.rs
@@ -38,7 +38,7 @@ pub(crate) trait QueryAccessors<CTX: QueryContext>: QueryConfig<CTX> {
     const EVAL_ALWAYS: bool;
     const DEP_KIND: DepKind;
 
-    type Cache: QueryCache<Key = Self::Key, Value = Self::Value>;
+    type Cache: QueryCache<CTX, Key = Self::Key, 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, Self::Cache>;
diff --git a/src/librustc/ty/query/plumbing.rs b/src/librustc/ty/query/plumbing.rs
index c0f2c4a11bc..8462610a6b6 100644
--- a/src/librustc/ty/query/plumbing.rs
+++ b/src/librustc/ty/query/plumbing.rs
@@ -51,14 +51,14 @@ impl<CTX: QueryContext, K, C: Default> Default for QueryStateShard<CTX, K, C> {
     }
 }
 
-pub(crate) struct QueryState<CTX: QueryContext, C: QueryCache> {
+pub(crate) struct QueryState<CTX: QueryContext, C: QueryCache<CTX>> {
     cache: C,
     shards: Sharded<QueryStateShard<CTX, C::Key, C::Sharded>>,
     #[cfg(debug_assertions)]
     pub(super) cache_hits: AtomicUsize,
 }
 
-impl<CTX: QueryContext, C: QueryCache> QueryState<CTX, C> {
+impl<CTX: QueryContext, C: QueryCache<CTX>> QueryState<CTX, C> {
     pub(super) fn get_lookup<K2: Hash>(
         &'tcx self,
         key: &K2,
@@ -86,7 +86,7 @@ enum QueryResult<CTX: QueryContext> {
     Poisoned,
 }
 
-impl<CTX: QueryContext, C: QueryCache> QueryState<CTX, C> {
+impl<CTX: QueryContext, C: QueryCache<CTX>> QueryState<CTX, C> {
     pub(super) fn iter_results<R>(
         &self,
         f: impl for<'a> FnOnce(
@@ -130,7 +130,7 @@ impl<CTX: QueryContext, C: QueryCache> QueryState<CTX, C> {
     }
 }
 
-impl<CTX: QueryContext, C: QueryCache> Default for QueryState<CTX, C> {
+impl<CTX: QueryContext, C: QueryCache<CTX>> Default for QueryState<CTX, C> {
     fn default() -> QueryState<CTX, C> {
         QueryState {
             cache: C::default(),
@@ -152,7 +152,7 @@ pub(crate) struct QueryLookup<'tcx, CTX: QueryContext, K, C> {
 /// This will poison the relevant query if dropped.
 struct JobOwner<'tcx, CTX: QueryContext, C>
 where
-    C: QueryCache,
+    C: QueryCache<CTX>,
     C::Key: Eq + Hash + Clone + Debug,
     C::Value: Clone,
 {
@@ -161,9 +161,9 @@ where
     id: QueryJobId,
 }
 
-impl<'tcx, C: QueryCache> JobOwner<'tcx, TyCtxt<'tcx>, C>
+impl<'tcx, C> JobOwner<'tcx, TyCtxt<'tcx>, C>
 where
-    C: QueryCache,
+    C: QueryCache<TyCtxt<'tcx>> + 'tcx,
     C::Key: Eq + Hash + Clone + Debug,
     C::Value: Clone,
 {
@@ -176,12 +176,12 @@ where
     /// This function is inlined because that results in a noticeable speed-up
     /// for some compile-time benchmarks.
     #[inline(always)]
-    fn try_start<Q>(
+    fn try_start<'a, 'b, Q>(
         tcx: TyCtxt<'tcx>,
         span: Span,
         key: &C::Key,
-        mut lookup: QueryLookup<'tcx, TyCtxt<'tcx>, C::Key, C::Sharded>,
-    ) -> TryGetJob<'tcx, C>
+        mut lookup: QueryLookup<'a, TyCtxt<'tcx>, C::Key, C::Sharded>,
+    ) -> TryGetJob<'b, TyCtxt<'tcx>, C>
     where
         Q: QueryDescription<TyCtxt<'tcx>, Key = C::Key, Value = C::Value, Cache = C>,
     {
@@ -262,16 +262,16 @@ where
     }
 }
 
-impl<'tcx, CTX: QueryContext, C: QueryCache> JobOwner<'tcx, CTX, C>
+impl<'tcx, CTX: QueryContext, C> JobOwner<'tcx, CTX, C>
 where
-    C: QueryCache,
+    C: QueryCache<CTX>,
     C::Key: Eq + Hash + Clone + Debug,
     C::Value: Clone,
 {
     /// 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
     #[inline(always)]
-    fn complete(self, tcx: TyCtxt<'tcx>, result: &C::Value, dep_node_index: DepNodeIndex) {
+    fn complete(self, tcx: CTX, result: &C::Value, dep_node_index: DepNodeIndex) {
         // We can move out of `self` here because we `mem::forget` it below
         let key = unsafe { ptr::read(&self.key) };
         let state = self.state;
@@ -304,7 +304,7 @@ where
     (result, diagnostics.into_inner())
 }
 
-impl<'tcx, CTX: QueryContext, C: QueryCache> Drop for JobOwner<'tcx, CTX, C>
+impl<'tcx, CTX: QueryContext, C: QueryCache<CTX>> Drop for JobOwner<'tcx, CTX, C>
 where
     C::Key: Eq + Hash + Clone + Debug,
     C::Value: Clone,
@@ -338,13 +338,13 @@ pub(crate) struct CycleError<CTX: QueryContext> {
 }
 
 /// The result of `try_start`.
-enum TryGetJob<'tcx, C: QueryCache>
+enum TryGetJob<'tcx, CTX: QueryContext, C: QueryCache<CTX>>
 where
     C::Key: Eq + Hash + Clone + Debug,
     C::Value: Clone,
 {
     /// The query is not yet started. Contains a guard to the cache eventually used to start it.
-    NotYetStarted(JobOwner<'tcx, TyCtxt<'tcx>, C>),
+    NotYetStarted(JobOwner<'tcx, CTX, C>),
 
     /// The query was already completed.
     /// Returns the result of the query and its dep-node index
@@ -504,9 +504,9 @@ impl<'tcx> TyCtxt<'tcx> {
         on_miss: OnMiss,
     ) -> R
     where
-        C: QueryCache,
+        C: QueryCache<TyCtxt<'tcx>>,
         OnHit: FnOnce(&C::Value, DepNodeIndex) -> R,
-        OnMiss: FnOnce(C::Key, QueryLookup<'tcx, TyCtxt<'tcx>, C::Key, C::Sharded>) -> R,
+        OnMiss: FnOnce(C::Key, QueryLookup<'_, TyCtxt<'tcx>, C::Key, C::Sharded>) -> R,
     {
         state.cache.lookup(
             state,
@@ -550,7 +550,12 @@ impl<'tcx> TyCtxt<'tcx> {
         self,
         span: Span,
         key: Q::Key,
-        lookup: QueryLookup<'tcx, TyCtxt<'tcx>, Q::Key, <Q::Cache as QueryCache>::Sharded>,
+        lookup: QueryLookup<
+            '_,
+            TyCtxt<'tcx>,
+            Q::Key,
+            <Q::Cache as QueryCache<TyCtxt<'tcx>>>::Sharded,
+        >,
     ) -> Q::Value {
         let job = match JobOwner::try_start::<Q>(self, span, &key, lookup) {
             TryGetJob::NotYetStarted(job) => job,
@@ -866,14 +871,14 @@ macro_rules! is_eval_always {
 }
 
 macro_rules! query_storage {
-    ([][$K:ty, $V:ty]) => {
-        <<$K as Key>::CacheSelector as CacheSelector<$K, $V>>::Cache
+    (<$tcx:tt>[][$K:ty, $V:ty]) => {
+        <<$K as Key>::CacheSelector as CacheSelector<TyCtxt<$tcx>, $K, $V>>::Cache
     };
-    ([storage($ty:ty) $($rest:tt)*][$K:ty, $V:ty]) => {
+    (<$tcx:tt>[storage($ty:ty) $($rest:tt)*][$K:ty, $V:ty]) => {
         $ty
     };
-    ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)*][$($args:tt)*]) => {
-        query_storage!([$($($modifiers)*)*][$($args)*])
+    (<$tcx:tt>[$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)*][$($args:tt)*]) => {
+        query_storage!(<$tcx>[$($($modifiers)*)*][$($args)*])
     };
 }
 
@@ -989,7 +994,7 @@ macro_rules! define_queries_inner {
             const EVAL_ALWAYS: bool = is_eval_always!([$($modifiers)*]);
             const DEP_KIND: dep_graph::DepKind = dep_graph::DepKind::$node;
 
-            type Cache = query_storage!([$($modifiers)*][$K, $V]);
+            type Cache = query_storage!(<$tcx>[$($modifiers)*][$K, $V]);
 
             #[inline(always)]
             fn query_state<'a>(tcx: TyCtxt<$tcx>) -> &'a QueryState<TyCtxt<$tcx>, Self::Cache> {
diff --git a/src/librustc/ty/query/profiling_support.rs b/src/librustc/ty/query/profiling_support.rs
index a540a18a19f..616fbaafab9 100644
--- a/src/librustc/ty/query/profiling_support.rs
+++ b/src/librustc/ty/query/profiling_support.rs
@@ -163,7 +163,7 @@ pub(super) fn alloc_self_profile_query_strings_for_query_cache<'tcx, C>(
     query_state: &QueryState<TyCtxt<'tcx>, C>,
     string_cache: &mut QueryKeyStringCache,
 ) where
-    C: QueryCache,
+    C: QueryCache<TyCtxt<'tcx>>,
     C::Key: Debug + Clone,
 {
     tcx.prof.with_profiler(|profiler| {
diff --git a/src/librustc/ty/query/stats.rs b/src/librustc/ty/query/stats.rs
index e6578d1eb5f..a13f00dc6d4 100644
--- a/src/librustc/ty/query/stats.rs
+++ b/src/librustc/ty/query/stats.rs
@@ -38,7 +38,7 @@ struct QueryStats {
     local_def_id_keys: Option<usize>,
 }
 
-fn stats<CTX: QueryContext, C: QueryCache>(
+fn stats<CTX: QueryContext, C: QueryCache<CTX>>(
     name: &'static str,
     map: &QueryState<CTX, C>,
 ) -> QueryStats {