Support variantful generators

This allows generators to overlap fields using variants.

Author: tmandry

src/librustc/mir/mod.rs

@@ -2028,6 +2028,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))
     }
@@ -2553,11 +2557,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()

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,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);
                     }
                 }
             }

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)
     }
 }
 

src/librustc_codegen_llvm/debuginfo/metadata.rs

@@ -691,9 +691,12 @@ pub fn type_metadata(
                                    usage_site_span).finalize(cx)
         }
         ty::Generator(def_id, substs,  _) => {
-            let upvar_tys : Vec<_> = substs.field_tys(def_id, cx.tcx).map(|t| {
+            // TODO handle variant fields
+            let upvar_tys : Vec<_> = substs.prefix_tys(def_id, cx.tcx).map(|t| {
                 cx.tcx.normalize_erasing_regions(ParamEnv::reveal_all(), t)
             }).collect();
+            // TODO use prepare_enum_metadata and update it to handle multiple
+            // fields in the outer layout.
             prepare_tuple_metadata(cx,
                                    t,
                                    &upvar_tys,
@@ -1818,6 +1821,7 @@ fn prepare_enum_metadata(
     };
 
     // The variant part must be wrapped in a struct according to DWARF.
+    // TODO create remaining fields here, if any.
     let type_array = create_DIArray(DIB(cx), &[Some(variant_part)]);
     let struct_wrapper = unsafe {
         llvm::LLVMRustDIBuilderCreateStructType(

src/librustc_codegen_llvm/type_of.rs

@@ -63,6 +63,11 @@ fn uncached_llvm_type<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>,
                     write!(&mut name, "::{}", def.variants[index].ident).unwrap();
                 }
             }
+            if let (&ty::Generator(..), &layout::Variants::Single { index })
+                 = (&layout.ty.sty, &layout.variants)
+            {
+                write!(&mut name, "::variant#{:?}", index).unwrap();
+            }
             Some(name)
         }
         _ => None

src/librustc_codegen_ssa/mir/mod.rs

@@ -4,6 +4,7 @@ use rustc::mir::{self, Mir};
 use rustc::session::config::DebugInfo;
 use rustc_mir::monomorphize::Instance;
 use rustc_target::abi::call::{FnType, PassMode, IgnoreMode};
+use rustc_target::abi::{Variants, VariantIdx};
 use crate::base;
 use crate::debuginfo::{self, VariableAccess, VariableKind, FunctionDebugContext};
 use crate::traits::*;
@@ -648,7 +649,7 @@ fn arg_local_refs<'a, 'tcx: 'a, Bx: BuilderMethods<'a, 'tcx>>(
                     .iter()
                     .zip(upvar_tys)
                     .enumerate()
-                    .map(|(i, (upvar, ty))| (i, upvar.debug_name, upvar.by_ref, ty));
+                    .map(|(i, (upvar, ty))| (None, i, upvar.debug_name, upvar.by_ref, ty));
 
                 let generator_fields = mir.generator_layout.as_ref().map(|generator_layout| {
                     let (def_id, gen_substs) = match closure_layout.ty.sty {
@@ -658,23 +659,39 @@ fn arg_local_refs<'a, 'tcx: 'a, Bx: BuilderMethods<'a, 'tcx>>(
                     // TODO handle variant scopes here
                     let state_tys = gen_substs.state_tys(def_id, tcx);
 
-                    // TODO remove assumption of only one variant
-                    let upvar_count = mir.upvar_decls.len();
-                    generator_layout.variant_fields[0]
-                        .iter()
+                    generator_layout.variant_fields.iter()
                         .zip(state_tys)
                         .enumerate()
-                        .filter_map(move |(i, (decl, ty))| {
-                            let ty = fx.monomorphize(&ty);
-                            decl.name.map(|name| (i + upvar_count + 1, name, false, ty))
+                        .flat_map(move |(variant_idx, (decls, tys))| {
+                            let variant_idx = Some(VariantIdx::from(variant_idx));
+                            decls.iter()
+                                .zip(tys)
+                                .enumerate()
+                                .filter_map(move |(i, (decl, ty))| {
+                                    let ty = fx.monomorphize(&ty);
+                                    decl.name.map(|name| {
+                                        (variant_idx, i, name, false, ty)
+                                })
+                            })
                         })
                 }).into_iter().flatten();
 
                 upvars.chain(generator_fields)
             };
 
-            for (field, name, by_ref, ty) in extra_locals {
-                let byte_offset_of_var_in_env = closure_layout.fields.offset(field).bytes();
+            for (variant_idx, field, name, by_ref, ty) in extra_locals {
+                let fields = match variant_idx {
+                    Some(variant_idx) => {
+                        match &closure_layout.variants {
+                            Variants::Multiple { variants, .. } => {
+                                &variants[variant_idx].fields
+                            },
+                            _ => bug!("variant index on univariant layout"),
+                        }
+                    }
+                    None => &closure_layout.fields,
+                };
+                let byte_offset_of_var_in_env = fields.offset(field).bytes();
 
                 let ops = bx.debuginfo_upvar_ops_sequence(byte_offset_of_var_in_env);
 

src/librustc_codegen_ssa/mir/place.rs

@@ -296,9 +296,15 @@ impl<'a, 'tcx: 'a, V: CodegenObject> PlaceRef<'tcx, V> {
                 ..
             } => {
                 let ptr = self.project_field(bx, discr_index);
-                let to = self.layout.ty.ty_adt_def().unwrap()
-                    .discriminant_for_variant(bx.tcx(), variant_index)
-                    .val;
+                let to = match self.layout.ty.sty {
+                    ty::TyKind::Adt(adt_def, _) => adt_def
+                        .discriminant_for_variant(bx.tcx(), variant_index)
+                        .val,
+                    // Generators don't support explicit discriminant values, so
+                    // they are the same as the variant index.
+                    ty::TyKind::Generator(..) => variant_index.as_u32() as u128,
+                    _ => bug!(),
+                };
                 bx.store(
                     bx.cx().const_uint_big(bx.cx().backend_type(ptr.layout), to),
                     ptr.llval,