Box data inside Steal

compiler/rustc_ast_lowering/src/lib.rs

@@ -715,7 +715,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             .resolver
             .disambiguators
             .get(&def_id)
-            .map(|s| s.steal())
+            .map(|s| *s.steal())
             .unwrap_or_else(|| PerParentDisambiguatorState::new(def_id));
 
         let disambiguator = std::mem::replace(&mut self.current_disambiguator, new_disambig);

compiler/rustc_data_structures/src/steal.rs

@@ -1,5 +1,5 @@
 use crate::stable_hash::{StableHash, StableHashCtxt, StableHasher};
-use crate::sync::{MappedReadGuard, MappedWriteGuard, ReadGuard, RwLock, WriteGuard};
+use crate::sync::{MappedReadGuard, ReadGuard, RwLock};
 
 /// The `Steal` struct is intended to used as the value for a query.
 /// Specifically, we sometimes have queries (*cough* MIR *cough*)
@@ -23,11 +23,15 @@ use crate::sync::{MappedReadGuard, MappedWriteGuard, ReadGuard, RwLock, WriteGua
 // FIXME(#41710): what is the best way to model linear queries?
 #[derive(Debug)]
 pub struct Steal<T> {
-    value: RwLock<Option<T>>,
+    value: RwLock<Option<Box<T>>>,
 }
 
 impl<T> Steal<T> {
     pub fn new(value: T) -> Self {
+        Steal { value: RwLock::new(Some(Box::new(value))) }
+    }
+
+    pub fn new_from_box(value: Box<T>) -> Self {
         Steal { value: RwLock::new(Some(value)) }
     }
 
@@ -37,24 +41,11 @@ impl<T> Steal<T> {
         if borrow.is_none() {
             panic!("attempted to read from stolen value: {}", std::any::type_name::<T>());
         }
-        ReadGuard::map(borrow, |opt| opt.as_ref().unwrap())
-    }
-
-    /// An escape hatch for rustc drivers to mutate `Steal` caches.
-    ///
-    /// Use at your own risk. This can badly break incremental compilation
-    /// and anything else that relies on the immutability of query caches.
-    #[track_caller]
-    pub fn risky_hack_borrow_mut(&self) -> MappedWriteGuard<'_, T> {
-        let borrow = self.value.borrow_mut();
-        if borrow.is_none() {
-            panic!("attempted to read from stolen value: {}", std::any::type_name::<T>());
-        }
-        WriteGuard::map(borrow, |opt| opt.as_mut().unwrap())
+        ReadGuard::map(borrow, |opt| opt.as_deref().unwrap())
     }
 
     #[track_caller]
-    pub fn steal(&self) -> T {
+    pub fn steal(&self) -> Box<T> {
         let value_ref = &mut *self.value.try_write().expect("stealing value which is locked");
         let value = value_ref.take();
         value.expect("attempt to steal from stolen value")

compiler/rustc_interface/src/passes.rs

@@ -472,7 +472,7 @@ fn early_lint_checks(tcx: TyCtxt<'_>, (): ()) {
         false,
         lint_store,
         tcx.registered_tools(()),
-        Some(lint_buffer),
+        Some(*lint_buffer),
         rustc_lint::BuiltinCombinedEarlyLintPass::new(),
         (&**krate, &*krate.attrs),
     )
@@ -784,7 +784,7 @@ fn resolver_for_lowering_raw<'tcx>(
 ) -> (&'tcx Steal<(ty::ResolverAstLowering<'tcx>, Arc<ast::Crate>)>, &'tcx ty::ResolverGlobalCtxt) {
     let arenas = Resolver::arenas();
     let _ = tcx.registered_tools(()); // Uses `crate_for_resolver`.
-    let (krate, pre_configured_attrs) = tcx.crate_for_resolver(()).steal();
+    let (krate, pre_configured_attrs) = *tcx.crate_for_resolver(()).steal();
     let mut resolver = Resolver::new(
         tcx,
         &pre_configured_attrs,
@@ -1056,7 +1056,7 @@ impl<'a, 'tcx> Diagnostic<'a, ()> for DiagCallback<'tcx> {
 pub fn emit_delayed_lints(tcx: TyCtxt<'_>) {
     for owner_id in tcx.hir_crate_items(()).delayed_lint_items() {
         if let Some(delayed_lints) = tcx.opt_ast_lowering_delayed_lints(owner_id) {
-            for lint in delayed_lints.steal() {
+            for lint in delayed_lints.steal().into_iter() {
                 tcx.emit_node_span_lint(
                     lint.lint_id.lint,
                     lint.id,

compiler/rustc_middle/src/arena.rs

@@ -13,6 +13,7 @@ macro_rules! arena_types {
             [] steal_thir: rustc_data_structures::steal::Steal<rustc_middle::thir::Thir<'tcx>>,
             [] steal_mir: rustc_data_structures::steal::Steal<rustc_middle::mir::Body<'tcx>>,
             [decode] mir: rustc_middle::mir::Body<'tcx>,
+            [] mir_box: Box<rustc_middle::mir::Body<'tcx>>,
             [] steal_promoted:
                 rustc_data_structures::steal::Steal<
                     rustc_index::IndexVec<

compiler/rustc_middle/src/ty/context.rs

@@ -843,8 +843,8 @@ impl<'tcx> TyCtxt<'tcx> {
         self.arena.alloc(Steal::new(thir))
     }
 
-    pub fn alloc_steal_mir(self, mir: Body<'tcx>) -> &'tcx Steal<Body<'tcx>> {
-        self.arena.alloc(Steal::new(mir))
+    pub fn alloc_steal_mir(self, mir: Box<Body<'tcx>>) -> &'tcx Steal<Body<'tcx>> {
+        self.arena.alloc(Steal::new_from_box(mir))
     }
 
     pub fn alloc_steal_promoted(

compiler/rustc_mir_build/src/check_unsafety.rs

@@ -183,10 +183,10 @@ impl<'tcx> UnsafetyVisitor<'_, 'tcx> {
             // Run all other queries that depend on THIR.
             self.tcx.ensure_done().mir_built(def);
             let inner_thir = if self.tcx.sess.opts.unstable_opts.no_steal_thir {
-                &inner_thir.borrow()
+                &*inner_thir.borrow()
             } else {
                 // We don't have other use for the THIR. Steal it to reduce memory usage.
-                &inner_thir.steal()
+                &*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);
@@ -1064,10 +1064,10 @@ pub(crate) fn check_unsafety(tcx: TyCtxt<'_>, def: LocalDefId) {
     // Runs all other queries that depend on THIR.
     tcx.ensure_done().mir_built(def);
     let thir = if tcx.sess.opts.unstable_opts.no_steal_thir {
-        &thir.borrow()
+        &*thir.borrow()
     } else {
         // We don't have other use for the THIR. Steal it to reduce memory usage.
-        &thir.steal()
+        &*thir.steal()
     };
 
     let hir_id = tcx.local_def_id_to_hir_id(def);

compiler/rustc_mir_transform/src/lib.rs

@@ -233,9 +233,9 @@ pub fn provide(providers: &mut Providers) {
 
 fn remap_mir_for_const_eval_select<'tcx>(
     tcx: TyCtxt<'tcx>,
-    mut body: Body<'tcx>,
+    body: &mut Body<'tcx>,
     context: hir::Constness,
-) -> Body<'tcx> {
+) {
     for bb in body.basic_blocks.as_mut().iter_mut() {
         let terminator = bb.terminator.as_mut().expect("invalid terminator");
         match terminator.kind {
@@ -300,7 +300,6 @@ fn remap_mir_for_const_eval_select<'tcx>(
             _ => {}
         }
     }
-    body
 }
 
 fn take_array<T, const N: usize>(b: &mut Box<[T]>) -> Result<[T; N], Box<[T]>> {
@@ -381,20 +380,18 @@ fn mir_const_qualif(tcx: TyCtxt<'_>, def: LocalDefId) -> ConstQualifs {
 fn mir_built(tcx: TyCtxt<'_>, def: LocalDefId) -> &Steal<Body<'_>> {
     // Delegate to the main MIR building code in the `rustc_mir_build` crate.
     // This is the one place that is allowed to call `build_mir_inner_impl`.
-    let mut body = tcx.build_mir_inner_impl(def);
+    let mut body = Box::new(tcx.build_mir_inner_impl(def));
+
+    pass_manager::dump_mir_for_phase_change(tcx, &body);
 
     // Identifying trivial consts based on their mir_built is easy, but a little wasteful.
     // Trying to push this logic earlier in the compiler and never even produce the Body would
     // probably improve compile time.
-    if trivial_const::trivial_const(tcx, def, || &body).is_some() {
+    if trivial_const::trivial_const(tcx, def, || &*body).is_some() {
         // Skip all the passes below for trivial consts.
-        let body = tcx.alloc_steal_mir(body);
-        pass_manager::dump_mir_for_phase_change(tcx, &body.borrow());
-        return body;
+        return tcx.alloc_steal_mir(body);
     }
 
-    pass_manager::dump_mir_for_phase_change(tcx, &body);
-
     pm::run_passes(
         tcx,
         &mut body,
@@ -480,28 +477,28 @@ fn mir_promoted(
 /// Compute the MIR that is used during CTFE (and thus has no optimizations run on it)
 fn mir_for_ctfe(tcx: TyCtxt<'_>, def_id: LocalDefId) -> &Body<'_> {
     debug_assert!(!tcx.is_trivial_const(def_id), "Tried to get mir_for_ctfe of a trivial const");
-    tcx.arena.alloc(inner_mir_for_ctfe(tcx, def_id))
+    &**tcx.arena.alloc(inner_mir_for_ctfe(tcx, def_id))
 }
 
-fn inner_mir_for_ctfe(tcx: TyCtxt<'_>, def: LocalDefId) -> Body<'_> {
+fn inner_mir_for_ctfe(tcx: TyCtxt<'_>, def: LocalDefId) -> Box<Body<'_>> {
     if tcx.is_constructor(def.to_def_id()) {
         // There's no reason to run all of the MIR passes on constructors when
         // we can just output the MIR we want directly. This also saves const
         // qualification and borrow checking the trouble of special casing
         // constructors.
-        return shim::build_adt_ctor(tcx, def.to_def_id());
+        return Box::new(shim::build_adt_ctor(tcx, def.to_def_id()));
     }
 
     let body = tcx.mir_drops_elaborated_and_const_checked(def);
-    let body = match tcx.hir_body_const_context(def) {
+    let mut body = match tcx.hir_body_const_context(def) {
         // consts and statics do not have `optimized_mir`, so we can steal the body instead of
         // cloning it.
         Some(hir::ConstContext::Const { .. } | hir::ConstContext::Static(_)) => body.steal(),
-        Some(hir::ConstContext::ConstFn) => body.borrow().clone(),
+        Some(hir::ConstContext::ConstFn) => Box::new(body.borrow().clone()),
         None => bug!("`mir_for_ctfe` called on non-const {def:?}"),
     };
 
-    let mut body = remap_mir_for_const_eval_select(tcx, body, hir::Constness::Const);
+    remap_mir_for_const_eval_select(tcx, &mut *body, hir::Constness::Const);
     pm::run_passes(tcx, &mut body, &[&ctfe_limit::CtfeLimit], None, pm::Optimizations::Allowed);
 
     body
@@ -779,10 +776,10 @@ fn optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> &Body<'_> {
         return tcx.mir_for_ctfe(did);
     }
 
-    tcx.arena.alloc(inner_optimized_mir(tcx, did))
+    &**tcx.arena.alloc(inner_optimized_mir(tcx, did))
 }
 
-fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> {
+fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Box<Body<'_>> {
     match tcx.hir_body_const_context(did) {
         // Run the `mir_for_ctfe` query, which depends on `mir_drops_elaborated_and_const_checked`
         // which we are going to steal below. Thus we need to run `mir_for_ctfe` first, so it
@@ -792,8 +789,8 @@ fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> {
         Some(other) => panic!("do not use `optimized_mir` for constants: {other:?}"),
     }
     debug!("about to call mir_drops_elaborated...");
-    let body = tcx.mir_drops_elaborated_and_const_checked(did).steal();
-    let mut body = remap_mir_for_const_eval_select(tcx, body, hir::Constness::NotConst);
+    let mut body = tcx.mir_drops_elaborated_and_const_checked(did).steal();
+    remap_mir_for_const_eval_select(tcx, &mut *body, hir::Constness::NotConst);
 
     if body.tainted_by_errors.is_some() {
         return body;
@@ -830,9 +827,9 @@ fn promoted_mir(tcx: TyCtxt<'_>, def: LocalDefId) -> &IndexVec<Promoted, Body<'_
     }
     let mut promoted = tcx.mir_promoted(def).1.steal();
 
-    for body in &mut promoted {
+    for body in promoted.iter_mut() {
         run_analysis_to_runtime_passes(tcx, body);
     }
 
-    tcx.arena.alloc(promoted)
+    tcx.arena.alloc(*promoted)
 }