Support variantful generators

This allows generators to overlap fields using variants.

2 files changed, 81 insertions, 28 deletions

diff --git a/src/librustc/ty/layout.rs b/src/librustc/ty/layout.rs
index fd1d3a91ede..10afaebc91a 100644
--- a/src/librustc/ty/layout.rs
+++ b/src/librustc/ty/layout.rs
@@ -604,12 +604,57 @@ 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<_>, _>>()?,
+                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) => {
@@ -1646,6 +1691,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 handle_discriminant = |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 |
@@ -1720,7 +1773,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 handle_discriminant(discr);
+                        }
+                        substs.prefix_tys(def_id, tcx).nth(i).unwrap()
+                    }
+                }
             }
 
             ty::Tuple(tys) => tys[i],
@@ -1740,11 +1805,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 handle_discriminant(discr);
                     }
                 }
             }
diff --git a/src/librustc/ty/sty.rs b/src/librustc/ty/sty.rs
index edd6014618e..842e49dfc08 100644
--- a/src/librustc/ty/sty.rs
+++ b/src/librustc/ty/sty.rs
@@ -15,7 +15,6 @@ use crate::util::captures::Captures;
 use crate::mir::interpret::{Scalar, Pointer};
 
 use smallvec::SmallVec;
-use std::iter;
 use std::cmp::Ordering;
 use std::marker::PhantomData;
 use rustc_target::spec::abi;
@@ -475,30 +474,23 @@ impl<'a, 'gcx, 'tcx> GeneratorSubsts<'tcx> {
     /// This returns the types of the MIR locals which had to be stored across suspension points.
     /// It is calculated in rustc_mir::transform::generator::StateTransform.
     /// All the types here must be in the tuple in GeneratorInterior.
+    ///
+    /// The locals are grouped by their variant number. Note that some locals may
+    /// be repeated in multiple variants.
     pub fn state_tys(self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) ->
-        impl Iterator<Item=Ty<'tcx>> + Captures<'gcx> + 'a
+        impl Iterator<Item=impl Iterator<Item=Ty<'tcx>> + Captures<'gcx> + 'a>
     {
-        // TODO remove so we can handle variants properly
         tcx.generator_layout(def_id)
-            .variant_fields[0].iter()
-            .map(move |d| d.ty.subst(tcx, self.substs))
+            .variant_fields.iter()
+            .map(move |v| v.iter().map(move |d| d.ty.subst(tcx, self.substs)))
     }
 
-    /// This is the types of the fields of a generator which
-    /// is available before the generator transformation.
-    /// It includes the upvars and the state discriminant.
-    pub fn pre_transforms_tys(self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) ->
+    /// This is the types of the fields of a generator which are not stored in a
+    /// variant.
+    pub fn prefix_tys(self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) ->
         impl Iterator<Item=Ty<'tcx>> + 'a
     {
-        self.upvar_tys(def_id, tcx).chain(iter::once(self.discr_ty(tcx)))
-    }
-
-    /// This is the types of all the fields stored in a generator.
-    /// It includes the upvars, state types and the state discriminant.
-    pub fn field_tys(self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) ->
-        impl Iterator<Item=Ty<'tcx>> + Captures<'gcx> + 'a
-    {
-        self.pre_transforms_tys(def_id, tcx).chain(self.state_tys(def_id, tcx))
+        self.upvar_tys(def_id, tcx)
     }
 }