Box generator-related Body fields

14 files changed, 89 insertions, 54 deletions

diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs
index 85d1b702399..b2e0a4311a7 100644
--- a/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs
+++ b/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs
@@ -1417,7 +1417,7 @@ fn generator_layout_and_saved_local_names(
     def_id: DefId,
 ) -> (&'tcx GeneratorLayout<'tcx>, IndexVec<mir::GeneratorSavedLocal, Option<Symbol>>) {
     let body = tcx.optimized_mir(def_id);
-    let generator_layout = body.generator_layout.as_ref().unwrap();
+    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);
diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index f0b52c69881..4d8615a215f 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/compiler/rustc_middle/src/mir/mod.rs
@@ -146,6 +146,22 @@ impl<'tcx> MirSource<'tcx> {
     }
 }
 
+#[derive(Clone, TyEncodable, TyDecodable, Debug, HashStable, TypeFoldable)]
+pub struct GeneratorInfo<'tcx> {
+    /// The yield type of the function, if it is a generator.
+    pub yield_ty: Option<Ty<'tcx>>,
+
+    /// Generator drop glue.
+    pub generator_drop: Option<Body<'tcx>>,
+
+    /// The layout of a generator. Produced by the state transformation.
+    pub generator_layout: Option<GeneratorLayout<'tcx>>,
+
+    /// If this is a generator then record the type of source expression that caused this generator
+    /// to be created.
+    pub generator_kind: GeneratorKind,
+}
+
 /// The lowered representation of a single function.
 #[derive(Clone, TyEncodable, TyDecodable, Debug, HashStable, TypeFoldable)]
 pub struct Body<'tcx> {
@@ -166,18 +182,7 @@ pub struct Body<'tcx> {
     /// and used for debuginfo. Indexed by a `SourceScope`.
     pub source_scopes: IndexVec<SourceScope, SourceScopeData<'tcx>>,
 
-    /// The yield type of the function, if it is a generator.
-    pub yield_ty: Option<Ty<'tcx>>,
-
-    /// Generator drop glue.
-    pub generator_drop: Option<Box<Body<'tcx>>>,
-
-    /// The layout of a generator. Produced by the state transformation.
-    pub generator_layout: Option<GeneratorLayout<'tcx>>,
-
-    /// If this is a generator then record the type of source expression that caused this generator
-    /// to be created.
-    pub generator_kind: Option<GeneratorKind>,
+    pub generator: Option<Box<GeneratorInfo<'tcx>>>,
 
     /// Declarations of locals.
     ///
