Auto merge of #69302 - jonas-schievink:yield-needs-storage, r=Zoxc

Fix generator miscompilations

Fixes https://github.com/rust-lang/rust/issues/69039

r? @Zoxc

src/librustc/mir/visit.rs

@@ -519,12 +519,12 @@ macro_rules! make_mir_visitor {
                         resume_arg,
                         drop: _,
                     } => {
+                        self.visit_operand(value, source_location);
                         self.visit_place(
                             resume_arg,
                             PlaceContext::MutatingUse(MutatingUseContext::Store),
                             source_location,
                         );
-                        self.visit_operand(value, source_location);
                     }
 
                 }

src/librustc_mir/dataflow/impls/storage_liveness.rs

@@ -118,18 +118,25 @@ impl<'mir, 'tcx> BitDenotation<'tcx> for RequiresStorage<'mir, 'tcx> {
         self.borrowed_locals.borrow().analysis().statement_effect(sets, stmt, loc);
 
         // If a place is assigned to in a statement, it needs storage for that statement.
-        match stmt.kind {
-            StatementKind::StorageDead(l) => sets.kill(l),
-            StatementKind::Assign(box (ref place, _))
-            | StatementKind::SetDiscriminant { box ref place, .. } => {
+        match &stmt.kind {
+            StatementKind::StorageDead(l) => sets.kill(*l),
+            StatementKind::Assign(box (place, _))
+            | StatementKind::SetDiscriminant { box place, .. } => {
                 sets.gen(place.local);
             }
-            StatementKind::InlineAsm(box InlineAsm { ref outputs, .. }) => {
+            StatementKind::InlineAsm(box InlineAsm { outputs, .. }) => {
                 for place in &**outputs {
                     sets.gen(place.local);
                 }
             }
-            _ => (),
+
+            // Nothing to do for these. Match exhaustively so this fails to compile when new
+            // variants are added.
+            StatementKind::AscribeUserType(..)
+            | StatementKind::FakeRead(..)
+            | StatementKind::Nop
+            | StatementKind::Retag(..)
+            | StatementKind::StorageLive(..) => {}
         }
     }
 
@@ -145,23 +152,58 @@ impl<'mir, 'tcx> BitDenotation<'tcx> for RequiresStorage<'mir, 'tcx> {
         // If a place is borrowed in a terminator, it needs storage for that terminator.
         self.borrowed_locals.borrow().analysis().terminator_effect(sets, terminator, loc);
 
-        if let TerminatorKind::Call { destination: Some((place, _)), .. } = terminator.kind {
-            sets.gen(place.local);
+        match &terminator.kind {
+            TerminatorKind::Call { destination: Some((Place { local, .. }, _)), .. }
+            | TerminatorKind::Yield { resume_arg: Place { local, .. }, .. } => {
+                sets.gen(*local);
+            }
+
+            // Nothing to do for these. Match exhaustively so this fails to compile when new
+            // variants are added.
+            TerminatorKind::Call { destination: None, .. }
+            | TerminatorKind::Abort
+            | TerminatorKind::Assert { .. }
+            | TerminatorKind::Drop { .. }
+            | TerminatorKind::DropAndReplace { .. }
+            | TerminatorKind::FalseEdges { .. }
+            | TerminatorKind::FalseUnwind { .. }
+            | TerminatorKind::GeneratorDrop
+            | TerminatorKind::Goto { .. }
+            | TerminatorKind::Resume
+            | TerminatorKind::Return
+            | TerminatorKind::SwitchInt { .. }
+            | TerminatorKind::Unreachable => {}
         }
     }
 
     fn terminator_effect(&self, sets: &mut GenKillSet<Local>, loc: Location) {
-        // For call terminators the destination requires storage for the call
-        // and after the call returns successfully, but not after a panic.
-        // Since `propagate_call_unwind` doesn't exist, we have to kill the
-        // destination here, and then gen it again in `propagate_call_return`.
-        if let TerminatorKind::Call { destination: Some((ref place, _)), .. } =
-            self.body[loc.block].terminator().kind
-        {
-            if let Some(local) = place.as_local() {
-                sets.kill(local);
+        match &self.body[loc.block].terminator().kind {
+            // For call terminators the destination requires storage for the call
+            // and after the call returns successfully, but not after a panic.
+            // Since `propagate_call_unwind` doesn't exist, we have to kill the
+            // destination here, and then gen it again in `propagate_call_return`.
+            TerminatorKind::Call { destination: Some((Place { local, .. }, _)), .. } => {
+                sets.kill(*local);
             }
+
+            // Nothing to do for these. Match exhaustively so this fails to compile when new
+            // variants are added.
+            TerminatorKind::Call { destination: None, .. }
+            | TerminatorKind::Yield { .. }
+            | TerminatorKind::Abort
+            | TerminatorKind::Assert { .. }
+            | TerminatorKind::Drop { .. }
+            | TerminatorKind::DropAndReplace { .. }
+            | TerminatorKind::FalseEdges { .. }
+            | TerminatorKind::FalseUnwind { .. }
+            | TerminatorKind::GeneratorDrop
+            | TerminatorKind::Goto { .. }
+            | TerminatorKind::Resume
+            | TerminatorKind::Return
+            | TerminatorKind::SwitchInt { .. }
+            | TerminatorKind::Unreachable => {}
         }
+
         self.check_for_move(sets, loc);
     }
 

src/librustc_mir/dataflow/move_paths/builder.rs

@@ -381,9 +381,9 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
             }
 
             TerminatorKind::Yield { ref value, resume_arg: ref place, .. } => {
+                self.gather_operand(value);
                 self.create_move_path(place);
                 self.gather_init(place.as_ref(), InitKind::Deep);
-                self.gather_operand(value);
             }
 
             TerminatorKind::Drop { ref location, target: _, unwind: _ } => {

src/librustc_mir/transform/generator.rs

@@ -186,18 +186,24 @@ fn self_arg() -> Local {
     Local::new(1)
 }
 
-/// Generator have not been resumed yet
+/// Generator has not been resumed yet.
 const UNRESUMED: usize = GeneratorSubsts::UNRESUMED;
-/// Generator has returned / is completed
+/// Generator has returned / is completed.
 const RETURNED: usize = GeneratorSubsts::RETURNED;
-/// Generator has been poisoned
+/// Generator has panicked and is poisoned.
 const POISONED: usize = GeneratorSubsts::POISONED;
 
+/// A `yield` point in the generator.
 struct SuspensionPoint<'tcx> {
+    /// State discriminant used when suspending or resuming at this point.
     state: usize,
+    /// The block to jump to after resumption.
     resume: BasicBlock,
+    /// Where to move the resume argument after resumption.
     resume_arg: Place<'tcx>,
+    /// Which block to jump to if the generator is dropped in this state.
     drop: Option<BasicBlock>,
+    /// Set of locals that have live storage while at this suspension point.
     storage_liveness: liveness::LiveVarSet,
 }
 
@@ -325,6 +331,15 @@ impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
                 // Yield
                 let state = 3 + self.suspension_points.len();
 
+                // The resume arg target location might itself be remapped if its base local is
+                // live across a yield.
+                let resume_arg =
+                    if let Some(&(ty, variant, idx)) = self.remap.get(&resume_arg.local) {
+                        self.make_field(variant, idx, ty)
+                    } else {
+                        resume_arg
+                    };
+
                 self.suspension_points.push(SuspensionPoint {
                     state,
                     resume,

src/test/ui/generator/issue-69039.rs

@@ -0,0 +1,30 @@
+// run-pass
+
+#![feature(generators, generator_trait)]
+
+use std::ops::{Generator, GeneratorState};
+
+fn my_scenario() -> impl Generator<String, Yield = &'static str, Return = String> {
+    |_arg: String| {
+        let my_name = yield "What is your name?";
+        let my_mood = yield "How are you feeling?";
+        format!("{} is {}", my_name.trim(), my_mood.trim())
+    }
+}
+
+fn main() {
+    let mut my_session = Box::pin(my_scenario());
+
+    assert_eq!(
+        my_session.as_mut().resume("_arg".to_string()),
+        GeneratorState::Yielded("What is your name?")
+    );
+    assert_eq!(
+        my_session.as_mut().resume("Your Name".to_string()),
+        GeneratorState::Yielded("How are you feeling?")
+    );
+    assert_eq!(
+        my_session.as_mut().resume("Sensory Organs".to_string()),
+        GeneratorState::Complete("Your Name is Sensory Organs".to_string())
+    );
+}