Separate generator info from MIR body.

Author: cjgillot

compiler/rustc_borrowck/src/type_check/input_output.rs

@@ -115,28 +115,26 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             }
         }
 
+        let body_yield_ty = body.yield_ty(self.tcx());
         debug!(
             "equate_inputs_and_outputs: body.yield_ty {:?}, universal_regions.yield_ty {:?}",
-            body.yield_ty(),
-            universal_regions.yield_ty
+            body_yield_ty, universal_regions.yield_ty
         );
 
         // We will not have a universal_regions.yield_ty if we yield (by accident)
         // outside of a generator and return an `impl Trait`, so emit a delay_span_bug
         // because we don't want to panic in an assert here if we've already got errors.
-        if body.yield_ty().is_some() != universal_regions.yield_ty.is_some() {
+        if body_yield_ty.is_some() != universal_regions.yield_ty.is_some() {
             self.tcx().sess.delay_span_bug(
                 body.span,
                 &format!(
                     "Expected body to have yield_ty ({:?}) iff we have a UR yield_ty ({:?})",
-                    body.yield_ty(),
-                    universal_regions.yield_ty,
+                    body_yield_ty, universal_regions.yield_ty,
                 ),
             );
         }
 
-        if let (Some(mir_yield_ty), Some(ur_yield_ty)) =
-            (body.yield_ty(), universal_regions.yield_ty)
+        if let (Some(mir_yield_ty), Some(ur_yield_ty)) = (body_yield_ty, universal_regions.yield_ty)
         {
             let yield_span = body.local_decls[RETURN_PLACE].source_info.span;
             self.equate_normalized_input_or_output(ur_yield_ty, mir_yield_ty, yield_span);

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -1473,7 +1473,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 self.check_operand(value, term_location);
 
                 let value_ty = value.ty(body, tcx);
-                match body.yield_ty() {
+                match body.yield_ty(tcx) {
                     None => span_mirbug!(self, term, "yield in non-generator"),
                     Some(ty) => {
                         if let Err(terr) = self.sub_types(

compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs

@@ -1235,8 +1235,8 @@ fn generator_layout_and_saved_local_names<'tcx>(
     tcx: TyCtxt<'tcx>,
     def_id: DefId,
 ) -> (&'tcx GeneratorLayout<'tcx>, IndexVec<mir::GeneratorSavedLocal, Option<Symbol>>) {
+    let generator_layout = &tcx.mir_generator_info(def_id).generator_layout;
     let body = tcx.optimized_mir(def_id);
-    let generator_layout = body.generator_layout().unwrap();
     let mut generator_saved_local_names = IndexVec::from_elem(None, &generator_layout.field_tys);
 
     let state_arg = mir::Local::new(1);

compiler/rustc_const_eval/src/transform/check_consts/check.rs

@@ -222,7 +222,7 @@ impl<'mir, 'tcx> Checker<'mir, 'tcx> {
 
         // `async` functions cannot be `const fn`. This is checked during AST lowering, so there's
         // no need to emit duplicate errors here.
-        if self.ccx.is_async() || body.generator.is_some() {
+        if self.ccx.is_async() || tcx.generator_kind(def_id).is_some() {
             tcx.sess.delay_span_bug(body.span, "`async` functions cannot be `const fn`");
             return;
         }

compiler/rustc_const_eval/src/transform/promote_consts.rs

@@ -973,7 +973,6 @@ pub fn promote_candidates<'tcx>(
             0,
             vec![],
             body.span,
-            body.generator_kind(),
             body.tainted_by_errors,
         );
         promoted.phase = MirPhase::Analysis(AnalysisPhase::Initial);

compiler/rustc_const_eval/src/transform/validate.rs

@@ -65,6 +65,7 @@ impl<'tcx> MirPass<'tcx> for Validator {
             storage_liveness,
             place_cache: Vec::new(),
             value_cache: Vec::new(),
+            is_generator: tcx.generator_kind(def_id).is_some(),
         }
         .visit_body(body);
     }
@@ -117,6 +118,7 @@ struct TypeChecker<'a, 'tcx> {
     storage_liveness: ResultsCursor<'a, 'tcx, MaybeStorageLive>,
     place_cache: Vec<PlaceRef<'tcx>>,
     value_cache: Vec<u128>,
+    is_generator: bool,
 }
 
 impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
@@ -323,16 +325,8 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                     }
                     &ty::Generator(def_id, substs, _) => {
                         let f_ty = if let Some(var) = parent_ty.variant_index {
-                            let gen_body = if def_id == self.body.source.def_id() {
-                                self.body
-                            } else {
-                                self.tcx.optimized_mir(def_id)
-                            };
-
-                            let Some(layout) = gen_body.generator_layout() else {
-                                self.fail(location, format!("No generator layout for {:?}", parent_ty));
-                                return;
-                            };
+                            let generator_info = self.tcx.mir_generator_info(def_id);
+                            let layout = &generator_info.generator_layout;
 
                             let Some(&local) = layout.variant_fields[var].get(f) else {
                                 fail_out_of_bounds(self, location);
@@ -836,10 +830,10 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                 }
             }
             TerminatorKind::Yield { resume, drop, .. } => {
-                if self.body.generator.is_none() {
+                if !self.is_generator {
                     self.fail(location, "`Yield` cannot appear outside generator bodies");
                 }
-                if self.mir_phase >= MirPhase::Runtime(RuntimePhase::Initial) {
+                if self.mir_phase >= MirPhase::Runtime(RuntimePhase::GeneratorsLowered) {
                     self.fail(location, "`Yield` should have been replaced by generator lowering");
                 }
                 self.check_edge(location, *resume, EdgeKind::Normal);
@@ -878,10 +872,10 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                 }
             }
             TerminatorKind::GeneratorDrop => {
-                if self.body.generator.is_none() {
+                if !self.is_generator {
                     self.fail(location, "`GeneratorDrop` cannot appear outside generator bodies");
                 }
-                if self.mir_phase >= MirPhase::Runtime(RuntimePhase::Initial) {
+                if self.mir_phase >= MirPhase::Runtime(RuntimePhase::GeneratorsLowered) {
                     self.fail(
                         location,
                         "`GeneratorDrop` should have been replaced by generator lowering",

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -203,6 +203,7 @@ provide! { tcx, def_id, other, cdata,
     thir_abstract_const => { table }
     optimized_mir => { table }
     mir_for_ctfe => { table }
+    mir_generator_info => { table }
     promoted_mir => { table }
     def_span => { table }
     def_ident_span => { table }

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -1387,6 +1387,10 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             debug!("EntryBuilder::encode_mir({:?})", def_id);
             if encode_opt {
                 record!(self.tables.optimized_mir[def_id.to_def_id()] <- tcx.optimized_mir(def_id));
+
+                if let DefKind::Generator = self.tcx.def_kind(def_id) {
+                    record!(self.tables.mir_generator_info[def_id.to_def_id()] <- tcx.mir_generator_info(def_id));
+                }
             }
             if encode_const {
                 record!(self.tables.mir_for_ctfe[def_id.to_def_id()] <- tcx.mir_for_ctfe(def_id));

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -361,6 +361,7 @@ define_tables! {
     object_lifetime_default: Table<DefIndex, LazyValue<ObjectLifetimeDefault>>,
     optimized_mir: Table<DefIndex, LazyValue<mir::Body<'static>>>,
     mir_for_ctfe: Table<DefIndex, LazyValue<mir::Body<'static>>>,
+    mir_generator_info: Table<DefIndex, LazyValue<mir::GeneratorInfo<'static>>>,
     promoted_mir: Table<DefIndex, LazyValue<IndexVec<mir::Promoted, mir::Body<'static>>>>,
     // FIXME(compiler-errors): Why isn't this a LazyArray?
     thir_abstract_const: Table<DefIndex, LazyValue<&'static [ty::abstract_const::Node<'static>]>>,

compiler/rustc_middle/src/arena.rs

@@ -13,6 +13,11 @@ 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>,
+            [decode] generator_info: rustc_middle::mir::GeneratorInfo<'tcx>,
+            [] mir_generator_lowered: (
+                rustc_data_structures::steal::Steal<rustc_middle::mir::Body<'tcx>>,
+                Option<rustc_middle::mir::GeneratorInfo<'tcx>>,
+            ),
             [] steal_promoted:
                 rustc_data_structures::steal::Steal<
                     rustc_index::vec::IndexVec<