optimized_mir

Author: lcnr

src/librustc_metadata/rmeta/decoder/cstore_impl.rs

@@ -111,7 +111,7 @@ provide! { <'tcx> tcx, def_id, other, cdata,
             bug!("coerce_unsized_info: `{:?}` is missing its info", def_id);
         })
     }
-    optimized_mir => { cdata.get_optimized_mir(tcx, def_id.index) }
+    optimized_mir => { tcx.arena.alloc(cdata.get_optimized_mir(tcx, def_id.index)) }
     promoted_mir => { cdata.get_promoted_mir(tcx, def_id.index) }
     mir_const_qualif => { cdata.mir_const_qualif(def_id.index) }
     fn_sig => { cdata.fn_sig(def_id.index, tcx) }

src/librustc_middle/arena.rs

@@ -14,6 +14,7 @@ macro_rules! arena_types {
             [] layouts: rustc_target::abi::Layout, rustc_target::abi::Layout;
             // AdtDef are interned and compared by address
             [] adt_def: rustc_middle::ty::AdtDef, rustc_middle::ty::AdtDef;
+            [] mir: rustc_middle::mir::Body<$tcx>, rustc_middle::mir::Body<$tcx>;
             [decode] tables: rustc_middle::ty::TypeckTables<$tcx>, rustc_middle::ty::TypeckTables<'_x>;
             [] const_allocs: rustc_middle::mir::interpret::Allocation, rustc_middle::mir::interpret::Allocation;
             // Required for the incremental on-disk cache

src/librustc_middle/query/mod.rs

@@ -244,11 +244,13 @@ rustc_queries! {
 
         /// MIR after our optimization passes have run. This is MIR that is ready
         /// for codegen. This is also the only query that can fetch non-local MIR, at present.
-        query optimized_mir(key: DefId) -> mir::Body<'tcx> {
+        query optimized_mir(key: DefId) -> &'tcx mir::Body<'tcx> {
             desc { |tcx| "optimizing MIR for `{}`", tcx.def_path_str(key) }
-            storage(ArenaCacheSelector<'tcx>)
             cache_on_disk_if { key.is_local() }
         }
+        query optimized_mir_of_const_arg(key: ty::WithOptParam<LocalDefId>) -> &'tcx mir::Body<'tcx> {
+            desc { |tcx| "optimizing MIR for `{}`", tcx.def_path_str(key.did.to_def_id()) }
+        }
 
         /// Returns coverage summary info for a function, after executing the `InstrumentCoverage`
         /// MIR pass (assuming the -Zinstrument-coverage option is enabled).

src/librustc_middle/ty/mod.rs

@@ -2880,7 +2880,13 @@ impl<'tcx> TyCtxt<'tcx> {
     /// Returns the possibly-auto-generated MIR of a `(DefId, Subst)` pair.
     pub fn instance_mir(self, instance: ty::InstanceDef<'tcx>) -> &'tcx Body<'tcx> {
         match instance {
-            ty::InstanceDef::Item(def) => self.optimized_mir(def.did),
+            ty::InstanceDef::Item(def) => {
+                if let Some(def) = def.as_local() {
+                    self.optimized_mir_of_const_arg(def)
+                } else {
+                    self.optimized_mir(def.did)
+                }
+            }
             ty::InstanceDef::VtableShim(..)
             | ty::InstanceDef::ReifyShim(..)
             | ty::InstanceDef::Intrinsic(..)

src/librustc_mir/transform/mod.rs

@@ -52,6 +52,7 @@ pub(crate) fn provide(providers: &mut Providers) {
         mir_validated,
         mir_drops_elaborated_and_const_checked,
         optimized_mir,
+        optimized_mir_of_const_arg,
         is_mir_available,
         promoted_mir,
         ..*providers
@@ -424,19 +425,41 @@ fn run_optimization_passes<'tcx>(
     );
 }
 
-fn optimized_mir(tcx: TyCtxt<'_>, def_id: DefId) -> Body<'_> {
-    if tcx.is_constructor(def_id) {
+fn optimized_mir<'tcx>(tcx: TyCtxt<'tcx>, did: DefId) -> &'tcx Body<'tcx> {
+    let did = did.expect_local();
+    if let param_did @ Some(_) = tcx.opt_const_param_of(did) {
+        tcx.optimized_mir_of_const_arg(ty::WithOptParam { did, param_did })
+    } else {
+        tcx.arena.alloc(inner_optimized_mir(tcx, ty::WithOptParam::dummy(did)))
+    }
+}
+
+fn optimized_mir_of_const_arg<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def: ty::WithOptParam<LocalDefId>,
+) -> &'tcx Body<'tcx> {
+    if def.param_did.is_none() {
+        if let param_did @ Some(_) = tcx.opt_const_param_of(def.did) {
+            tcx.optimized_mir_of_const_arg(ty::WithOptParam { param_did, ..def })
+        } else {
+            tcx.optimized_mir(def.did)
+        }
+    } else {
+        tcx.arena.alloc(inner_optimized_mir(tcx, def))
+    }
+}
+
+fn inner_optimized_mir(tcx: TyCtxt<'_>, def: ty::WithOptParam<LocalDefId>) -> Body<'_> {
+    if tcx.is_constructor(def.did.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_id);
+        return shim::build_adt_ctor(tcx, def.did.to_def_id());
     }
 
-    let def_id = def_id.expect_local();
-
-    let mut body = tcx.mir_drops_elaborated_and_const_checked(def_id).steal();
-    run_optimization_passes(tcx, &mut body, def_id, None);
+    let mut body = tcx.mir_drops_elaborated_and_const_checked(def.did).steal();
+    run_optimization_passes(tcx, &mut body, def.did, None);
 
     debug_assert!(!body.has_free_regions(), "Free regions in optimized MIR");