Support variantful generators

This allows generators to overlap fields using variants.

2 files changed, 36 insertions, 13 deletions

diff --git a/src/librustc_codegen_ssa/mir/mod.rs b/src/librustc_codegen_ssa/mir/mod.rs
index 4387d77a925..52429294852 100644
--- a/src/librustc_codegen_ssa/mir/mod.rs
+++ b/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);
 
diff --git a/src/librustc_codegen_ssa/mir/place.rs b/src/librustc_codegen_ssa/mir/place.rs
index 1134707f96c..2875468127e 100644
--- a/src/librustc_codegen_ssa/mir/place.rs
+++ b/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,