Auto merge of #59897 - tmandry:variantful-generators, r=eddyb

Multi-variant layouts for generators

This allows generators to overlap fields using variants, but doesn't do any such overlapping yet. It creates one variant for every state of the generator (unresumed, returned, panicked, plus one for every yield), and puts every stored local in each of the yield-point variants.

Required for optimizing generator layouts (#52924).

There was quite a lot of refactoring needed for this change. I've done my best in later commits to eliminate assumptions in the code that only certain kinds of types are multi-variant, and to centralize knowledge of the inner mechanics of generators in as few places as possible.

This change also emits debuginfo about the fields contained in each variant, as well as preserving debuginfo about stored locals while running in the generator.

Also, fixes #59972.

Future work:
- Use this change for an optimization pass that actually overlaps locals within the generator struct (#52924)
- In the type layout fields, don't include locals that are uninitialized for a particular variant, so miri and UB sanitizers can check our memory (see #59972 (comment))
- Preserve debuginfo scopes across generator yield points

src/librustc/mir/mod.rs

@@ -2030,6 +2030,10 @@ impl<'tcx> Place<'tcx> {
             variant_index))
     }
 
+    pub fn downcast_unnamed(self, variant_index: VariantIdx) -> Place<'tcx> {
+        self.elem(ProjectionElem::Downcast(None, variant_index))
+    }
+
     pub fn index(self, index: Local) -> Place<'tcx> {
         self.elem(ProjectionElem::Index(index))
     }
@@ -2589,11 +2593,6 @@ impl<'tcx> Debug for Rvalue<'tcx> {
                                     let var_name = tcx.hir().name_by_hir_id(freevar.var_id());
                                     struct_fmt.field(&var_name.as_str(), place);
                                 }
-                                struct_fmt.field("$state", &places[freevars.len()]);
-                                for i in (freevars.len() + 1)..places.len() {
-                                    struct_fmt
-                                        .field(&format!("${}", i - freevars.len() - 1), &places[i]);
-                                }
                             });
 
                             struct_fmt.finish()
@@ -3031,10 +3030,29 @@ pub struct UnsafetyCheckResult {
     pub unsafe_blocks: Lrc<[(hir::HirId, bool)]>,
 }
 
+newtype_index! {
+    pub struct GeneratorSavedLocal {
+        derive [HashStable]
+        DEBUG_FORMAT = "_{}",
+    }
+}
+
 /// The layout of generator state
 #[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable)]
 pub struct GeneratorLayout<'tcx> {
-    pub fields: Vec<LocalDecl<'tcx>>,
+    /// The type of every local stored inside the generator.
+    pub field_tys: IndexVec<GeneratorSavedLocal, Ty<'tcx>>,
+
+    /// Which of the above fields are in each variant. Note that one field may
+    /// be stored in multiple variants.
+    pub variant_fields: IndexVec<VariantIdx, IndexVec<Field, GeneratorSavedLocal>>,
+
+    /// Names and scopes of all the stored generator locals.
+    /// NOTE(tmandry) This is *strictly* a temporary hack for codegen
+    /// debuginfo generation, and will be removed at some point.
+    /// Do **NOT** use it for anything else, local information should not be
+    /// in the MIR, please rely on local crate HIR or other side-channels.
+    pub __local_debuginfo_codegen_only_do_not_use: IndexVec<GeneratorSavedLocal, LocalDecl<'tcx>>,
 }
 
 #[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable)]
@@ -3223,7 +3241,9 @@ BraceStructTypeFoldableImpl! {
 
 BraceStructTypeFoldableImpl! {
     impl<'tcx> TypeFoldable<'tcx> for GeneratorLayout<'tcx> {
-        fields
+        field_tys,
+        variant_fields,
+        __local_debuginfo_codegen_only_do_not_use,
     }
 }
 
@@ -3598,6 +3618,15 @@ impl<'tcx> TypeFoldable<'tcx> for Field {
     }
 }
 