@@ -259,10 +264,14 @@ impl<'tcx> Body<'tcx> {
             source,
             basic_blocks,
             source_scopes,
-            yield_ty: None,
-            generator_drop: None,
-            generator_layout: None,
-            generator_kind,
+            generator: generator_kind.map(|generator_kind| {
+                Box::new(GeneratorInfo {
+                    yield_ty: None,
+                    generator_drop: None,
+                    generator_layout: None,
+                    generator_kind,
+                })
+            }),
             local_decls,
             user_type_annotations,
             arg_count,
@@ -289,16 +298,13 @@ impl<'tcx> Body<'tcx> {
             source: MirSource::item(DefId::local(CRATE_DEF_INDEX)),
             basic_blocks,
             source_scopes: IndexVec::new(),
-            yield_ty: None,
-            generator_drop: None,
-            generator_layout: None,
+            generator: None,
             local_decls: IndexVec::new(),
             user_type_annotations: IndexVec::new(),
             arg_count: 0,
             spread_arg: None,
             span: DUMMY_SP,
             required_consts: Vec::new(),
-            generator_kind: None,
             var_debug_info: Vec::new(),
             is_polymorphic: false,
             predecessor_cache: PredecessorCache::new(),
@@ -480,6 +486,26 @@ impl<'tcx> Body<'tcx> {
     pub fn dominators(&self) -> Dominators<BasicBlock> {
         dominators(self)
     }
+
+    #[inline]
+    pub fn yield_ty(&self) -> Option<Ty<'tcx>> {
+        self.generator.as_ref().and_then(|generator| generator.yield_ty)
+    }
+
+    #[inline]
+    pub fn generator_layout(&self) -> Option<&GeneratorLayout<'tcx>> {
+        self.generator.as_ref().and_then(|generator| generator.generator_layout.as_ref())
+    }
+
+    #[inline]
+    pub fn generator_drop(&self) -> Option<&Body<'tcx>> {
+        self.generator.as_ref().and_then(|generator| generator.generator_drop.as_ref())
+    }
+
+    #[inline]
+    pub fn generator_kind(&self) -> Option<GeneratorKind> {
+        self.generator.as_ref().map(|generator| generator.generator_kind)
+    }
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, Debug, TyEncodable, TyDecodable, HashStable)]
diff --git a/compiler/rustc_middle/src/mir/visit.rs b/compiler/rustc_middle/src/mir/visit.rs
index bb8c5f175b8..3e92844c4f9 100644
--- a/compiler/rustc_middle/src/mir/visit.rs
+++ b/compiler/rustc_middle/src/mir/visit.rs
@@ -241,11 +241,13 @@ macro_rules! make_mir_visitor {
                 body: &$($mutability)? Body<'tcx>,
             ) {
                 let span = body.span;
-                if let Some(yield_ty) = &$($mutability)? body.yield_ty {
-                    self.visit_ty(
-                        yield_ty,
-                        TyContext::YieldTy(SourceInfo::outermost(span))
-                    );
+                if let Some(gen) = &$($mutability)? body.generator {
+                    if let Some(yield_ty) = &$($mutability)? gen.yield_ty {
+                        self.visit_ty(
+                            yield_ty,
+                            TyContext::YieldTy(SourceInfo::outermost(span))
+                        );
+                    }
                 }
 
                 // for best performance, we want to use an iterator rather
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
index 154a5e54a53..2e077827873 100644
--- a/compiler/rustc_middle/src/ty/mod.rs
+++ b/compiler/rustc_middle/src/ty/mod.rs
@@ -2998,7 +2998,7 @@ impl<'tcx> TyCtxt<'tcx> {
     /// Returns layout of a generator. Layout might be unavailable if the
     /// generator is tainted by errors.
     pub fn generator_layout(self, def_id: DefId) -> Option<&'tcx GeneratorLayout<'tcx>> {
-        self.optimized_mir(def_id).generator_layout.as_ref()
+        self.optimized_mir(def_id).generator_layout()
     }
 
     /// Given the `DefId` of an impl, returns the `DefId` of the trait it implements.
diff --git a/compiler/rustc_mir/src/borrow_check/type_check/input_output.rs b/compiler/rustc_mir/src/borrow_check/type_check/input_output.rs
index 157959b1159..77d91366224 100644
--- a/compiler/rustc_mir/src/borrow_check/type_check/input_output.rs
+++ b/compiler/rustc_mir/src/borrow_check/type_check/input_output.rs
@@ -103,11 +103,8 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             }
         }
 
-        assert!(
-            body.yield_ty.is_some() && universal_regions.yield_ty.is_some()
-                || body.yield_ty.is_none() && universal_regions.yield_ty.is_none()
-        );
-        if let Some(mir_yield_ty) = body.yield_ty {
+        assert!(body.yield_ty().is_some() == universal_regions.yield_ty.is_some());
+        if let Some(mir_yield_ty) = body.yield_ty() {
             let ur_yield_ty = universal_regions.yield_ty.unwrap();
             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);
diff --git a/compiler/rustc_mir/src/borrow_check/type_check/mod.rs b/compiler/rustc_mir/src/borrow_check/type_check/mod.rs
index aa946fdafa9..fa5b342c6e2 100644
--- a/compiler/rustc_mir/src/borrow_check/type_check/mod.rs
+++ b/compiler/rustc_mir/src/borrow_check/type_check/mod.rs
@@ -1651,7 +1651,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             }
             TerminatorKind::Yield { ref value, .. } => {
                 let value_ty = value.ty(body, tcx);
-                match body.yield_ty {
+                match body.yield_ty() {
                     None => span_mirbug!(self, term, "yield in non-generator"),
                     Some(ty) => {
                         if let Err(terr) = self.sub_types(
diff --git a/compiler/rustc_mir/src/shim.rs b/compiler/rustc_mir/src/shim.rs
index 6aaf27bdcb5..b246ec5c814 100644
--- a/compiler/rustc_mir/src/shim.rs
+++ b/compiler/rustc_mir/src/shim.rs
@@ -134,7 +134,7 @@ fn build_drop_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, ty: Option<Ty<'tcx>>)
 
     // Check if this is a generator, if so, return the drop glue for it
     if let Some(&ty::Generator(gen_def_id, substs, _)) = ty.map(|ty| ty.kind()) {
-        let body = &**tcx.optimized_mir(gen_def_id).generator_drop.as_ref().unwrap();
+        let body = tcx.optimized_mir(gen_def_id).generator_drop().unwrap();
         return body.clone().subst(tcx, substs);
     }
 
diff --git a/compiler/rustc_mir/src/transform/check_consts/validation.rs b/compiler/rustc_mir/src/transform/check_consts/validation.rs
index a8263683712..a8569eb4f48 100644
--- a/compiler/rustc_mir/src/transform/check_consts/validation.rs
+++ b/compiler/rustc_mir/src/transform/check_consts/validation.rs
@@ -222,7 +222,7 @@ impl Validator<'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 is_async_fn(self.ccx) || body.generator_kind.is_some() {
+        if is_async_fn(self.ccx) || body.generator.is_some() {
             tcx.sess.delay_span_bug(body.span, "`async` functions cannot be `const fn`");
             return;
         }
diff --git a/compiler/rustc_mir/src/transform/const_prop.rs b/compiler/rustc_mir/src/transform/const_prop.rs
index 3b27b544310..2a0e5e2c9c5 100644
--- a/compiler/rustc_mir/src/transform/const_prop.rs
+++ b/compiler/rustc_mir/src/transform/const_prop.rs
@@ -140,7 +140,7 @@ impl<'tcx> MirPass<'tcx> for ConstProp {
             body.arg_count,
             Default::default(),
             body.span,
-            body.generator_kind,
+            body.generator_kind(),
         );
 
         // FIXME(oli-obk, eddyb) Optimize locals (or even local paths) to hold
diff --git a/compiler/rustc_mir/src/transform/generator.rs b/compiler/rustc_mir/src/transform/generator.rs
index 9f996e84528..7a1f3d44a5e 100644
--- a/compiler/rustc_mir/src/transform/generator.rs
+++ b/compiler/rustc_mir/src/transform/generator.rs
@@ -1111,7 +1111,7 @@ fn create_generator_resume_function<'tcx>(
     cases.insert(0, (UNRESUMED, BasicBlock::new(0)));
 
     // Panic when resumed on the returned or poisoned state
-    let generator_kind = body.generator_kind.unwrap();
+    let generator_kind = body.generator_kind().unwrap();
 
     if can_unwind {
         cases.insert(
@@ -1236,14 +1236,14 @@ fn create_cases<'tcx>(
 
 impl<'tcx> MirPass<'tcx> for StateTransform {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-        let yield_ty = if let Some(yield_ty) = body.yield_ty {
+        let yield_ty = if let Some(yield_ty) = body.yield_ty() {
             yield_ty
         } else {
             // This only applies to generators
             return;
         };
 
-        assert!(body.generator_drop.is_none());
+        assert!(body.generator_drop().is_none());
 
         // The first argument is the generator type passed by value
         let gen_ty = body.local_decls.raw[1].ty;
@@ -1340,10 +1340,11 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
         transform.visit_body(body);
 
         // Update our MIR struct to reflect the changes we've made
-        body.yield_ty = None;
         body.arg_count = 2; // self, resume arg
         body.spread_arg = None;
-        body.generator_layout = Some(layout);
+
+        body.generator.as_mut().unwrap().yield_ty = None;
+        body.generator.as_mut().unwrap().generator_layout = Some(layout);
 
         // Insert `drop(generator_struct)` which is used to drop upvars for generators in
         // the unresumed state.
@@ -1362,7 +1363,7 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
         // Create a copy of our MIR and use it to create the drop shim for the generator
         let drop_shim = create_generator_drop_shim(tcx, &transform, gen_ty, body, drop_clean);
 
-        body.generator_drop = Some(box drop_shim);
+        body.generator.as_mut().unwrap().generator_drop = Some(drop_shim);
 
         // Create the Generator::resume function
         create_generator_resume_function(tcx, transform, body, can_return);
diff --git a/compiler/rustc_mir/src/transform/inline.rs b/compiler/rustc_mir/src/transform/inline.rs
index 16279040a1c..518f7b41ed8 100644
--- a/compiler/rustc_mir/src/transform/inline.rs
+++ b/compiler/rustc_mir/src/transform/inline.rs
@@ -214,7 +214,7 @@ impl Inliner<'tcx> {
             // since their `optimized_mir` is used for layout computation, which can
             // create a cycle, even when no attempt is made to inline the function
             // in the other direction.
-            if caller_body.generator_kind.is_some() {
+            if caller_body.generator.is_some() {
                 return Err("local generator (query cycle avoidance)");
             }
 
diff --git a/compiler/rustc_mir/src/transform/promote_consts.rs b/compiler/rustc_mir/src/transform/promote_consts.rs
index 1d4438d80c9..47d976281a4 100644
--- a/compiler/rustc_mir/src/transform/promote_consts.rs
+++ b/compiler/rustc_mir/src/transform/promote_consts.rs
@@ -1177,7 +1177,7 @@ pub fn promote_candidates<'tcx>(
             0,
             vec![],
             body.span,
-            body.generator_kind,
+            body.generator_kind(),
         );
 
         let promoter = Promoter {
diff --git a/compiler/rustc_mir/src/util/pretty.rs b/compiler/rustc_mir/src/util/pretty.rs
index f46e94eea96..247a0beccaf 100644
--- a/compiler/rustc_mir/src/util/pretty.rs
+++ b/compiler/rustc_mir/src/util/pretty.rs
@@ -131,7 +131,7 @@ fn dump_matched_mir_node<'tcx, F>(
             Some(promoted) => write!(file, "::{:?}`", promoted)?,
         }
         writeln!(file, " {} {}", disambiguator, pass_name)?;
-        if let Some(ref layout) = body.generator_layout {
+        if let Some(ref layout) = body.generator_layout() {
             writeln!(file, "/* generator_layout = {:#?} */", layout)?;
         }
         writeln!(file)?;
@@ -956,7 +956,7 @@ fn write_mir_sig(tcx: TyCtxt<'_>, body: &Body<'_>, w: &mut dyn Write) -> io::Res
         write!(w, ": {} =", body.return_ty())?;
     }
 
-    if let Some(yield_ty) = body.yield_ty {
+    if let Some(yield_ty) = body.yield_ty() {
         writeln!(w)?;
         writeln!(w, "yields {}", yield_ty)?;
     }
diff --git a/compiler/rustc_mir_build/src/build/mod.rs b/compiler/rustc_mir_build/src/build/mod.rs
index 5f6c8d26402..b928458df8e 100644
--- a/compiler/rustc_mir_build/src/build/mod.rs
+++ b/compiler/rustc_mir_build/src/build/mod.rs
@@ -188,7 +188,9 @@ fn mir_build(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> Body<'_
                 body,
                 span_with_body,
             );
-            mir.yield_ty = yield_ty;
+            if yield_ty.is_some() {
+                mir.generator.as_mut().unwrap().yield_ty = yield_ty;
+            }
             mir
         } else {
             // Get the revealed type of this const. This is *not* the adjusted
@@ -218,7 +220,7 @@ fn mir_build(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> Body<'_
             !(body.local_decls.has_free_regions()
                 || body.basic_blocks().has_free_regions()
                 || body.var_debug_info.has_free_regions()
-                || body.yield_ty.has_free_regions()),
+                || body.yield_ty().has_free_regions()),
             "Unexpected free regions in MIR: {:?}",
             body,
         );
@@ -686,10 +688,11 @@ fn construct_error<'a, 'tcx>(hir: Cx<'a, 'tcx>, body_id: hir::BodyId) -> Body<'t
     let def_id = tcx.hir().local_def_id(owner_id);
     let span = tcx.hir().span(owner_id);
     let ty = tcx.ty_error();
+    let generator_kind = tcx.hir().body(body_id).generator_kind;
     let num_params = match hir.body_owner_kind {
         hir::BodyOwnerKind::Fn => tcx.hir().fn_decl_by_hir_id(owner_id).unwrap().inputs.len(),
         hir::BodyOwnerKind::Closure => {
-            if tcx.hir().body(body_id).generator_kind().is_some() {
+            if generator_kind.is_some() {
                 // Generators have an implicit `self` parameter *and* a possibly
                 // implicit resume parameter.
                 2
@@ -701,8 +704,16 @@ fn construct_error<'a, 'tcx>(hir: Cx<'a, 'tcx>, body_id: hir::BodyId) -> Body<'t
         hir::BodyOwnerKind::Const => 0,
         hir::BodyOwnerKind::Static(_) => 0,
     };
-    let mut builder =
-        Builder::new(hir, def_id.to_def_id(), span, num_params, Safety::Safe, ty, span, None);
+    let mut builder = Builder::new(
+        hir,
+        def_id.to_def_id(),
+        span,
+        num_params,
+        Safety::Safe,
+        ty,
+        span,
+        generator_kind,
+    );
     let source_info = builder.source_info(span);
     // Some MIR passes will expect the number of parameters to match the
     // function declaration.
@@ -711,9 +722,7 @@ fn construct_error<'a, 'tcx>(hir: Cx<'a, 'tcx>, body_id: hir::BodyId) -> Body<'t
     }
     builder.cfg.terminate(START_BLOCK, source_info, TerminatorKind::Unreachable);
     let mut body = builder.finish();
-    if tcx.hir().body(body_id).generator_kind.is_some() {
-        body.yield_ty = Some(ty);
-    }
+    body.generator.as_mut().map(|gen| gen.yield_ty = Some(ty));
     body
 }