Rollup merge of #107831 - nnethercote:query-refactoring, r=oli-obk

Query refactoring

Just some cleanups I found when learning about the query system.

Best reviewed one commit at a time.

r? `@oli-obk`

compiler/rustc_const_eval/src/const_eval/eval_queries.rs

@@ -54,7 +54,7 @@ fn eval_body_using_ecx<'mir, 'tcx>(
 
     trace!(
         "eval_body_using_ecx: pushing stack frame for global: {}{}",
-        with_no_trimmed_paths!(ty::tls::with(|tcx| tcx.def_path_str(cid.instance.def_id()))),
+        with_no_trimmed_paths!(ecx.tcx.def_path_str(cid.instance.def_id())),
         cid.promoted.map_or_else(String::new, |p| format!("::promoted[{:?}]", p))
     );
 

compiler/rustc_infer/src/infer/canonical/canonicalizer.rs

@@ -203,12 +203,10 @@ impl CanonicalizeMode for CanonicalizeQueryResponse {
                 // rust-lang/rust#57464: `impl Trait` can leak local
                 // scopes (in manner violating typeck). Therefore, use
                 // `delay_span_bug` to allow type error over an ICE.
-                ty::tls::with(|tcx| {
-                    tcx.sess.delay_span_bug(
-                        rustc_span::DUMMY_SP,
-                        &format!("unexpected region in query response: `{:?}`", r),
-                    );
-                });
+                canonicalizer.tcx.sess.delay_span_bug(
+                    rustc_span::DUMMY_SP,
+                    &format!("unexpected region in query response: `{:?}`", r),
+                );
                 r
             }
         }

compiler/rustc_interface/src/callbacks.rs

@@ -38,7 +38,7 @@ fn track_diagnostic(diagnostic: &mut Diagnostic, f: &mut dyn FnMut(&mut Diagnost
 
             // 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));
+            return tls::enter_context(&icx, move || (*f)(diagnostic));
         }
 
         // In any other case, invoke diagnostics anyway.

compiler/rustc_interface/src/passes.rs

@@ -738,30 +738,16 @@ pub static DEFAULT_EXTERN_QUERY_PROVIDERS: LazyLock<ExternProviders> = LazyLock:
     extern_providers
 });
 
-pub struct QueryContext<'tcx> {
-    gcx: &'tcx GlobalCtxt<'tcx>,
-}
-
-impl<'tcx> QueryContext<'tcx> {
-    pub fn enter<F, R>(&mut self, f: F) -> R
-    where
-        F: FnOnce(TyCtxt<'tcx>) -> R,
-    {
-        let icx = ty::tls::ImplicitCtxt::new(self.gcx);
-        ty::tls::enter_context(&icx, |_| f(icx.tcx))
-    }
-}
-
 pub fn create_global_ctxt<'tcx>(
     compiler: &'tcx Compiler,
     lint_store: Lrc<LintStore>,
     dep_graph: DepGraph,
     untracked: Untracked,
     queries: &'tcx OnceCell<TcxQueries<'tcx>>,
-    global_ctxt: &'tcx OnceCell<GlobalCtxt<'tcx>>,
+    gcx_cell: &'tcx OnceCell<GlobalCtxt<'tcx>>,
     arena: &'tcx WorkerLocal<Arena<'tcx>>,
     hir_arena: &'tcx WorkerLocal<rustc_hir::Arena<'tcx>>,
-) -> QueryContext<'tcx> {
+) -> &'tcx GlobalCtxt<'tcx> {
     // We're constructing the HIR here; we don't care what we will
     // read, since we haven't even constructed the *input* to
     // incr. comp. yet.
@@ -785,8 +771,8 @@ pub fn create_global_ctxt<'tcx>(
         TcxQueries::new(local_providers, extern_providers, query_result_on_disk_cache)
     });
 
-    let gcx = sess.time("setup_global_ctxt", || {
-        global_ctxt.get_or_init(move || {
+    sess.time("setup_global_ctxt", || {
+        gcx_cell.get_or_init(move || {
             TyCtxt::create_global_ctxt(
                 sess,
                 lint_store,
@@ -799,9 +785,7 @@ pub fn create_global_ctxt<'tcx>(
                 rustc_query_impl::query_callbacks(arena),
             )
         })
-    });
-
-    QueryContext { gcx }
+    })
 }
 
 /// Runs the resolution, type-checking, region checking and other

compiler/rustc_interface/src/queries.rs

@@ -1,6 +1,6 @@
 use crate::errors::{FailedWritingFile, RustcErrorFatal, RustcErrorUnexpectedAnnotation};
 use crate::interface::{Compiler, Result};
-use crate::passes::{self, BoxedResolver, QueryContext};
+use crate::passes::{self, BoxedResolver};
 
 use rustc_ast as ast;
 use rustc_codegen_ssa::traits::CodegenBackend;
@@ -64,7 +64,7 @@ impl<'a, T> std::ops::DerefMut for QueryResult<'a, T> {
     }
 }
 
