Move query TLS to rustc_query_system.

Author: cjgillot

compiler/rustc_interface/src/callbacks.rs

@@ -9,8 +9,7 @@
 //! The functions in this file should fall back to the default set in their
 //! origin crate when the `TyCtxt` is not present in TLS.
 
-use rustc_errors::{Diagnostic, TRACK_DIAGNOSTIC};
-use rustc_middle::dep_graph::{DepNodeExt, TaskDepsRef};
+use rustc_middle::dep_graph::DepNodeExt;
 use rustc_middle::ty::tls;
 use rustc_query_system::dep_graph::dep_node::default_dep_kind_debug;
 use rustc_query_system::dep_graph::{DepContext, DepKind, DepNode};
@@ -26,26 +25,6 @@ fn track_span_parent(def_id: rustc_span::def_id::LocalDefId) {
     })
 }
 
-/// This is a callback from `rustc_errors` as it cannot access the implicit state
-/// in `rustc_middle` otherwise. It is used when diagnostic messages are
-/// emitted and stores them in the current query, if there is one.
-fn track_diagnostic(diagnostic: Diagnostic, f: &mut dyn FnMut(Diagnostic)) {
-    tls::with_context_opt(|icx| {
-        if let Some(icx) = icx {
-            if let Some(diagnostics) = icx.diagnostics {
-                diagnostics.lock().extend(Some(diagnostic.clone()));
-            }
-
-            // Diagnostics are tracked, we can ignore the dependency.
-            let icx = tls::ImplicitCtxt { task_deps: TaskDepsRef::Ignore, ..icx.clone() };
-            return tls::enter_context(&icx, move || (*f)(diagnostic));
-        }
-
-        // In any other case, invoke diagnostics anyway.
-        (*f)(diagnostic);
-    })
-}
-
 /// This is a callback from `rustc_hir` as it cannot access the implicit state
 /// in `rustc_middle` otherwise.
 fn def_id_debug(def_id: rustc_hir::def_id::DefId, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -103,5 +82,5 @@ pub fn setup_callbacks() {
         .swap(&(dep_kind_debug as fn(_, &mut fmt::Formatter<'_>) -> _));
     rustc_query_system::dep_graph::dep_node::DEP_NODE_DEBUG
         .swap(&(dep_node_debug as fn(_, &mut fmt::Formatter<'_>) -> _));
-    TRACK_DIAGNOSTIC.swap(&(track_diagnostic as _));
+    rustc_errors::TRACK_DIAGNOSTIC.swap(&(rustc_query_system::tls::track_diagnostic as _));
 }

compiler/rustc_interface/src/interface.rs

@@ -450,11 +450,10 @@ pub fn try_print_query_stack(
     // Be careful relying on global state here: this code is called from
     // a panic hook, which means that the global `DiagCtxt` may be in a weird
     // state if it was responsible for triggering the panic.
-    let i = ty::tls::with_context_opt(|icx| {
-        if let Some(icx) = icx {
+    let i = ty::tls::with_opt(|tcx| {
+        if let Some(tcx) = tcx {
             ty::print::with_no_queries!(print_query_stack(
-                QueryCtxt::new(icx.tcx),
-                icx.query,
+                QueryCtxt::new(tcx),
                 dcx,
                 num_frames,
                 file,

compiler/rustc_middle/src/dep_graph/mod.rs

@@ -1,4 +1,4 @@
-use crate::ty::{self, TyCtxt};
+use crate::ty::TyCtxt;
 use rustc_data_structures::profiling::SelfProfilerRef;
 use rustc_query_system::ich::StableHashingContext;
 use rustc_session::Session;
@@ -24,27 +24,6 @@ pub type DepKindStruct<'tcx> = rustc_query_system::dep_graph::DepKindStruct<TyCt
 pub struct DepsType;
 
 impl Deps for DepsType {
-    fn with_deps<OP, R>(task_deps: TaskDepsRef<'_>, op: OP) -> R
-    where
-        OP: FnOnce() -> R,
-    {
-        ty::tls::with_context(|icx| {
-            let icx = ty::tls::ImplicitCtxt { task_deps, ..icx.clone() };
-
-            ty::tls::enter_context(&icx, op)
-        })
-    }
-
-    fn read_deps<OP>(op: OP)
-    where
-        OP: for<'a> FnOnce(TaskDepsRef<'a>),
-    {
-        ty::tls::with_context_opt(|icx| {
-            let Some(icx) = icx else { return };
-            op(icx.task_deps)
-        })
-    }
-
     const DEP_KIND_NULL: DepKind = dep_kinds::Null;
     const DEP_KIND_RED: DepKind = dep_kinds::Red;
     const DEP_KIND_MAX: u16 = dep_node::DEP_KIND_VARIANTS - 1;

compiler/rustc_middle/src/query/plumbing.rs

@@ -16,7 +16,6 @@ use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::hir_id::OwnerId;
 use rustc_query_system::dep_graph::DepNodeIndex;
 use rustc_query_system::dep_graph::SerializedDepNodeIndex;
-pub(crate) use rustc_query_system::query::QueryJobId;
 use rustc_query_system::query::*;
 use rustc_query_system::HandleCycleError;
 use rustc_span::{ErrorGuaranteed, Span, DUMMY_SP};

compiler/rustc_middle/src/ty/context.rs

@@ -2,8 +2,6 @@
 
 #![allow(rustc::usage_of_ty_tykind)]
 
-pub mod tls;
-
 use crate::arena::Arena;
 use crate::dep_graph::{DepGraph, DepKindStruct};
 use crate::infer::canonical::{CanonicalParamEnvCache, CanonicalVarInfo, CanonicalVarInfos};
@@ -662,15 +660,40 @@ impl<'tcx> GlobalCtxt<'tcx> {
     where
         F: FnOnce(TyCtxt<'tcx>) -> R,
     {
-        let icx = tls::ImplicitCtxt::new(self);
-        tls::enter_context(&icx, || f(icx.tcx))
+        rustc_query_system::tls::create_and_enter_context(self, || f(TyCtxt { gcx: self }))
     }
 
     pub fn finish(&self) -> FileEncodeResult {
         self.dep_graph.finish_encoding(&self.sess.prof)
     }
 }
 
+pub mod tls {
+    use super::TyCtxt;
+
+    /// Allows access to the `TyCtxt` in the current `ImplicitCtxt`.
+    /// Panics if there is no `ImplicitCtxt` available.
+    #[inline]
+    pub fn with<F, R>(f: F) -> R
+    where
+        F: for<'tcx> FnOnce(TyCtxt<'tcx>) -> R,
+    {
+        rustc_query_system::tls::with(|gcx| f(TyCtxt { gcx: unsafe { &*gcx.cast() } }))
+    }
+
+    /// Allows access to the `TyCtxt` in the current `ImplicitCtxt`.
+    /// The closure is passed None if there is no `ImplicitCtxt` available.
+    #[inline]
+    pub fn with_opt<F, R>(f: F) -> R
+    where
+        F: for<'tcx> FnOnce(Option<TyCtxt<'tcx>>) -> R,
+    {
+        rustc_query_system::tls::with_opt(|opt_context| {
+            f(opt_context.map(|gcx| TyCtxt { gcx: unsafe { &*gcx.cast() } }))
+        })
+    }
+}
+
 impl<'tcx> TyCtxt<'tcx> {
     /// Expects a body and returns its codegen attributes.
     ///

compiler/rustc_middle/src/ty/context/tls.rs

@@ -6,8 +6,6 @@ use crate::rustc_middle::dep_graph::DepContext;
 use crate::rustc_middle::ty::TyEncoder;
 use crate::QueryConfigRestored;
 use rustc_data_structures::stable_hasher::{Hash64, HashStable, StableHasher};
-use rustc_data_structures::sync::Lock;
-use rustc_errors::Diagnostic;
 
 use rustc_index::Idx;
 use rustc_middle::dep_graph::dep_kinds;
@@ -17,7 +15,6 @@ use rustc_middle::dep_graph::{
 use rustc_middle::query::on_disk_cache::AbsoluteBytePos;
 use rustc_middle::query::on_disk_cache::{CacheDecoder, CacheEncoder, EncodedDepNodeIndex};
 use rustc_middle::query::Key;
-use rustc_middle::ty::tls::{self, ImplicitCtxt};
 use rustc_middle::ty::{self, print::with_no_queries, TyCtxt};
 use rustc_query_system::dep_graph::{DepNodeParams, HasDepContext};
 use rustc_query_system::ich::StableHashingContext;
@@ -31,7 +28,6 @@ use rustc_serialize::Encodable;
 use rustc_session::Limit;
 use rustc_span::def_id::LOCAL_CRATE;
 use std::num::NonZeroU64;
-use thin_vec::ThinVec;
 
 #[derive(Copy, Clone)]
 pub struct QueryCtxt<'tcx> {
@@ -75,11 +71,6 @@ impl QueryContext for QueryCtxt<'_> {
         )
     }
 
-    #[inline]
-    fn current_query_job(self) -> Option<QueryJobId> {
-        tls::with_context(|icx| icx.query)
-    }
-
     fn collect_active_jobs(self) -> QueryMap {
         let mut jobs = QueryMap::default();
 
@@ -119,37 +110,8 @@ impl QueryContext for QueryCtxt<'_> {
         }
     }
 
-    /// Executes a job by changing the `ImplicitCtxt` to point to the
-    /// new query job while it executes. It returns the diagnostics
-    /// captured during execution and the actual result.
-    #[inline(always)]
-    fn start_query<R>(
-        self,
-        token: QueryJobId,
-        depth_limit: bool,
-        diagnostics: Option<&Lock<ThinVec<Diagnostic>>>,
-        compute: impl FnOnce() -> R,
-    ) -> R {
-        // The `TyCtxt` stored in TLS has the same global interner lifetime
-        // as `self`, so we use `with_context` to relate the 'tcx lifetimes
-        // when accessing the `ImplicitCtxt`.
-        tls::with_context(move |current_icx| {
-            if depth_limit && !self.recursion_limit().value_within_limit(current_icx.query_depth) {
-                self.depth_limit_error(token);
-            }
-
-            // Update the `ImplicitCtxt` to point to our new query job.
-            let new_icx = ImplicitCtxt {
-                tcx: self.tcx,
-                query: Some(token),
-                diagnostics,
-                query_depth: current_icx.query_depth + depth_limit as usize,
-                task_deps: current_icx.task_deps,
-            };
-
-            // Use the `ImplicitCtxt` while we execute the query.
-            tls::enter_context(&new_icx, compute)
-        })
+    fn recursion_limit(self) -> Limit {
+        self.tcx.recursion_limit()
     }
 
     fn depth_limit_error(self, job: QueryJobId) {