Give GeneratorLayout a list of fields for each variant

But don't really use it yet.

4 files changed, 16 insertions, 13 deletions

diff --git a/src/librustc/mir/mod.rs b/src/librustc/mir/mod.rs
index bf2a1eaafd6..5372f868aa9 100644
--- a/src/librustc/mir/mod.rs
+++ b/src/librustc/mir/mod.rs
@@ -2998,7 +2998,7 @@ pub struct UnsafetyCheckResult {
 /// The layout of generator state
 #[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable)]
 pub struct GeneratorLayout<'tcx> {
-    pub fields: Vec<LocalDecl<'tcx>>,
+    pub variant_fields: Vec<Vec<LocalDecl<'tcx>>>,
 }
 
 #[derive(Clone, Debug, RustcEncodable, RustcDecodable, HashStable)]
@@ -3187,7 +3187,7 @@ BraceStructTypeFoldableImpl! {
 
 BraceStructTypeFoldableImpl! {
     impl<'tcx> TypeFoldable<'tcx> for GeneratorLayout<'tcx> {
-        fields
+        variant_fields
     }
 }
 
diff --git a/src/librustc/ty/sty.rs b/src/librustc/ty/sty.rs
index 1ac18f6117f..edd6014618e 100644
--- a/src/librustc/ty/sty.rs
+++ b/src/librustc/ty/sty.rs
@@ -475,16 +475,16 @@ 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.
-    pub fn state_tys(
-        self,
-        def_id: DefId,
-        tcx: TyCtxt<'a, 'gcx, 'tcx>,
-    ) -> impl Iterator<Item=Ty<'tcx>> + Captures<'gcx> + 'a {
-        let state = tcx.generator_layout(def_id).fields.iter();
-        state.map(move |d| d.ty.subst(tcx, self.substs))
+    pub fn state_tys(self, def_id: DefId, tcx: TyCtxt<'a, 'gcx, 'tcx>) ->
+        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))
     }
 
-    /// This is the types of the fields of a generate which
+    /// 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>) ->
diff --git a/src/librustc_codegen_ssa/mir/mod.rs b/src/librustc_codegen_ssa/mir/mod.rs
index 621c4e5d448..4387d77a925 100644
--- a/src/librustc_codegen_ssa/mir/mod.rs
+++ b/src/librustc_codegen_ssa/mir/mod.rs
@@ -655,10 +655,12 @@ fn arg_local_refs<'a, 'tcx: 'a, Bx: BuilderMethods<'a, 'tcx>>(
                         ty::Generator(def_id, substs, _) => (def_id, substs),
                         _ => bug!("generator layout without generator substs"),
                     };
+                    // TODO handle variant scopes here
                     let state_tys = gen_substs.state_tys(def_id, tcx);
 
-                    let upvar_count = upvar_debuginfo.len();
-                    generator_layout.fields
+                    // TODO remove assumption of only one variant
+                    let upvar_count = mir.upvar_decls.len();
+                    generator_layout.variant_fields[0]
                         .iter()
                         .zip(state_tys)
                         .enumerate()
diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index cf7397ca488..253038fd030 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -549,7 +549,8 @@ fn compute_layout<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
     }).unzip();
 
     let layout = GeneratorLayout {
-        fields: vars
+        // Put everything in one variant, for now.
+        variant_fields: vec![vars]
     };
 
     (remap, layout, storage_liveness)