Fix generator transform bug rust-lang#69039

Author: cynecx

src/librustc/mir/mod.rs

@@ -1468,21 +1468,21 @@ impl<'tcx> TerminatorKind<'tcx> {
     /// successors, which may be rendered differently between the text and the graphviz format.
     pub fn fmt_head<W: Write>(&self, fmt: &mut W) -> fmt::Result {
         use self::TerminatorKind::*;
-        match *self {
+        match self {
             Goto { .. } => write!(fmt, "goto"),
-            SwitchInt { discr: ref place, .. } => write!(fmt, "switchInt({:?})", place),
+            SwitchInt { discr, .. } => write!(fmt, "switchInt({:?})", discr),
             Return => write!(fmt, "return"),
             GeneratorDrop => write!(fmt, "generator_drop"),
             Resume => write!(fmt, "resume"),
             Abort => write!(fmt, "abort"),
-            Yield { ref value, .. } => write!(fmt, "_1 = suspend({:?})", value),
+            Yield { value, resume_arg, .. } => write!(fmt, "{:?} = yield({:?})", resume_arg, value),
             Unreachable => write!(fmt, "unreachable"),
-            Drop { ref location, .. } => write!(fmt, "drop({:?})", location),
-            DropAndReplace { ref location, ref value, .. } => {
+            Drop { location, .. } => write!(fmt, "drop({:?})", location),
+            DropAndReplace { location, value, .. } => {
                 write!(fmt, "replace({:?} <- {:?})", location, value)
             }
-            Call { ref func, ref args, ref destination, .. } => {
-                if let Some((ref destination, _)) = *destination {
+            Call { func, args, destination, .. } => {
+                if let Some((destination, _)) = destination {
                     write!(fmt, "{:?} = ", destination)?;
                 }
                 write!(fmt, "{:?}(", func)?;
@@ -1494,7 +1494,7 @@ impl<'tcx> TerminatorKind<'tcx> {
                 }
                 write!(fmt, ")")
             }
-            Assert { ref cond, expected, ref msg, .. } => {
+            Assert { cond, expected, msg, .. } => {
                 write!(fmt, "assert(")?;
                 if !expected {
                     write!(fmt, "!")?;

src/librustc_mir/dataflow/impls/storage_liveness.rs

@@ -139,11 +139,12 @@ impl<'mir, 'tcx> BitDenotation<'tcx> for RequiresStorage<'mir, 'tcx> {
 
     fn before_terminator_effect(&self, sets: &mut GenKillSet<Local>, loc: Location) {
         self.check_for_borrow(sets, loc);
-
-        if let TerminatorKind::Call { destination: Some((Place { local, .. }, _)), .. } =
-            self.body[loc.block].terminator().kind
-        {
-            sets.gen(local);
+        match &self.body[loc.block].terminator().kind {
+            TerminatorKind::Call { destination: Some((Place { local, .. }, _)), .. }
+            | TerminatorKind::Yield { resume_arg: Place { local, .. }, .. } => {
+                sets.gen(*local);
+            }
+            _ => {}
         }
     }
 

src/librustc_mir/transform/generator.rs

@@ -283,7 +283,9 @@ impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
         _location: Location,
     ) {
         // Replace an Local in the remap with a generator struct access
+        debug!("generator visit_place {:?}", place);
         if let Some(&(ty, variant_index, idx)) = self.remap.get(&place.local) {
+            debug!("generator visit_place replacing with ({:?}, {:?})", variant_index, idx);
             replace_base(place, self.make_field(variant_index, idx, ty), self.tcx);
         }
     }
@@ -798,6 +800,8 @@ fn compute_layout<'tcx>(
 
     let layout = GeneratorLayout { field_tys: tys, variant_fields, storage_conflicts };
 
+    debug!("generator remap = {:?}", remap);
+
     (remap, layout, storage_liveness)
 }
 
@@ -1007,7 +1011,7 @@ fn insert_panic_block<'tcx>(
 
 fn create_generator_resume_function<'tcx>(
     tcx: TyCtxt<'tcx>,
-    transform: TransformVisitor<'tcx>,
+    transform: &TransformVisitor<'tcx>,
     def_id: DefId,
     source: MirSource<'tcx>,
     body: &mut BodyAndCache<'tcx>,
@@ -1121,10 +1125,16 @@ fn create_cases<'tcx>(
                 if operation == Operation::Resume {
                     // Move the resume argument to the destination place of the `Yield` terminator
                     let resume_arg = Local::new(2); // 0 = return, 1 = self
+                    let dst_place = match transform.remap.get(&point.resume_arg.local) {
+                        Some(&(ty, variant_index, idx)) => {
+                            transform.make_field(variant_index, idx, ty)
+                        }
+                        None => point.resume_arg,
+                    };
                     statements.push(Statement {
                         source_info,
                         kind: StatementKind::Assign(box (
-                            point.resume_arg,
+                            dst_place,
                             Rvalue::Use(Operand::Move(resume_arg.into())),
                         )),
                     });
@@ -1257,6 +1267,6 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
         body.generator_drop = Some(box drop_shim);
 
         // Create the Generator::resume function
-        create_generator_resume_function(tcx, transform, def_id, source, body);
+        create_generator_resume_function(tcx, &transform, def_id, source, body);
     }
 }

src/test/mir-opt/generator-storage-dead-unwind.rs

@@ -49,7 +49,7 @@ fn main() {
 //     StorageLive(_4);
 //     _4 = Bar(const 6i32,);
 //     ...
-//     _1 = suspend(move _6) -> [resume: bb2, drop: bb4];
+//     _5 = yield(move _6) -> [resume: bb2, drop: bb4];
 // }
 // bb1 (cleanup): {
 //     resume;