Wrap TyCtxt inside a QueryCtxt for queries.

cjgillot · cjgillot · commit 2db2776589ec · 2021-02-19T17:51:50.000+01:00
diff --git a/compiler/rustc_macros/src/query.rs b/compiler/rustc_macros/src/query.rs
@@ -354,9 +354,10 @@ fn add_query_description_impl(
             quote! {
                 #[inline]
                 fn try_load_from_disk(
-                    #tcx: TyCtxt<'tcx>,
-                    #id: SerializedDepNodeIndex
+                    tcx: QueryCtxt<'tcx>,
+                    id: SerializedDepNodeIndex
                 ) -> Option<Self::Value> {
+                    let (#tcx, #id) = (*tcx, id);
                     #block
                 }
             }
@@ -365,10 +366,10 @@ fn add_query_description_impl(
             quote! {
                 #[inline]
                 fn try_load_from_disk(
-                    tcx: TyCtxt<'tcx>,
+                    tcx: QueryCtxt<'tcx>,
                     id: SerializedDepNodeIndex
                 ) -> Option<Self::Value> {
-                    tcx.on_disk_cache.as_ref()?.try_load_query_result(tcx, id)
+                    tcx.on_disk_cache.as_ref()?.try_load_query_result(*tcx, id)
                 }
             }
         };
@@ -393,10 +394,11 @@ fn add_query_description_impl(
             #[inline]
             #[allow(unused_variables, unused_braces)]
             fn cache_on_disk(
-                #tcx: TyCtxt<'tcx>,
-                #key: &Self::Key,
-                #value: Option<&Self::Value>
+                tcx: QueryCtxt<'tcx>,
+                key: &Self::Key,
+                value: Option<&Self::Value>
             ) -> bool {
+                let (#tcx, #key, #value) = (*tcx, key, value);
                 #expr
             }
 
@@ -414,16 +416,14 @@ fn add_query_description_impl(
 
     let desc = quote! {
         #[allow(unused_variables)]
-        fn describe(
-            #tcx: TyCtxt<'tcx>,
-            #key: #arg,
-        ) -> String {
-            ::rustc_middle::ty::print::with_no_trimmed_paths(|| format!(#desc))
+        fn describe(tcx: QueryCtxt<'tcx>, key: #arg) -> String {
+            let (#tcx, #key) = (*tcx, key);
+            ::rustc_middle::ty::print::with_no_trimmed_paths(|| format!(#desc).into())
         }
     };
 
     impls.extend(quote! {
-        impl<'tcx> QueryDescription<TyCtxt<'tcx>> for queries::#name<'tcx> {
+        impl<'tcx> QueryDescription<QueryCtxt<'tcx>> for queries::#name<'tcx> {
             #desc
             #cache
         }
diff --git a/compiler/rustc_middle/src/dep_graph/dep_node.rs b/compiler/rustc_middle/src/dep_graph/dep_node.rs
@@ -55,6 +55,7 @@
 //!
 //! [dependency graph]: https://rustc-dev-guide.rust-lang.org/query.html
 
+use crate::ty::query::QueryCtxt;
 use crate::ty::TyCtxt;
 
 use rustc_data_structures::fingerprint::Fingerprint;
@@ -261,7 +262,7 @@ pub mod dep_kind {
 
                     if let Some(key) = recover(tcx, dep_node) {
                         force_query::<queries::$variant<'_>, _>(
-                            tcx,
+                            QueryCtxt(tcx),
                             key,
                             DUMMY_SP,
                             *dep_node
@@ -287,7 +288,7 @@ pub mod dep_kind {
                                      .unwrap_or(false));
 
                     let key = recover(tcx, dep_node).unwrap_or_else(|| panic!("Failed to recover key for {:?} with hash {}", dep_node, dep_node.hash));
-                    if queries::$variant::cache_on_disk(tcx, &key, None) {
+                    if queries::$variant::cache_on_disk(QueryCtxt(tcx), &key, None) {
                         let _ = tcx.$variant(key);
                     }
                 }
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
@@ -7,6 +7,7 @@ use crate::traits::query::{
     CanonicalTypeOpProvePredicateGoal, CanonicalTypeOpSubtypeGoal,
 };
 use crate::ty::query::queries;
+use crate::ty::query::QueryCtxt;
 use crate::ty::subst::{GenericArg, SubstsRef};
 use crate::ty::{self, ParamEnvAnd, Ty, TyCtxt};
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId};
diff --git a/compiler/rustc_middle/src/ty/query/job.rs b/compiler/rustc_middle/src/ty/query/job.rs
@@ -1,3 +1,4 @@
+use crate::ty::query::QueryCtxt;
 use crate::ty::tls;
 
 use rustc_query_system::query::deadlock;
@@ -20,7 +21,7 @@ pub unsafe fn handle_deadlock() {
     thread::spawn(move || {
         tls::enter_context(icx, |_| {
             rustc_span::SESSION_GLOBALS
-                .set(session_globals, || tls::with(|tcx| deadlock(tcx, &registry)))
-        })
+                .set(session_globals, || tls::with(|tcx| deadlock(QueryCtxt(tcx), &registry)))
+        });
     });
 }
diff --git a/compiler/rustc_middle/src/ty/query/mod.rs b/compiler/rustc_middle/src/ty/query/mod.rs
@@ -60,6 +60,7 @@ use std::sync::Arc;
 
 #[macro_use]
 mod plumbing;
+pub use plumbing::QueryCtxt;
 pub(crate) use rustc_query_system::query::CycleError;
 use rustc_query_system::query::*;
 
diff --git a/compiler/rustc_middle/src/ty/query/on_disk_cache.rs b/compiler/rustc_middle/src/ty/query/on_disk_cache.rs
@@ -3,6 +3,7 @@ use crate::mir::interpret::{AllocDecodingSession, AllocDecodingState};
 use crate::mir::{self, interpret};
 use crate::ty::codec::{RefDecodable, TyDecoder, TyEncoder};
 use crate::ty::context::TyCtxt;
+use crate::ty::query::QueryCtxt;
 use crate::ty::{self, Ty};
 use rustc_data_structures::fingerprint::{Fingerprint, FingerprintDecoder, FingerprintEncoder};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexSet};
@@ -312,7 +313,7 @@ impl<'sess> OnDiskCache<'sess> {
                     ($($query:ident,)*) => {
                         $(
                             encode_query_results::<ty::query::queries::$query<'_>>(
-                                tcx,
+                                QueryCtxt(tcx),
                                 enc,
                                 qri
                             )?;
@@ -1230,12 +1231,12 @@ impl<'a> Decodable<opaque::Decoder<'a>> for IntEncodedWithFixedSize {
 }
 
 fn encode_query_results<'a, 'tcx, Q>(
-    tcx: TyCtxt<'tcx>,
+    tcx: QueryCtxt<'tcx>,
     encoder: &mut CacheEncoder<'a, 'tcx, FileEncoder>,
     query_result_index: &mut EncodedQueryResultIndex,
 ) -> FileEncodeResult
 where
-    Q: super::QueryDescription<TyCtxt<'tcx>> + super::QueryAccessors<TyCtxt<'tcx>>,
+    Q: super::QueryDescription<QueryCtxt<'tcx>> + super::QueryAccessors<QueryCtxt<'tcx>>,
     Q::Value: Encodable<CacheEncoder<'a, 'tcx, FileEncoder>>,
 {
     let _timer = tcx
diff --git a/compiler/rustc_middle/src/ty/query/plumbing.rs b/compiler/rustc_middle/src/ty/query/plumbing.rs
@@ -5,6 +5,7 @@
 use crate::ty::query::Query;
 use crate::ty::tls::{self, ImplicitCtxt};
 use crate::ty::{self, TyCtxt};
+use rustc_query_system::dep_graph::HasDepContext;
 use rustc_query_system::query::QueryContext;
 use rustc_query_system::query::{CycleError, QueryJobId, QueryJobInfo};
 
@@ -15,7 +16,29 @@ use rustc_errors::{struct_span_err, Diagnostic, DiagnosticBuilder, Handler, Leve
 use rustc_span::def_id::DefId;
 use rustc_span::Span;
 
-impl QueryContext for TyCtxt<'tcx> {
+#[derive(Copy, Clone)]
+pub struct QueryCtxt<'tcx>(pub TyCtxt<'tcx>);
+
+impl<'tcx> std::ops::Deref for QueryCtxt<'tcx> {
+    type Target = TyCtxt<'tcx>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl HasDepContext for QueryCtxt<'tcx> {
+    type DepKind = crate::dep_graph::DepKind;
+    type StableHashingContext = crate::ich::StableHashingContext<'tcx>;
+    type DepContext = TyCtxt<'tcx>;
+
+    #[inline]
+    fn dep_context(&self) -> &Self::DepContext {
+        &self.0
+    }
+}
+
+impl QueryContext for QueryCtxt<'tcx> {
     type Query = Query<'tcx>;
 
     fn incremental_verify_ich(&self) -> bool {
@@ -26,11 +49,11 @@ impl QueryContext for TyCtxt<'tcx> {
     }
 
     fn def_path_str(&self, def_id: DefId) -> String {
-        TyCtxt::def_path_str(*self, def_id)
+        self.0.def_path_str(def_id)
     }
 
     fn current_query_job(&self) -> Option<QueryJobId<Self::DepKind>> {
-        tls::with_related_context(*self, |icx| icx.query)
+        tls::with_related_context(**self, |icx| icx.query)
     }
 
     fn try_collect_active_jobs(
@@ -53,10 +76,10 @@ impl QueryContext for TyCtxt<'tcx> {
         // The `TyCtxt` stored in TLS has the same global interner lifetime
         // as `self`, so we use `with_related_context` to relate the 'tcx lifetimes
         // when accessing the `ImplicitCtxt`.
-        tls::with_related_context(*self, move |current_icx| {
+        tls::with_related_context(**self, move |current_icx| {
             // Update the `ImplicitCtxt` to point to our new query job.
             let new_icx = ImplicitCtxt {
-                tcx: *self,
+                tcx: **self,
                 query: Some(token),
                 diagnostics,
                 layout_depth: current_icx.layout_depth,
@@ -71,7 +94,7 @@ impl QueryContext for TyCtxt<'tcx> {
     }
 }
 
-impl<'tcx> TyCtxt<'tcx> {
+impl<'tcx> QueryCtxt<'tcx> {
     #[inline(never)]
     #[cold]
     pub(super) fn report_cycle(
@@ -81,7 +104,7 @@ impl<'tcx> TyCtxt<'tcx> {
         assert!(!stack.is_empty());
 
         let fix_span = |span: Span, query: &Query<'tcx>| {
-            self.sess.source_map().guess_head_span(query.default_span(self, span))
+            self.sess.source_map().guess_head_span(query.default_span(*self, span))
         };
 
         // Disable naming impls with types in this path, since that
@@ -119,7 +142,9 @@ impl<'tcx> TyCtxt<'tcx> {
             err
         })
     }
+}
 
+impl<'tcx> TyCtxt<'tcx> {
     pub fn try_print_query_stack(handler: &Handler, num_frames: Option<usize>) {
         eprintln!("query stack during panic:");
 
@@ -149,7 +174,7 @@ impl<'tcx> TyCtxt<'tcx> {
                             "#{} [{}] {}",
                             i,
                             query_info.info.query.name(),
-                            query_info.info.query.describe(icx.tcx)
+                            query_info.info.query.describe(QueryCtxt(icx.tcx))
                         ),
                     );
                     diag.span =
@@ -272,7 +297,7 @@ macro_rules! define_queries {
                 }
             }
 
-            pub fn describe(&self, tcx: TyCtxt<$tcx>) -> String {
+            pub(crate) fn describe(&self, tcx: QueryCtxt<$tcx>) -> String {
                 let (r, name) = match *self {
                     $(Query::$name(key) => {
                         (queries::$name::describe(tcx, key), stringify!($name))
@@ -362,35 +387,35 @@ macro_rules! define_queries {
             const NAME: &'static str = stringify!($name);
         }
 
-        impl<$tcx> QueryAccessors<TyCtxt<$tcx>> for queries::$name<$tcx> {
+        impl<$tcx> QueryAccessors<QueryCtxt<$tcx>> for queries::$name<$tcx> {
             const ANON: bool = is_anon!([$($modifiers)*]);
             const EVAL_ALWAYS: bool = is_eval_always!([$($modifiers)*]);
             const DEP_KIND: dep_graph::DepKind = dep_graph::DepKind::$name;
 
             type Cache = query_storage::$name<$tcx>;
 
             #[inline(always)]
-            fn query_state<'a>(tcx: TyCtxt<$tcx>) -> &'a QueryState<crate::dep_graph::DepKind, Query<$tcx>, Self::Key> {
+            fn query_state<'a>(tcx: QueryCtxt<$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>
+            fn query_cache<'a>(tcx: QueryCtxt<$tcx>) -> &'a QueryCacheStore<Self::Cache>
                 where 'tcx:'a
             {
                 &tcx.query_caches.$name
             }
 
             #[inline]
-            fn compute(tcx: TyCtxt<'tcx>, key: Self::Key) -> Self::Value {
+            fn compute(tcx: QueryCtxt<'tcx>, key: Self::Key) -> Self::Value {
                 let provider = tcx.queries.providers.get(key.query_crate())
                     // HACK(eddyb) it's possible crates may be loaded after
                     // the query engine is created, and because crate loading
                     // is not yet integrated with the query engine, such crates
                     // would be missing appropriate entries in `providers`.
                     .unwrap_or(&tcx.queries.fallback_extern_providers)
                     .$name;
-                provider(tcx, key)
+                provider(*tcx, key)
             }
 
             fn hash_result(
@@ -401,7 +426,7 @@ macro_rules! define_queries {
             }
 
             fn handle_cycle_error(
-                tcx: TyCtxt<'tcx>,
+                tcx: QueryCtxt<'tcx>,
                 error: CycleError<Query<'tcx>>
             ) -> Self::Value {
                 handle_cycle_error!([$($modifiers)*][tcx, error])
@@ -425,7 +450,8 @@ macro_rules! define_queries {
                     Err(lookup) => lookup,
                 };
 
-                get_query::<queries::$name<'_>, _>(self.tcx, DUMMY_SP, key, lookup, QueryMode::Ensure);
+                let qcx = QueryCtxt(self.tcx);
+                get_query::<queries::$name<'_>, _>(qcx, DUMMY_SP, key, lookup, QueryMode::Ensure);
             })*
         }
 
@@ -516,7 +542,8 @@ macro_rules! define_queries {
                     Err(lookup) => lookup,
                 };
 
-                get_query::<queries::$name<'_>, _>(self.tcx, self.span, key, lookup, QueryMode::Get).unwrap()
+                let qcx = QueryCtxt(self.tcx);
+                get_query::<queries::$name<'_>, _>(qcx, self.span, key, lookup, QueryMode::Get).unwrap()
             })*
         }
 
@@ -558,12 +585,12 @@ macro_rules! define_queries_struct {
 
             pub(crate) fn try_collect_active_jobs(
                 &self
-            ) -> Option<FxHashMap<QueryJobId<crate::dep_graph::DepKind>, QueryJobInfo<crate::dep_graph::DepKind, <TyCtxt<$tcx> as QueryContext>::Query>>> {
+            ) -> Option<FxHashMap<QueryJobId<crate::dep_graph::DepKind>, QueryJobInfo<crate::dep_graph::DepKind, Query<$tcx>>>> {
                 let mut jobs = FxHashMap::default();
 
                 $(
                     self.$name.try_collect_active_jobs(
-                        <queries::$name<'tcx> as QueryAccessors<TyCtxt<'tcx>>>::DEP_KIND,
+                        <queries::$name<'tcx> as QueryAccessors<QueryCtxt<'tcx>>>::DEP_KIND,
                         Query::$name,
                         &mut jobs,
                     )?;
diff --git a/compiler/rustc_middle/src/ty/query/stats.rs b/compiler/rustc_middle/src/ty/query/stats.rs
@@ -1,7 +1,7 @@
-use crate::ty::query::queries;
+use crate::ty::query::query_storage;
 use crate::ty::TyCtxt;
 use rustc_hir::def_id::{DefId, LOCAL_CRATE};
-use rustc_query_system::query::{QueryAccessors, QueryCache, QueryCacheStore};
+use rustc_query_system::query::{QueryCache, QueryCacheStore};
 
 use std::any::type_name;
 use std::mem;
@@ -125,7 +125,7 @@ macro_rules! print_stats {
 
             $(
                 queries.push(stats::<
-                    <queries::$name<'_> as QueryAccessors<TyCtxt<'_>>>::Cache,
+                    query_storage::$name<'_>,
                 >(
                     stringify!($name),
                     &tcx.query_caches.$name,
diff --git a/compiler/rustc_middle/src/ty/query/values.rs b/compiler/rustc_middle/src/ty/query/values.rs