+impl<'tcx> TypeFoldable<'tcx> for GeneratorSavedLocal {
+    fn super_fold_with<'gcx: 'tcx, F: TypeFolder<'gcx, 'tcx>>(&self, _: &mut F) -> Self {
+        *self
+    }
+    fn super_visit_with<V: TypeVisitor<'tcx>>(&self, _: &mut V) -> bool {
+        false
+    }
+}
+
 impl<'tcx> TypeFoldable<'tcx> for Constant<'tcx> {
     fn super_fold_with<'gcx: 'tcx, F: TypeFolder<'gcx, 'tcx>>(&self, folder: &mut F) -> Self {
         Constant {

src/librustc/mir/tcx.rs

@@ -177,11 +177,13 @@ impl<'tcx> Rvalue<'tcx> {
             }
             Rvalue::Discriminant(ref place) => {
                 let ty = place.ty(local_decls, tcx).ty;
-                if let ty::Adt(adt_def, _) = ty.sty {
-                    adt_def.repr.discr_type().to_ty(tcx)
-                } else {
-                    // This can only be `0`, for now, so `u8` will suffice.
-                    tcx.types.u8
+                match ty.sty {
+                    ty::Adt(adt_def, _) => adt_def.repr.discr_type().to_ty(tcx),
+                    ty::Generator(_, substs, _) => substs.discr_ty(tcx),
+                    _ => {
+                        // This can only be `0`, for now, so `u8` will suffice.
+                        tcx.types.u8
+                    }
                 }
             }
             Rvalue::NullaryOp(NullOp::Box, t) => tcx.mk_box(t),

src/librustc/ty/layout.rs

@@ -604,12 +604,63 @@ impl<'a, 'tcx> LayoutCx<'tcx, TyCtxt<'a, 'tcx, 'tcx>> {
                 tcx.intern_layout(unit)
             }
 
-            // Tuples, generators and closures.
             ty::Generator(def_id, ref substs, _) => {
-                let tys = substs.field_tys(def_id, tcx);
-                univariant(&tys.map(|ty| self.layout_of(ty)).collect::<Result<Vec<_>, _>>()?,
+                // FIXME(tmandry): For fields that are repeated in multiple
+                // variants in the GeneratorLayout, we need code to ensure that
+                // the offset of these fields never change. Right now this is
+                // not an issue since every variant has every field, but once we
+                // optimize this we have to be more careful.
+
+                let discr_index = substs.prefix_tys(def_id, tcx).count();
+                let prefix_tys = substs.prefix_tys(def_id, tcx)
+                    .chain(iter::once(substs.discr_ty(tcx)));
+                let prefix = univariant_uninterned(
+                    &prefix_tys.map(|ty| self.layout_of(ty)).collect::<Result<Vec<_>, _>>()?,
                     &ReprOptions::default(),
-                    StructKind::AlwaysSized)?
+                    StructKind::AlwaysSized)?;
+
+                let mut size = prefix.size;
+                let mut align = prefix.align;
+                let variants_tys = substs.state_tys(def_id, tcx);
+                let variants = variants_tys.enumerate().map(|(i, variant_tys)| {
+                    let mut variant = univariant_uninterned(
+                        &variant_tys.map(|ty| self.layout_of(ty)).collect::<Result<Vec<_>, _>>()?,
+                        &ReprOptions::default(),
+                        StructKind::Prefixed(prefix.size, prefix.align.abi))?;
+
+                    variant.variants = Variants::Single { index: VariantIdx::new(i) };
+
+                    size = size.max(variant.size);
+                    align = align.max(variant.align);
+
+                    Ok(variant)
+                }).collect::<Result<IndexVec<VariantIdx, _>, _>>()?;
+
+                let abi = if prefix.abi.is_uninhabited() ||
+                             variants.iter().all(|v| v.abi.is_uninhabited()) {
+                    Abi::Uninhabited
+                } else {
+                    Abi::Aggregate { sized: true }
+                };
+                let discr = match &self.layout_of(substs.discr_ty(tcx))?.abi {
+                    Abi::Scalar(s) => s.clone(),
+                    _ => bug!(),
+                };
+
+                let layout = tcx.intern_layout(LayoutDetails {
+                    variants: Variants::Multiple {
+                        discr,
+                        discr_kind: DiscriminantKind::Tag,
+                        discr_index,
+                        variants,
+                    },
+                    fields: prefix.fields,
+                    abi,
+                    size,
+                    align,
+                });
+                debug!("generator layout: {:#?}", layout);
+                layout
             }
 
             ty::Closure(def_id, ref substs) => {
@@ -1647,6 +1698,14 @@ impl<'a, 'tcx, C> TyLayoutMethods<'tcx, C> for Ty<'tcx>
 
     fn field(this: TyLayout<'tcx>, cx: &C, i: usize) -> C::TyLayout {
         let tcx = cx.tcx();
+        let discr_layout = |discr: &Scalar| -> C::TyLayout {
+            let layout = LayoutDetails::scalar(cx, discr.clone());
+            MaybeResult::from_ok(TyLayout {
+                details: tcx.intern_layout(layout),
+                ty: discr.value.to_ty(tcx)
+            })
+        };
+
         cx.layout_of(match this.ty.sty {
             ty::Bool |
             ty::Char |
@@ -1721,7 +1780,19 @@ impl<'a, 'tcx, C> TyLayoutMethods<'tcx, C> for Ty<'tcx>
             }
 
             ty::Generator(def_id, ref substs, _) => {
-                substs.field_tys(def_id, tcx).nth(i).unwrap()
+                match this.variants {
+                    Variants::Single { index } => {
+                        substs.state_tys(def_id, tcx)
+                            .nth(index.as_usize()).unwrap()
+                            .nth(i).unwrap()
+                    }
+                    Variants::Multiple { ref discr, discr_index, .. } => {
+                        if i == discr_index {
+                            return discr_layout(discr);
+                        }
+                        substs.prefix_tys(def_id, tcx).nth(i).unwrap()
+                    }
+                }
             }
 
             ty::Tuple(tys) => tys[i].expect_ty(),
@@ -1741,11 +1812,7 @@ impl<'a, 'tcx, C> TyLayoutMethods<'tcx, C> for Ty<'tcx>
                     // Discriminant field for enums (where applicable).
                     Variants::Multiple { ref discr, .. } => {
                         assert_eq!(i, 0);
-                        let layout = LayoutDetails::scalar(cx, discr.clone());
-                        return MaybeResult::from_ok(TyLayout {
-                            details: tcx.intern_layout(layout),
-                            ty: discr.value.to_ty(tcx)
-                        });
+                        return discr_layout(discr);
                     }
                 }
             }

src/librustc/ty/mod.rs

@@ -43,6 +43,7 @@ use std::ops::Deref;
 use rustc_data_structures::sync::{self, Lrc, ParallelIterator, par_iter};
 use std::slice;
 use std::{mem, ptr};
+use std::ops::Range;
 use syntax::ast::{self, Name, Ident, NodeId};
 use syntax::attr;
 use syntax::ext::hygiene::Mark;
@@ -2416,11 +2417,17 @@ impl<'a, 'gcx, 'tcx> AdtDef {
         })
     }
 
+    #[inline]
+    pub fn variant_range(&self) -> Range<VariantIdx> {
+        (VariantIdx::new(0)..VariantIdx::new(self.variants.len()))
+    }
+
     /// Computes the discriminant value used by a specific variant.
     /// Unlike `discriminants`, this is (amortized) constant-time,
     /// only doing at most one query for evaluating an explicit
     /// discriminant (the last one before the requested variant),
     /// assuming there are no constant-evaluation errors there.
+    #[inline]
     pub fn discriminant_for_variant(&self,
                                     tcx: TyCtxt<'a, 'gcx, 'tcx>,
                                     variant_index: VariantIdx)