-impl<'a, 'tcx> QueryResult<'a, QueryContext<'tcx>> {
+impl<'a, 'tcx> QueryResult<'a, &'tcx GlobalCtxt<'tcx>> {
     pub fn enter<T>(&mut self, f: impl FnOnce(TyCtxt<'tcx>) -> T) -> T {
         (*self.0).get_mut().enter(f)
     }
@@ -78,7 +78,7 @@ impl<T> Default for Query<T> {
 
 pub struct Queries<'tcx> {
     compiler: &'tcx Compiler,
-    gcx: OnceCell<GlobalCtxt<'tcx>>,
+    gcx_cell: OnceCell<GlobalCtxt<'tcx>>,
     queries: OnceCell<TcxQueries<'tcx>>,
 
     arena: WorkerLocal<Arena<'tcx>>,
@@ -90,15 +90,16 @@ pub struct Queries<'tcx> {
     register_plugins: Query<(ast::Crate, Lrc<LintStore>)>,
     expansion: Query<(Lrc<ast::Crate>, Rc<RefCell<BoxedResolver>>, Lrc<LintStore>)>,
     dep_graph: Query<DepGraph>,
-    global_ctxt: Query<QueryContext<'tcx>>,
+    // This just points to what's in `gcx_cell`.
+    gcx: Query<&'tcx GlobalCtxt<'tcx>>,
     ongoing_codegen: Query<Box<dyn Any>>,
 }
 
 impl<'tcx> Queries<'tcx> {
     pub fn new(compiler: &'tcx Compiler) -> Queries<'tcx> {
         Queries {
             compiler,
-            gcx: OnceCell::new(),
+            gcx_cell: OnceCell::new(),
             queries: OnceCell::new(),
             arena: WorkerLocal::new(|_| Arena::default()),
             hir_arena: WorkerLocal::new(|_| rustc_hir::Arena::default()),
@@ -108,7 +109,7 @@ impl<'tcx> Queries<'tcx> {
             register_plugins: Default::default(),
             expansion: Default::default(),
             dep_graph: Default::default(),
-            global_ctxt: Default::default(),
+            gcx: Default::default(),
             ongoing_codegen: Default::default(),
         }
     }
@@ -207,8 +208,8 @@ impl<'tcx> Queries<'tcx> {
         })
     }
 
-    pub fn global_ctxt(&'tcx self) -> Result<QueryResult<'_, QueryContext<'tcx>>> {
-        self.global_ctxt.compute(|| {
+    pub fn global_ctxt(&'tcx self) -> Result<QueryResult<'_, &'tcx GlobalCtxt<'tcx>>> {
+        self.gcx.compute(|| {
             let crate_name = *self.crate_name()?.borrow();
             let (krate, resolver, lint_store) = self.expansion()?.steal();
 
@@ -218,18 +219,18 @@ impl<'tcx> Queries<'tcx> {
                 ast_lowering: untracked_resolver_for_lowering,
             } = BoxedResolver::to_resolver_outputs(resolver);
 
-            let mut qcx = passes::create_global_ctxt(
+            let gcx = passes::create_global_ctxt(
                 self.compiler,
                 lint_store,
                 self.dep_graph()?.steal(),
                 untracked,
                 &self.queries,
-                &self.gcx,
+                &self.gcx_cell,
                 &self.arena,
                 &self.hir_arena,
             );
 
-            qcx.enter(|tcx| {
+            gcx.enter(|tcx| {
                 let feed = tcx.feed_unit_query();
                 feed.resolver_for_lowering(
                     tcx.arena.alloc(Steal::new((untracked_resolver_for_lowering, krate))),
@@ -239,7 +240,7 @@ impl<'tcx> Queries<'tcx> {
                 let feed = tcx.feed_local_crate();
                 feed.crate_name(crate_name);
             });
-            Ok(qcx)
+            Ok(gcx)
         })
     }
 
@@ -387,7 +388,7 @@ impl Compiler {
 
         // NOTE: intentionally does not compute the global context if it hasn't been built yet,
         // since that likely means there was a parse error.
-        if let Some(Ok(gcx)) = &mut *queries.global_ctxt.result.borrow_mut() {
+        if let Some(Ok(gcx)) = &mut *queries.gcx.result.borrow_mut() {
             let gcx = gcx.get_mut();
             // We assume that no queries are run past here. If there are new queries
             // after this point, they'll show up as "<unknown>" in self-profiling data.

compiler/rustc_middle/src/dep_graph/mod.rs

@@ -55,7 +55,7 @@ impl rustc_query_system::dep_graph::DepKind for DepKind {
         ty::tls::with_context(|icx| {
             let icx = ty::tls::ImplicitCtxt { task_deps, ..icx.clone() };
 
-            ty::tls::enter_context(&icx, |_| op())
+            ty::tls::enter_context(&icx, op)
         })
     }
 

compiler/rustc_middle/src/lib.rs

@@ -34,6 +34,7 @@
 #![feature(get_mut_unchecked)]
 #![feature(if_let_guard)]
 #![feature(iter_from_generator)]
+#![feature(local_key_cell_methods)]
 #![feature(negative_impls)]
 #![feature(never_type)]
 #![feature(extern_types)]

compiler/rustc_middle/src/ty/context.rs

@@ -468,6 +468,18 @@ pub struct GlobalCtxt<'tcx> {
     pub(crate) alloc_map: Lock<interpret::AllocMap<'tcx>>,
 }
 
+impl<'tcx> GlobalCtxt<'tcx> {
+    /// Installs `self` in a `TyCtxt` and `ImplicitCtxt` for the duration of
+    /// `f`.
+    pub fn enter<'a: 'tcx, F, R>(&'a self, f: F) -> R
+    where
+        F: FnOnce(TyCtxt<'tcx>) -> R,
+    {
+        let icx = tls::ImplicitCtxt::new(self);
+        tls::enter_context(&icx, || f(icx.tcx))
+    }
+}
+
 impl<'tcx> TyCtxt<'tcx> {
     /// Expects a body and returns its codegen attributes.
     ///

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

@@ -89,9 +89,8 @@ mod tlv {
     /// This is used to set the pointer to the new `ImplicitCtxt`.
     #[inline]
     pub(super) fn with_tlv<F: FnOnce() -> R, R>(value: *const (), f: F) -> R {
-        let old = get_tlv();
-        let _reset = rustc_data_structures::OnDrop(move || TLV.with(|tlv| tlv.set(old)));
-        TLV.with(|tlv| tlv.set(value));
+        let old = TLV.replace(value);
+        let _reset = rustc_data_structures::OnDrop(move || TLV.set(old));
         f()
     }
 }
@@ -110,9 +109,9 @@ unsafe fn downcast<'a, 'tcx>(context: *const ()) -> &'a ImplicitCtxt<'a, 'tcx> {
 #[inline]
 pub fn enter_context<'a, 'tcx, F, R>(context: &ImplicitCtxt<'a, 'tcx>, f: F) -> R
 where
-    F: FnOnce(&ImplicitCtxt<'a, 'tcx>) -> R,
+    F: FnOnce() -> R,
 {
-    tlv::with_tlv(erase(context), || f(&context))
+    tlv::with_tlv(erase(context), f)
 }
 
 /// Allows access to the current `ImplicitCtxt` in a closure if one is available.

compiler/rustc_query_impl/src/plumbing.rs

@@ -124,7 +124,7 @@ impl QueryContext for QueryCtxt<'_> {
             };
 
             // Use the `ImplicitCtxt` while we execute the query.
-            tls::enter_context(&new_icx, |_| {
+            tls::enter_context(&new_icx, || {
                 rustc_data_structures::stack::ensure_sufficient_stack(compute)
             })
         })

compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs

@@ -98,6 +98,7 @@ impl<'tcx, 'a> GeneratorData<'tcx, 'a> {
     // obligation
     fn get_from_await_ty<F>(
         &self,
+        tcx: TyCtxt<'tcx>,
         visitor: AwaitsVisitor,
         hir: map::Map<'tcx>,
         ty_matches: F,
@@ -134,9 +135,7 @@ impl<'tcx, 'a> GeneratorData<'tcx, 'a> {
                                         .unwrap_or_else(|| {
                                             bug!(
                                                 "node_type: no type for node {}",
-                                                ty::tls::with(|tcx| tcx
-                                                    .hir()
-                                                    .node_to_string(await_expr.hir_id))
+                                                tcx.hir().node_to_string(await_expr.hir_id)
                                             )
                                         })
                                 },
@@ -2351,7 +2350,7 @@ impl<'tcx> TypeErrCtxtExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
 
         let mut interior_or_upvar_span = None;
 
-        let from_awaited_ty = generator_data.get_from_await_ty(visitor, hir, ty_matches);
+        let from_awaited_ty = generator_data.get_from_await_ty(self.tcx, visitor, hir, ty_matches);
         debug!(?from_awaited_ty);
 
         // The generator interior types share the same binders

src/librustdoc/lib.rs

@@ -814,9 +814,9 @@ fn main_args(at_args: &[String]) -> MainResult {
                 sess.fatal("Compilation failed, aborting rustdoc");
             }
 
-            let mut global_ctxt = abort_on_err(queries.global_ctxt(), sess);
+            let mut gcx = abort_on_err(queries.global_ctxt(), sess);
 
-            global_ctxt.enter(|tcx| {
+            gcx.enter(|tcx| {
                 let (krate, render_opts, mut cache) = sess.time("run_global_ctxt", || {
                     core::run_global_ctxt(
                         tcx,