Perform generator field computation before optimizations

Author: Aaron1011

src/librustc/mir/mod.rs

@@ -38,7 +38,7 @@ use syntax::symbol::Symbol;
 use syntax_pos::{Span, DUMMY_SP};
 
 pub use crate::mir::interpret::AssertMessage;
-pub use crate::mir::cache::{BodyCache, ReadOnlyBodyCache};
+pub use crate::mir::cache::{Cache, BodyCache, ReadOnlyBodyCache};
 pub use crate::read_only;
 
 mod cache;
@@ -117,6 +117,8 @@ pub struct Body<'tcx> {
     /// to be created.
     pub generator_kind: Option<GeneratorKind>,
 
+    pub generator_interior_tys: Option<Vec<Ty<'tcx>>>,
+
     /// Declarations of locals.
     ///
     /// The first local is the return value pointer, followed by `arg_count`
@@ -184,6 +186,7 @@ impl<'tcx> Body<'tcx> {
             generator_drop: None,
             generator_layout: None,
             generator_kind,
+            generator_interior_tys: None,
             local_decls,
             user_type_annotations,
             arg_count,

src/librustc/traits/select.rs

@@ -2755,9 +2755,9 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             ty::GeneratorWitness(did, types) => {
                 // Note that we need to use optimized_mir here,
                 // in order to have the `StateTransform` pass run
-                let gen_mir = self.tcx().optimized_mir(did);
+                /*let gen_mir = self.tcx().optimized_mir(did);
                 let gen_layout = gen_mir.generator_layout.as_ref()
-                    .expect("Missing generator layout!");
+                    .expect("Missing generator layout!");*/
 
                 // We need to compare the types from the GeneratoWitness
                 // to the types from the MIR. Since the generator MIR (specifically
@@ -2806,7 +2806,14 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
 
                 let mut used_types = Vec::new();
 
-                for ty in &gen_layout.field_tys {
+                // Note - while we use `optimized_mir`, the .generator_interior_tys
+                // field is only set during construction of the original MIR,
+                // and is unaffected by optimizations
+                let interior_tys = self.tcx().optimized_mir(did).generator_interior_tys
+                    .as_ref().expect("Missing generator interior types!");
+
+                //for ty in &gen_layout.field_tys {
+                for ty in interior_tys {
                     if let Some(witness_ty) = erased_types.get(&self.tcx().erase_regions(ty)) {
                         used_types.push(**witness_ty);
                     }

src/librustc_mir/transform/generator.rs

@@ -757,10 +757,11 @@ fn compute_layout<'tcx>(
     GeneratorLayout<'tcx>,
     FxHashMap<BasicBlock, liveness::LiveVarSet>,
 ) {
+
     // Use a liveness analysis to compute locals which are live across a suspension point
-    let LivenessInfo {
+    let (LivenessInfo {
         live_locals, live_locals_at_suspension_points, storage_conflicts, storage_liveness
-    } = locals_live_across_suspend_points(tcx, read_only!(body), source, movable);
+    }, tys) = compute_field_tys(tcx, source, movable, read_only!(body));
 
     // Erase regions from the types passed in from typeck so we can compare them with
     // MIR types
@@ -789,10 +790,8 @@ fn compute_layout<'tcx>(
 
     // Gather live local types and their indices.
     let mut locals = IndexVec::<GeneratorSavedLocal, _>::new();
-    let mut tys = IndexVec::<GeneratorSavedLocal, _>::new();
     for (idx, local) in live_locals.iter().enumerate() {
         locals.push(local);
-        tys.push(body.local_decls[local].ty);
         debug!("generator saved local {:?} => {:?}", GeneratorSavedLocal::from(idx), local);
     }
 
@@ -1168,6 +1167,40 @@ where
     }).collect()
 }
 
+pub(super) fn generator_interior_tys<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: DefId,
+    body: &BodyCache<'tcx>
+) -> Vec<Ty<'tcx>> {
+    // The first argument is the generator type passed by value
+    //let gen_ty = body.local_decls.raw[1].tyi;
+    let gen_ty = tcx.type_of(def_id);
+    let movable = match gen_ty.kind {
+        ty::Generator(_, _, movability) => {
+            movability == hir::Movability::Movable
+        },
+        _ => bug!("Unexpected type {:?}", gen_ty)
+    };
+
+    compute_field_tys(tcx, MirSource::item(def_id), movable, body.unwrap_read_only()).1.raw
+}
+
+fn compute_field_tys(
+    tcx: TyCtxt<'tcx>,
+    source: MirSource<'tcx>,
+    movable: bool,
+    body: ReadOnlyBodyCache<'_, 'tcx>
+) -> (LivenessInfo, IndexVec<GeneratorSavedLocal, Ty<'tcx>>) {
+    let liveness_info = locals_live_across_suspend_points(
+        tcx, body, source, movable
+    );
+    let mut tys = IndexVec::<GeneratorSavedLocal, _>::new();
+    for local in liveness_info.live_locals.iter() {
+        tys.push(body.local_decls[local].ty);
+    }
+    (liveness_info, tys)
+}
+
 impl<'tcx> MirPass<'tcx> for StateTransform {
     fn run_pass(
         &self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut BodyCache<'tcx>

src/librustc_mir/transform/mod.rs

@@ -2,7 +2,7 @@ use crate::{build, shim};
 use rustc_index::vec::IndexVec;
 use rustc::hir::def_id::{CrateNum, DefId, LOCAL_CRATE};
 use rustc::mir::{BodyCache, MirPhase, Promoted, ConstQualifs};
-use rustc::ty::{TyCtxt, InstanceDef, TypeFoldable};
+use rustc::ty::{self, TyCtxt, InstanceDef, TypeFoldable};
 use rustc::ty::query::Providers;
 use rustc::ty::steal::Steal;
 use rustc::hir;
@@ -98,7 +98,12 @@ fn mir_keys(tcx: TyCtxt<'_>, krate: CrateNum) -> &DefIdSet {
 }
 
 fn mir_built(tcx: TyCtxt<'_>, def_id: DefId) -> &Steal<BodyCache<'_>> {
-    let mir = build::mir_build(tcx, def_id);
+    let mut mir = build::mir_build(tcx, def_id);
+    if let ty::Generator(..) = tcx.type_of(def_id).kind {
+        let interior_types = generator::generator_interior_tys(tcx, def_id, &mir);
+        mir.generator_interior_tys = Some(interior_types);
+    }
+
     tcx.alloc_steal_mir(mir)
 }