Rollup merge of #132410 - bjorn3:yet_another_driver_refactor_round, r=cjgillot

Some more refactorings towards removing driver queries

Follow up to https://github.com/rust-lang/rust/pull/127184

## Custom driver breaking change

The `after_analysis` callback is changed to accept `TyCtxt` instead of `Queries`. The only safe query in `Queries` to call at this point is `global_ctxt()` which allows you to enter the `TyCtxt` either way. To fix your custom driver, replace the `queries: &'tcx Queries<'tcx>` argument with `tcx: TyCtxt<'tcx>` and remove your `queries.global_ctxt().unwrap().enter(|tcx| { ... })` call and only keep the contents of the closure.

## Custom driver deprecation

The `after_crate_root_parsing` callback is now deprecated. Several custom drivers are incorrectly calling `queries.global_ctxt()` from inside of it, which causes some driver code to be skipped. As such I would like to either remove it in the future or if custom drivers still need it, change it to accept an `&rustc_ast::Crate` instead.

33 files changed, 227 insertions, 195 deletions

diff --git a/compiler/rustc_driver_impl/src/lib.rs b/compiler/rustc_driver_impl/src/lib.rs
index 256266d2965..005053e4508 100644
--- a/compiler/rustc_driver_impl/src/lib.rs
+++ b/compiler/rustc_driver_impl/src/lib.rs
@@ -160,6 +160,9 @@ pub trait Callbacks {
     /// Called after parsing the crate root. Submodules are not yet parsed when
     /// this callback is called. Return value instructs the compiler whether to
     /// continue the compilation afterwards (defaults to `Compilation::Continue`)
+    #[deprecated = "This callback will likely be removed or stop giving access \
+                    to the TyCtxt in the future. Use either the after_expansion \
+                    or the after_analysis callback instead."]
     fn after_crate_root_parsing<'tcx>(
         &mut self,
         _compiler: &interface::Compiler,
@@ -181,7 +184,7 @@ pub trait Callbacks {
     fn after_analysis<'tcx>(
         &mut self,
         _compiler: &interface::Compiler,
-        _queries: &'tcx Queries<'tcx>,
+        _tcx: TyCtxt<'tcx>,
     ) -> Compilation {
         Compilation::Continue
     }
@@ -335,19 +338,12 @@ fn run_compiler(
         expanded_args: args,
     };
 
-    let has_input = match make_input(&default_early_dcx, &matches.free) {
-        Err(reported) => return Err(reported),
-        Ok(Some(input)) => {
+    let has_input = match make_input(&default_early_dcx, &matches.free)? {
+        Some(input) => {
             config.input = input;
             true // has input: normal compilation
         }
-        Ok(None) => match matches.free.as_slice() {
-            [] => false, // no input: we will exit early
-            [_] => panic!("make_input should have provided valid inputs"),
-            [fst, snd, ..] => default_early_dcx.early_fatal(format!(
-                "multiple input filenames provided (first two filenames are `{fst}` and `{snd}`)"
-            )),
-        },
+        None => false, // no input: we will exit early
     };
 
     drop(default_early_dcx);
@@ -405,10 +401,6 @@ fn run_compiler(
                     queries.global_ctxt()?.enter(|tcx| {
                         tcx.ensure().early_lint_checks(());
                         pretty::print(sess, pp_mode, pretty::PrintExtra::NeedsAstMap { tcx });
-                        Ok(())
-                    })?;
-
-                    queries.global_ctxt()?.enter(|tcx| {
                         passes::write_dep_info(tcx);
                     });
                 } else {
@@ -421,6 +413,7 @@ fn run_compiler(
                 return early_exit();
             }
 
+            #[allow(deprecated)]
             if callbacks.after_crate_root_parsing(compiler, queries) == Compilation::Stop {
                 return early_exit();
             }
@@ -442,25 +435,23 @@ fn run_compiler(
 
             queries.global_ctxt()?.enter(|tcx| {
                 passes::write_dep_info(tcx);
-            });
 
-            if sess.opts.output_types.contains_key(&OutputType::DepInfo)
-                && sess.opts.output_types.len() == 1
-            {
-                return early_exit();
-            }
+                if sess.opts.output_types.contains_key(&OutputType::DepInfo)
+                    && sess.opts.output_types.len() == 1
+                {
+                    return early_exit();
+                }
 
-            if sess.opts.unstable_opts.no_analysis {
-                return early_exit();
-            }
+                if sess.opts.unstable_opts.no_analysis {
+                    return early_exit();
+                }
 
-            queries.global_ctxt()?.enter(|tcx| tcx.analysis(()))?;
+                tcx.analysis(())?;
 
-            if callbacks.after_analysis(compiler, queries) == Compilation::Stop {
-                return early_exit();
-            }
+                if callbacks.after_analysis(compiler, tcx) == Compilation::Stop {
+                    return early_exit();
+                }
 
-            queries.global_ctxt()?.enter(|tcx| {
                 Ok(Some(Linker::codegen_and_build_linker(tcx, &*compiler.codegen_backend)?))
             })
         })?;
@@ -509,37 +500,40 @@ fn make_input(
     early_dcx: &EarlyDiagCtxt,
     free_matches: &[String],
 ) -> Result<Option<Input>, ErrorGuaranteed> {
-    let [input_file] = free_matches else { return Ok(None) };
-
-    if input_file != "-" {
-        // Normal `Input::File`
-        return Ok(Some(Input::File(PathBuf::from(input_file))));
-    }
-
-    // read from stdin as `Input::Str`
-    let mut input = String::new();
-    if io::stdin().read_to_string(&mut input).is_err() {
-        // Immediately stop compilation if there was an issue reading
-        // the input (for example if the input stream is not UTF-8).
-        let reported =
-            early_dcx.early_err("couldn't read from stdin, as it did not contain valid UTF-8");
-        return Err(reported);
-    }
+    match free_matches {
+        [] => Ok(None), // no input: we will exit early,
+        [ifile] if ifile == "-" => {
+            // read from stdin as `Input::Str`
+            let mut input = String::new();
+            if io::stdin().read_to_string(&mut input).is_err() {
+                // Immediately stop compilation if there was an issue reading
+                // the input (for example if the input stream is not UTF-8).
+                let reported = early_dcx
+                    .early_err("couldn't read from stdin, as it did not contain valid UTF-8");
+                return Err(reported);
+            }
 
-    let name = match env::var("UNSTABLE_RUSTDOC_TEST_PATH") {
-        Ok(path) => {
-            let line = env::var("UNSTABLE_RUSTDOC_TEST_LINE").expect(
-                "when UNSTABLE_RUSTDOC_TEST_PATH is set \
+            let name = match env::var("UNSTABLE_RUSTDOC_TEST_PATH") {
+                Ok(path) => {
+                    let line = env::var("UNSTABLE_RUSTDOC_TEST_LINE").expect(
+                        "when UNSTABLE_RUSTDOC_TEST_PATH is set \
                                     UNSTABLE_RUSTDOC_TEST_LINE also needs to be set",
-            );
-            let line = isize::from_str_radix(&line, 10)
-                .expect("UNSTABLE_RUSTDOC_TEST_LINE needs to be an number");
-            FileName::doc_test_source_code(PathBuf::from(path), line)
-        }
-        Err(_) => FileName::anon_source_code(&input),
-    };
+                    );
+                    let line = isize::from_str_radix(&line, 10)
+                        .expect("UNSTABLE_RUSTDOC_TEST_LINE needs to be an number");
+                    FileName::doc_test_source_code(PathBuf::from(path), line)
+                }
+                Err(_) => FileName::anon_source_code(&input),
+            };
 
-    Ok(Some(Input::Str { name, input }))
+            Ok(Some(Input::Str { name, input }))
+        }
+        [ifile] => Ok(Some(Input::File(PathBuf::from(ifile)))),
+        [ifile1, ifile2, ..] => early_dcx.early_fatal(format!(
+            "multiple input filenames provided (first two filenames are `{}` and `{}`)",
+            ifile1, ifile2
+        )),
+    }
 }
 
 /// Whether to stop or continue compilation.
diff --git a/compiler/rustc_incremental/src/lib.rs b/compiler/rustc_incremental/src/lib.rs
index e8735cba9bd..6dab6468870 100644
--- a/compiler/rustc_incremental/src/lib.rs
+++ b/compiler/rustc_incremental/src/lib.rs
@@ -16,8 +16,15 @@ mod persist;
 
 pub use persist::{
     LoadResult, copy_cgu_workproduct_to_incr_comp_cache_dir, finalize_session_directory,
-    in_incr_comp_dir, in_incr_comp_dir_sess, load_query_result_cache, save_dep_graph,
-    save_work_product_index, setup_dep_graph,
+    in_incr_comp_dir, in_incr_comp_dir_sess, load_query_result_cache, save_work_product_index,
+    setup_dep_graph,
 };
+use rustc_middle::util::Providers;
+
+#[allow(missing_docs)]
+pub fn provide(providers: &mut Providers) {
+    providers.hooks.save_dep_graph =
+        |tcx| tcx.sess.time("serialize_dep_graph", || persist::save_dep_graph(tcx.tcx));
+}
 
 rustc_fluent_macro::fluent_messages! { "../messages.ftl" }
diff --git a/compiler/rustc_incremental/src/persist/mod.rs b/compiler/rustc_incremental/src/persist/mod.rs
index 186db0a60a3..f5d5167e0e2 100644
--- a/compiler/rustc_incremental/src/persist/mod.rs
+++ b/compiler/rustc_incremental/src/persist/mod.rs
@@ -12,5 +12,6 @@ mod work_product;
 
 pub use fs::{finalize_session_directory, in_incr_comp_dir, in_incr_comp_dir_sess};
 pub use load::{LoadResult, load_query_result_cache, setup_dep_graph};
-pub use save::{save_dep_graph, save_work_product_index};
+pub(crate) use save::save_dep_graph;
+pub use save::save_work_product_index;
 pub use work_product::copy_cgu_workproduct_to_incr_comp_cache_dir;
diff --git a/compiler/rustc_incremental/src/persist/save.rs b/compiler/rustc_incremental/src/persist/save.rs
index 53726df4597..8fc50ca1b43 100644
--- a/compiler/rustc_incremental/src/persist/save.rs
+++ b/compiler/rustc_incremental/src/persist/save.rs
@@ -25,7 +25,7 @@ use crate::errors;
 ///
 /// This function should only run after all queries have completed.
 /// Trying to execute a query afterwards would attempt to read the result cache we just dropped.
-pub fn save_dep_graph(tcx: TyCtxt<'_>) {
+pub(crate) fn save_dep_graph(tcx: TyCtxt<'_>) {
     debug!("save_dep_graph()");
     tcx.dep_graph.with_ignore(|| {
         let sess = tcx.sess;
diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
index 43a98782016..2905fe688b5 100644
--- a/compiler/rustc_interface/src/passes.rs
+++ b/compiler/rustc_interface/src/passes.rs
@@ -689,10 +689,12 @@ pub static DEFAULT_QUERY_PROVIDERS: LazyLock<Providers> = LazyLock::new(|| {
     rustc_const_eval::provide(providers);
     rustc_middle::hir::provide(providers);
     rustc_borrowck::provide(providers);
+    rustc_incremental::provide(providers);
     rustc_mir_build::provide(providers);
     rustc_mir_transform::provide(providers);
     rustc_monomorphize::provide(providers);
     rustc_privacy::provide(providers);
+    rustc_query_impl::provide(providers);
     rustc_resolve::provide(providers);
     rustc_hir_analysis::provide(providers);
     rustc_hir_typeck::provide(providers);
diff --git a/compiler/rustc_interface/src/queries.rs b/compiler/rustc_interface/src/queries.rs
index b6837ec764f..cd3a2fb7049 100644
--- a/compiler/rustc_interface/src/queries.rs
+++ b/compiler/rustc_interface/src/queries.rs
@@ -12,13 +12,12 @@ use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_middle::arena::Arena;
 use rustc_middle::dep_graph::DepGraph;
 use rustc_middle::ty::{GlobalCtxt, TyCtxt};
-use rustc_serialize::opaque::FileEncodeResult;
 use rustc_session::Session;
 use rustc_session::config::{self, OutputFilenames, OutputType};
 
 use crate::errors::FailedWritingFile;
 use crate::interface::{Compiler, Result};
-use crate::{errors, passes};
+use crate::passes;
 
 /// Represent the result of a query.
 ///
@@ -62,7 +61,7 @@ impl<'a, T> std::ops::DerefMut for QueryResult<'a, T> {
 
 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)
+        (*self.0).borrow().enter(f)
     }
 }
 
@@ -90,8 +89,10 @@ impl<'tcx> Queries<'tcx> {
         }
     }
 
-    pub fn finish(&self) -> FileEncodeResult {
-        if let Some(gcx) = self.gcx_cell.get() { gcx.finish() } else { Ok(0) }
+    pub fn finish(&'tcx self) {
+        if let Some(gcx) = self.gcx_cell.get() {
+            gcx.finish();
+        }
     }
 
     pub fn parse(&self) -> Result<QueryResult<'_, ast::Crate>> {
@@ -209,29 +210,10 @@ impl Compiler {
         let queries = Queries::new(self);
         let ret = f(&queries);
 
-        // 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.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.
-            {
-                let _prof_timer =
-                    queries.compiler.sess.prof.generic_activity("self_profile_alloc_query_strings");
-                gcx.enter(rustc_query_impl::alloc_self_profile_query_strings);
-            }
-
-            self.sess.time("serialize_dep_graph", || gcx.enter(rustc_incremental::save_dep_graph));
-
-            gcx.enter(rustc_query_impl::query_key_hash_verify_all);
-        }
-
         // The timer's lifetime spans the dropping of `queries`, which contains
         // the global context.
         _timer = self.sess.timer("free_global_ctxt");
-        if let Err((path, error)) = queries.finish() {
-            self.sess.dcx().emit_fatal(errors::FailedWritingFile { path: &path, error });
-        }
+        queries.finish();
 
         ret
     }
diff --git a/compiler/rustc_middle/messages.ftl b/compiler/rustc_middle/messages.ftl
index 52c3212ab80..5dd85978f00 100644
--- a/compiler/rustc_middle/messages.ftl
+++ b/compiler/rustc_middle/messages.ftl
@@ -73,6 +73,9 @@ middle_drop_check_overflow =
 
 middle_erroneous_constant = erroneous constant encountered
 
+middle_failed_writing_file =
+    failed to write file {$path}: {$error}"
+
 middle_layout_references_error =
     the type has an unknown layout
 
diff --git a/compiler/rustc_middle/src/error.rs b/compiler/rustc_middle/src/error.rs
index 5c2aa0005d4..6300d856393 100644
--- a/compiler/rustc_middle/src/error.rs
+++ b/compiler/rustc_middle/src/error.rs
@@ -1,5 +1,5 @@
-use std::fmt;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
+use std::{fmt, io};
 
 use rustc_errors::codes::*;
 use rustc_errors::{DiagArgName, DiagArgValue, DiagMessage};
@@ -19,6 +19,13 @@ pub struct DropCheckOverflow<'tcx> {
 }
 
 #[derive(Diagnostic)]
+#[diag(middle_failed_writing_file)]
+pub struct FailedWritingFile<'a> {
+    pub path: &'a Path,
+    pub error: io::Error,
+}
+
+#[derive(Diagnostic)]
 #[diag(middle_opaque_hidden_type_mismatch)]
 pub struct OpaqueHiddenTypeMismatch<'tcx> {
     pub self_ty: Ty<'tcx>,
diff --git a/compiler/rustc_middle/src/hooks/mod.rs b/compiler/rustc_middle/src/hooks/mod.rs
index 742797fdeef..92fa64b0987 100644
--- a/compiler/rustc_middle/src/hooks/mod.rs
+++ b/compiler/rustc_middle/src/hooks/mod.rs
@@ -108,6 +108,19 @@ declare_hooks! {
     /// Returns `true` if we should codegen an instance in the local crate, or returns `false` if we
     /// can just link to the upstream crate and therefore don't need a mono item.
     hook should_codegen_locally(instance: crate::ty::Instance<'tcx>) -> bool;
+
+    hook alloc_self_profile_query_strings() -> ();
+
+    /// Saves and writes the DepGraph to the file system.
+    ///
+    /// This function saves both the dep-graph and the query result cache,
+    /// and drops the result cache.
+    ///
+    /// This function should only run after all queries have completed.
+    /// Trying to execute a query afterwards would attempt to read the result cache we just dropped.
+    hook save_dep_graph() -> ();
+
+    hook query_key_hash_verify_all() -> ();
 }
 
 #[cold]
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index 5bf62a17c8e..2c867e6b8ca 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -1371,8 +1371,17 @@ impl<'tcx> GlobalCtxt<'tcx> {
         tls::enter_context(&icx, || f(icx.tcx))
     }
 
-    pub fn finish(&self) -> FileEncodeResult {
-        self.dep_graph.finish_encoding()
+    pub fn finish(&'tcx self) {
+        // 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.
+        self.enter(|tcx| tcx.alloc_self_profile_query_strings());
+
+        self.enter(|tcx| tcx.save_dep_graph());
+        self.enter(|tcx| tcx.query_key_hash_verify_all());
+
+        if let Err((path, error)) = self.dep_graph.finish_encoding() {
+            self.sess.dcx().emit_fatal(crate::error::FailedWritingFile { path: &path, error });
+        }
     }
 }
 
diff --git a/compiler/rustc_query_impl/src/lib.rs b/compiler/rustc_query_impl/src/lib.rs
index df898e0587f..a1917fed4d9 100644
--- a/compiler/rustc_query_impl/src/lib.rs
+++ b/compiler/rustc_query_impl/src/lib.rs
@@ -222,3 +222,9 @@ pub fn query_system<'tcx>(
 }
 
 rustc_middle::rustc_query_append! { define_queries! }
+
+pub fn provide(providers: &mut rustc_middle::util::Providers) {
+    providers.hooks.alloc_self_profile_query_strings =
+        |tcx| alloc_self_profile_query_strings(tcx.tcx);
+    providers.hooks.query_key_hash_verify_all = |tcx| query_key_hash_verify_all(tcx.tcx);
+}
diff --git a/compiler/rustc_query_impl/src/profiling_support.rs b/compiler/rustc_query_impl/src/profiling_support.rs
index 5f989b264aa..7d80e54bf87 100644
--- a/compiler/rustc_query_impl/src/profiling_support.rs
+++ b/compiler/rustc_query_impl/src/profiling_support.rs
@@ -252,6 +252,8 @@ pub fn alloc_self_profile_query_strings(tcx: TyCtxt<'_>) {
         return;
     }
 
+    let _prof_timer = tcx.sess.prof.generic_activity("self_profile_alloc_query_strings");
+
     let mut string_cache = QueryKeyStringCache::new();
 
     for alloc in super::ALLOC_SELF_PROFILE_QUERY_STRINGS.iter() {
diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index a326e8583ea..d4b034ea219 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -313,6 +313,7 @@ macro_rules! optional {
 macro_rules! run_driver {
     ($args:expr, $callback:expr $(, $with_tcx:ident)?) => {{
         use rustc_driver::{Callbacks, Compilation, RunCompiler};
+        use rustc_middle::ty::TyCtxt;
         use rustc_interface::{interface, Queries};
         use stable_mir::CompilerError;
         use std::ops::ControlFlow;
@@ -373,23 +374,21 @@ macro_rules! run_driver {
             fn after_analysis<'tcx>(
                 &mut self,
                 _compiler: &interface::Compiler,
-                queries: &'tcx Queries<'tcx>,
+                tcx: TyCtxt<'tcx>,
             ) -> Compilation {
-                queries.global_ctxt().unwrap().enter(|tcx| {
-                    if let Some(callback) = self.callback.take() {
-                        rustc_internal::run(tcx, || {
-                            self.result = Some(callback($(optional!($with_tcx tcx))?));
-                        })
-                        .unwrap();
-                        if self.result.as_ref().is_some_and(|val| val.is_continue()) {
-                            Compilation::Continue
-                        } else {
-                            Compilation::Stop
-                        }
-                    } else {
+                if let Some(callback) = self.callback.take() {
+                    rustc_internal::run(tcx, || {
+                        self.result = Some(callback($(optional!($with_tcx tcx))?));
+                    })
+                    .unwrap();
+                    if self.result.as_ref().is_some_and(|val| val.is_continue()) {
                         Compilation::Continue
+                    } else {
+                        Compilation::Stop
                     }
-                })
+                } else {
+                    Compilation::Continue
+                }
             }
         }
 
diff --git a/src/tools/miri/src/bin/miri.rs b/src/tools/miri/src/bin/miri.rs
index 357c50889c4..c61c62c73da 100644
--- a/src/tools/miri/src/bin/miri.rs
+++ b/src/tools/miri/src/bin/miri.rs
@@ -73,51 +73,47 @@ impl rustc_driver::Callbacks for MiriCompilerCalls {
     fn after_analysis<'tcx>(
         &mut self,
         _: &rustc_interface::interface::Compiler,
-        queries: &'tcx rustc_interface::Queries<'tcx>,
+        tcx: TyCtxt<'tcx>,
     ) -> Compilation {
-        queries.global_ctxt().unwrap().enter(|tcx| {
-            if tcx.sess.dcx().has_errors_or_delayed_bugs().is_some() {
-                tcx.dcx().fatal("miri cannot be run on programs that fail compilation");
-            }
+        if tcx.sess.dcx().has_errors_or_delayed_bugs().is_some() {
+            tcx.dcx().fatal("miri cannot be run on programs that fail compilation");
+        }
 
-            let early_dcx = EarlyDiagCtxt::new(tcx.sess.opts.error_format);
-            init_late_loggers(&early_dcx, tcx);
-            if !tcx.crate_types().contains(&CrateType::Executable) {
-                tcx.dcx().fatal("miri only makes sense on bin crates");
-            }
+        let early_dcx = EarlyDiagCtxt::new(tcx.sess.opts.error_format);
+        init_late_loggers(&early_dcx, tcx);
+        if !tcx.crate_types().contains(&CrateType::Executable) {
+            tcx.dcx().fatal("miri only makes sense on bin crates");
+        }
 
-            let (entry_def_id, entry_type) = entry_fn(tcx);
-            let mut config = self.miri_config.clone();
+        let (entry_def_id, entry_type) = entry_fn(tcx);
+        let mut config = self.miri_config.clone();
 
-            // Add filename to `miri` arguments.
-            config.args.insert(0, tcx.sess.io.input.filestem().to_string());
+        // Add filename to `miri` arguments.
+        config.args.insert(0, tcx.sess.io.input.filestem().to_string());
 
-            // Adjust working directory for interpretation.
-            if let Some(cwd) = env::var_os("MIRI_CWD") {
-                env::set_current_dir(cwd).unwrap();
-            }
+        // Adjust working directory for interpretation.
+        if let Some(cwd) = env::var_os("MIRI_CWD") {
+            env::set_current_dir(cwd).unwrap();
+        }
 
-            if tcx.sess.opts.optimize != OptLevel::No {
-                tcx.dcx().warn("Miri does not support optimizations: the opt-level is ignored. The only effect \
+        if tcx.sess.opts.optimize != OptLevel::No {
+            tcx.dcx().warn("Miri does not support optimizations: the opt-level is ignored. The only effect \
                     of selecting a Cargo profile that enables optimizations (such as --release) is to apply \
                     its remaining settings, such as whether debug assertions and overflow checks are enabled.");
-            }
-            if tcx.sess.mir_opt_level() > 0 {
-                tcx.dcx().warn("You have explicitly enabled MIR optimizations, overriding Miri's default \
+        }
+        if tcx.sess.mir_opt_level() > 0 {
+            tcx.dcx().warn("You have explicitly enabled MIR optimizations, overriding Miri's default \
                     which is to completely disable them. Any optimizations may hide UB that Miri would \
                     otherwise detect, and it is not necessarily possible to predict what kind of UB will \
                     be missed. If you are enabling optimizations to make Miri run faster, we advise using \
                     cfg(miri) to shrink your workload instead. The performance benefit of enabling MIR \
                     optimizations is usually marginal at best.");
-            }
+        }
 
-            if let Some(return_code) = miri::eval_entry(tcx, entry_def_id, entry_type, config) {
-                std::process::exit(
-                    i32::try_from(return_code).expect("Return value was too large!"),
-                );
-            }
-            tcx.dcx().abort_if_errors();
-        });
+        if let Some(return_code) = miri::eval_entry(tcx, entry_def_id, entry_type, config) {
+            std::process::exit(i32::try_from(return_code).expect("Return value was too large!"));
+        }
+        tcx.dcx().abort_if_errors();
 
         Compilation::Stop
     }
@@ -193,20 +189,18 @@ impl rustc_driver::Callbacks for MiriBeRustCompilerCalls {
     fn after_analysis<'tcx>(
         &mut self,
         _: &rustc_interface::interface::Compiler,
-        queries: &'tcx rustc_interface::Queries<'tcx>,
+        tcx: TyCtxt<'tcx>,
     ) -> Compilation {
-        queries.global_ctxt().unwrap().enter(|tcx| {
-            if self.target_crate {
-                // cargo-miri has patched the compiler flags to make these into check-only builds,
-                // but we are still emulating regular rustc builds, which would perform post-mono
-                // const-eval during collection. So let's also do that here, even if we might be
-                // running with `--emit=metadata`. In particular this is needed to make
-                // `compile_fail` doc tests trigger post-mono errors.
-                // In general `collect_and_partition_mono_items` is not safe to call in check-only
-                // builds, but we are setting `-Zalways-encode-mir` which avoids those issues.
-                let _ = tcx.collect_and_partition_mono_items(());
-            }
-        });
+        if self.target_crate {
+            // cargo-miri has patched the compiler flags to make these into check-only builds,
+            // but we are still emulating regular rustc builds, which would perform post-mono
+            // const-eval during collection. So let's also do that here, even if we might be
+            // running with `--emit=metadata`. In particular this is needed to make
+            // `compile_fail` doc tests trigger post-mono errors.
+            // In general `collect_and_partition_mono_items` is not safe to call in check-only
+            // builds, but we are setting `-Zalways-encode-mir` which avoids those issues.
+            let _ = tcx.collect_and_partition_mono_items(());
+        }
         Compilation::Continue
     }
 }
diff --git a/tests/ui-fulldeps/obtain-borrowck.rs b/tests/ui-fulldeps/obtain-borrowck.rs
index e6c703addd9..af98f93297b 100644
--- a/tests/ui-fulldeps/obtain-borrowck.rs
+++ b/tests/ui-fulldeps/obtain-borrowck.rs
@@ -25,19 +25,20 @@ extern crate rustc_interface;
 extern crate rustc_middle;
 extern crate rustc_session;
 
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::thread_local;
+
 use rustc_borrowck::consumers::{self, BodyWithBorrowckFacts, ConsumerOptions};
 use rustc_driver::Compilation;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::LocalDefId;
+use rustc_interface::Config;
 use rustc_interface::interface::Compiler;
-use rustc_interface::{Config, Queries};
 use rustc_middle::query::queries::mir_borrowck::ProvidedValue;
 use rustc_middle::ty::TyCtxt;
 use rustc_middle::util::Providers;
 use rustc_session::Session;
-use std::cell::RefCell;
-use std::collections::HashMap;
-use std::thread_local;
 
 fn main() {
     let exit_code = rustc_driver::catch_with_exit_code(move || {
@@ -63,55 +64,49 @@ impl rustc_driver::Callbacks for CompilerCalls {
 
     // In this callback we trigger borrow checking of all functions and obtain
     // the result.
-    fn after_analysis<'tcx>(
-        &mut self,
-        compiler: &Compiler,
-        queries: &'tcx Queries<'tcx>,
-    ) -> Compilation {
-        compiler.sess.dcx().abort_if_errors();
-        queries.global_ctxt().unwrap().enter(|tcx| {
-            // Collect definition ids of MIR bodies.
-            let hir = tcx.hir();
-            let mut bodies = Vec::new();
-
-            let crate_items = tcx.hir_crate_items(());
-            for id in crate_items.free_items() {
-                if matches!(tcx.def_kind(id.owner_id), DefKind::Fn) {
-                    bodies.push(id.owner_id);
-                }
+    fn after_analysis<'tcx>(&mut self, _compiler: &Compiler, tcx: TyCtxt<'tcx>) -> Compilation {
+        tcx.sess.dcx().abort_if_errors();
+        // Collect definition ids of MIR bodies.
+        let hir = tcx.hir();
+        let mut bodies = Vec::new();
+
+        let crate_items = tcx.hir_crate_items(());
+        for id in crate_items.free_items() {
+            if matches!(tcx.def_kind(id.owner_id), DefKind::Fn) {
+                bodies.push(id.owner_id);
             }
-
-            for id in crate_items.trait_items() {
-                if matches!(tcx.def_kind(id.owner_id), DefKind::AssocFn) {
-                    let trait_item = hir.trait_item(id);
-                    if let rustc_hir::TraitItemKind::Fn(_, trait_fn) = &trait_item.kind {
-                        if let rustc_hir::TraitFn::Provided(_) = trait_fn {
-                            bodies.push(trait_item.owner_id);
-                        }
+        }
+
+        for id in crate_items.trait_items() {
+            if matches!(tcx.def_kind(id.owner_id), DefKind::AssocFn) {
+                let trait_item = hir.trait_item(id);
+                if let rustc_hir::TraitItemKind::Fn(_, trait_fn) = &trait_item.kind {
+                    if let rustc_hir::TraitFn::Provided(_) = trait_fn {
+                        bodies.push(trait_item.owner_id);
                     }
                 }
             }
+        }
 
-            for id in crate_items.impl_items() {
-                if matches!(tcx.def_kind(id.owner_id), DefKind::AssocFn) {
-                    bodies.push(id.owner_id);
-                }
-            }
-
-            // Trigger borrow checking of all bodies.
-            for def_id in bodies {
-                let _ = tcx.optimized_mir(def_id);
-            }
-
-            // See what bodies were borrow checked.
-            let mut bodies = get_bodies(tcx);
-            bodies.sort_by(|(def_id1, _), (def_id2, _)| def_id1.cmp(def_id2));
-            println!("Bodies retrieved for:");
-            for (def_id, body) in bodies {
-                println!("{}", def_id);
-                assert!(body.input_facts.unwrap().cfg_edge.len() > 0);
+        for id in crate_items.impl_items() {
+            if matches!(tcx.def_kind(id.owner_id), DefKind::AssocFn) {
+                bodies.push(id.owner_id);
             }
-        });
+        }
+
+        // Trigger borrow checking of all bodies.
+        for def_id in bodies {
+            let _ = tcx.optimized_mir(def_id);
+        }
+
+        // See what bodies were borrow checked.
+        let mut bodies = get_bodies(tcx);
+        bodies.sort_by(|(def_id1, _), (def_id2, _)| def_id1.cmp(def_id2));
+        println!("Bodies retrieved for:");
+        for (def_id, body) in bodies {
+            println!("{}", def_id);
+            assert!(body.input_facts.unwrap().cfg_edge.len() > 0);
+        }
 
         Compilation::Continue
     }
diff --git a/tests/ui-fulldeps/stable-mir/check_abi.rs b/tests/ui-fulldeps/stable-mir/check_abi.rs
index 5b7da7bb129..8caf3032afc 100644
--- a/tests/ui-fulldeps/stable-mir/check_abi.rs
+++ b/tests/ui-fulldeps/stable-mir/check_abi.rs
@@ -11,6 +11,7 @@
 #![feature(ascii_char, ascii_char_variants)]
 
 extern crate rustc_hir;
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_allocation.rs b/tests/ui-fulldeps/stable-mir/check_allocation.rs
index 1e2f640f39f..072c8ba6a44 100644
--- a/tests/ui-fulldeps/stable-mir/check_allocation.rs
+++ b/tests/ui-fulldeps/stable-mir/check_allocation.rs
@@ -13,6 +13,7 @@
 #![feature(ascii_char, ascii_char_variants)]
 
 extern crate rustc_hir;
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_attribute.rs b/tests/ui-fulldeps/stable-mir/check_attribute.rs
index 131fd99ebaa..22481e275a9 100644
--- a/tests/ui-fulldeps/stable-mir/check_attribute.rs
+++ b/tests/ui-fulldeps/stable-mir/check_attribute.rs
@@ -9,6 +9,7 @@
 #![feature(rustc_private)]
 
 extern crate rustc_hir;
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_binop.rs b/tests/ui-fulldeps/stable-mir/check_binop.rs
index 3b52d88de3c..8c44e285108 100644
--- a/tests/ui-fulldeps/stable-mir/check_binop.rs
+++ b/tests/ui-fulldeps/stable-mir/check_binop.rs
@@ -9,6 +9,7 @@
 #![feature(rustc_private)]
 
 extern crate rustc_hir;
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_crate_defs.rs b/tests/ui-fulldeps/stable-mir/check_crate_defs.rs
index e039ca07dd4..ed093903381 100644
--- a/tests/ui-fulldeps/stable-mir/check_crate_defs.rs
+++ b/tests/ui-fulldeps/stable-mir/check_crate_defs.rs
@@ -10,6 +10,7 @@
 #![feature(assert_matches)]
 
 extern crate rustc_hir;
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_def_ty.rs b/tests/ui-fulldeps/stable-mir/check_def_ty.rs
index ec3cf1753e2..482dbd22d5f 100644
--- a/tests/ui-fulldeps/stable-mir/check_def_ty.rs
+++ b/tests/ui-fulldeps/stable-mir/check_def_ty.rs
@@ -11,6 +11,7 @@
 #![feature(rustc_private)]
 #![feature(assert_matches)]
 
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_defs.rs b/tests/ui-fulldeps/stable-mir/check_defs.rs
index 3402b345818..bf1f1a2ceab 100644
--- a/tests/ui-fulldeps/stable-mir/check_defs.rs
+++ b/tests/ui-fulldeps/stable-mir/check_defs.rs
@@ -10,6 +10,7 @@
 #![feature(rustc_private)]
 #![feature(assert_matches)]
 
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_instance.rs b/tests/ui-fulldeps/stable-mir/check_instance.rs
index 7d63e202fa6..464350b1045 100644
--- a/tests/ui-fulldeps/stable-mir/check_instance.rs
+++ b/tests/ui-fulldeps/stable-mir/check_instance.rs
@@ -10,6 +10,7 @@
 #![feature(rustc_private)]
 #![feature(assert_matches)]
 
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_intrinsics.rs b/tests/ui-fulldeps/stable-mir/check_intrinsics.rs
index 3534228f73e..6edebaf756c 100644
--- a/tests/ui-fulldeps/stable-mir/check_intrinsics.rs
+++ b/tests/ui-fulldeps/stable-mir/check_intrinsics.rs
@@ -13,6 +13,7 @@
 #![feature(rustc_private)]
 #![feature(assert_matches)]
 
+extern crate rustc_middle;
 extern crate rustc_hir;
 #[macro_use]
 extern crate rustc_smir;
diff --git a/tests/ui-fulldeps/stable-mir/check_item_kind.rs b/tests/ui-fulldeps/stable-mir/check_item_kind.rs
index 91baa074c10..23b54e6c60b 100644
--- a/tests/ui-fulldeps/stable-mir/check_item_kind.rs
+++ b/tests/ui-fulldeps/stable-mir/check_item_kind.rs
@@ -10,6 +10,7 @@
 #![feature(rustc_private)]
 #![feature(assert_matches)]
 
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_normalization.rs b/tests/ui-fulldeps/stable-mir/check_normalization.rs
index 72e410f8080..928173b154b 100644
--- a/tests/ui-fulldeps/stable-mir/check_normalization.rs
+++ b/tests/ui-fulldeps/stable-mir/check_normalization.rs
@@ -9,6 +9,7 @@
 
 #![feature(rustc_private)]
 
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_trait_queries.rs b/tests/ui-fulldeps/stable-mir/check_trait_queries.rs
index 8721f243587..304a7ce9255 100644
--- a/tests/ui-fulldeps/stable-mir/check_trait_queries.rs
+++ b/tests/ui-fulldeps/stable-mir/check_trait_queries.rs
@@ -10,6 +10,7 @@
 #![feature(rustc_private)]
 #![feature(assert_matches)]
 
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_transform.rs b/tests/ui-fulldeps/stable-mir/check_transform.rs
index 40217b9aa95..bcf79c456b0 100644
--- a/tests/ui-fulldeps/stable-mir/check_transform.rs
+++ b/tests/ui-fulldeps/stable-mir/check_transform.rs
@@ -11,6 +11,7 @@
 #![feature(ascii_char, ascii_char_variants)]
 
 extern crate rustc_hir;
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/check_ty_fold.rs b/tests/ui-fulldeps/stable-mir/check_ty_fold.rs
index 0715e0cfc52..e21508c9b46 100644
--- a/tests/ui-fulldeps/stable-mir/check_ty_fold.rs
+++ b/tests/ui-fulldeps/stable-mir/check_ty_fold.rs
@@ -11,6 +11,7 @@
 #![feature(rustc_private)]
 #![feature(assert_matches)]
 
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/compilation-result.rs b/tests/ui-fulldeps/stable-mir/compilation-result.rs
index 286bbd7c594..d921de73f43 100644
--- a/tests/ui-fulldeps/stable-mir/compilation-result.rs
+++ b/tests/ui-fulldeps/stable-mir/compilation-result.rs
@@ -10,6 +10,7 @@
 #![feature(rustc_private)]
 #![feature(assert_matches)]
 
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/crate-info.rs b/tests/ui-fulldeps/stable-mir/crate-info.rs
index 6b458c5d923..53be8eb10c1 100644
--- a/tests/ui-fulldeps/stable-mir/crate-info.rs
+++ b/tests/ui-fulldeps/stable-mir/crate-info.rs
@@ -11,6 +11,7 @@
 #![feature(assert_matches)]
 
 extern crate rustc_hir;
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/projections.rs b/tests/ui-fulldeps/stable-mir/projections.rs
index a8bf4c1d399..fdb7eeed1b0 100644
--- a/tests/ui-fulldeps/stable-mir/projections.rs
+++ b/tests/ui-fulldeps/stable-mir/projections.rs
@@ -11,6 +11,7 @@
 #![feature(assert_matches)]
 
 extern crate rustc_hir;
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;
diff --git a/tests/ui-fulldeps/stable-mir/smir_visitor.rs b/tests/ui-fulldeps/stable-mir/smir_visitor.rs
index f1bc03781b9..666000d3b07 100644
--- a/tests/ui-fulldeps/stable-mir/smir_visitor.rs
+++ b/tests/ui-fulldeps/stable-mir/smir_visitor.rs
@@ -10,6 +10,7 @@
 #![feature(rustc_private)]
 #![feature(assert_matches)]
 
+extern crate rustc_middle;
 #[macro_use]
 extern crate rustc_smir;
 extern crate rustc_driver;