Optionally don't steal the THIR

5 files changed, 25 insertions, 4 deletions

diff --git a/compiler/rustc_interface/src/tests.rs b/compiler/rustc_interface/src/tests.rs
index 068d96c860f..558f13a832c 100644
--- a/compiler/rustc_interface/src/tests.rs
+++ b/compiler/rustc_interface/src/tests.rs
@@ -715,6 +715,7 @@ fn test_unstable_options_tracking_hash() {
     untracked!(no_analysis, true);
     untracked!(no_leak_check, true);
     untracked!(no_parallel_backend, true);
+    untracked!(no_steal_thir, true);
     untracked!(parse_crate_root_only, true);
     // `pre_link_arg` is omitted because it just forwards to `pre_link_args`.
     untracked!(pre_link_args, vec![String::from("abc"), String::from("def")]);
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 279033ee072..542653efd30 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -535,7 +535,8 @@ rustc_queries! {
         separate_provide_extern
     }
 
-    /// Fetch the THIR for a given body.
+    /// Fetch the THIR for a given body. The THIR body gets stolen by unsafety checking unless
+    /// `-Zno-steal-thir` is on.
     query thir_body(key: LocalDefId) -> Result<(&'tcx Steal<thir::Thir<'tcx>>, thir::ExprId), ErrorGuaranteed> {
         // Perf tests revealed that hashing THIR is inefficient (see #85729).
         no_hash
diff --git a/compiler/rustc_mir_build/src/check_unsafety.rs b/compiler/rustc_mir_build/src/check_unsafety.rs
index 9d0681b19b9..d5061b71699 100644
--- a/compiler/rustc_mir_build/src/check_unsafety.rs
+++ b/compiler/rustc_mir_build/src/check_unsafety.rs
@@ -201,9 +201,14 @@ impl<'tcx> UnsafetyVisitor<'_, 'tcx> {
     /// Handle closures/coroutines/inline-consts, which is unsafecked with their parent body.
     fn visit_inner_body(&mut self, def: LocalDefId) {
         if let Ok((inner_thir, expr)) = self.tcx.thir_body(def) {
-            // Runs all other queries that depend on THIR.
+            // Run all other queries that depend on THIR.
             self.tcx.ensure_done().mir_built(def);
-            let inner_thir = &inner_thir.steal();
+            let inner_thir = if self.tcx.sess.opts.unstable_opts.no_steal_thir {
+                &inner_thir.borrow()
+            } else {
+                // We don't have other use for the THIR. Steal it to reduce memory usage.
+                &inner_thir.steal()
+            };
             let hir_context = self.tcx.local_def_id_to_hir_id(def);
             let safety_context = mem::replace(&mut self.safety_context, SafetyContext::Safe);
             let mut inner_visitor = UnsafetyVisitor {
@@ -1157,7 +1162,12 @@ pub(crate) fn check_unsafety(tcx: TyCtxt<'_>, def: LocalDefId) {
     let Ok((thir, expr)) = tcx.thir_body(def) else { return };
     // Runs all other queries that depend on THIR.
     tcx.ensure_done().mir_built(def);
-    let thir = &thir.steal();
+    let thir = if tcx.sess.opts.unstable_opts.no_steal_thir {
+        &thir.borrow()
+    } else {
+        // We don't have other use for the THIR. Steal it to reduce memory usage.
+        &thir.steal()
+    };
 
     let hir_id = tcx.local_def_id_to_hir_id(def);
     let safety_context = tcx.hir_fn_sig_by_hir_id(hir_id).map_or(SafetyContext::Safe, |fn_sig| {
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
index 5b4068740a1..12fa05118ca 100644
--- a/compiler/rustc_session/src/options.rs
+++ b/compiler/rustc_session/src/options.rs
@@ -2366,6 +2366,8 @@ options! {
         "run LLVM in non-parallel mode (while keeping codegen-units and ThinLTO)"),
     no_profiler_runtime: bool = (false, parse_no_value, [TRACKED],
         "prevent automatic injection of the profiler_builtins crate"),
+    no_steal_thir: bool = (false, parse_bool, [UNTRACKED],
+        "don't steal the THIR when we're done with it; useful for rustc drivers (default: no)"),
     no_trait_vptr: bool = (false, parse_no_value, [TRACKED],
         "disable generation of trait vptr in vtable for upcasting"),
     no_unique_section_names: bool = (false, parse_bool, [TRACKED],
diff --git a/src/doc/unstable-book/src/compiler-flags/no-steal-thir.md b/src/doc/unstable-book/src/compiler-flags/no-steal-thir.md
new file mode 100644
index 00000000000..d83677d1711
--- /dev/null
+++ b/src/doc/unstable-book/src/compiler-flags/no-steal-thir.md
@@ -0,0 +1,7 @@
+# `no-steal-thir`
+
+By default, to save on memory, the THIR body (obtained from the `tcx.thir_body` query) is stolen
+once no longer used. This is inconvenient for authors of rustc drivers who want to access the THIR.
+
+This option disables the stealing. This has no observable effect on compiler behavior, only on
+memory usage.