Auto merge of #69716 - jonas-schievink:generator-size, r=tmandry

Don't store locals in generators that are immediately overwritten with the resume argument

This fixes https://github.com/rust-lang/rust/issues/69672 and makes https://github.com/rust-lang/rust/pull/69033 pass the async fn size tests again (in other words, there will be no size regression of async fn if both this and https://github.com/rust-lang/rust/pull/69033 land).

~~This is a small botch and I'd rather have a more precise analysis, but that seems much harder to pull off, so this special-cases `Yield` terminators that store the resume argument into a simple local (ie. without any field projections) and explicitly marks that local as "not live" in the suspend point of that yield. We know that this local does not need to be stored in the generator for this suspend point because the next resume would immediately overwrite it with the passed-in resume argument anyways. The local might still end up in the state if it is used across another yield.~~ (this now properly updates the dataflow framework to handle this case)

src/librustc_mir/dataflow/generic/cursor.rs

@@ -2,7 +2,7 @@
 
 use std::borrow::Borrow;
 
-use rustc::mir::{self, BasicBlock, Location};
+use rustc::mir::{self, BasicBlock, Location, TerminatorKind};
 use rustc_index::bit_set::BitSet;
 
 use super::{Analysis, Results};
@@ -29,14 +29,14 @@ where
 
     pos: CursorPosition,
 
-    /// When this flag is set, the cursor is pointing at a `Call` terminator whose call return
+    /// When this flag is set, the cursor is pointing at a `Call` or `Yield` terminator whose call
-    /// effect has been applied to `state`.
+    /// return or resume effect has been applied to `state`.
     ///
-    /// This flag helps to ensure that multiple calls to `seek_after_assume_call_returns` with the
+    /// This flag helps to ensure that multiple calls to `seek_after_assume_success` with the
     /// same target will result in exactly one invocation of `apply_call_return_effect`. It is
     /// sufficient to clear this only in `seek_to_block_start`, since seeking away from a
     /// terminator will always require a cursor reset.
-    call_return_effect_applied: bool,
+    success_effect_applied: bool,
 }
 
 impl<'mir, 'tcx, A, R> ResultsCursor<'mir, 'tcx, A, R>
@@ -50,7 +50,7 @@ where
             body,
             pos: CursorPosition::BlockStart(mir::START_BLOCK),
             state: results.borrow().entry_sets[mir::START_BLOCK].clone(),
-            call_return_effect_applied: false,
+            success_effect_applied: false,
             results,
         }
     }
@@ -76,14 +76,14 @@ where
     pub fn seek_to_block_start(&mut self, block: BasicBlock) {
         self.state.overwrite(&self.results.borrow().entry_sets[block]);
         self.pos = CursorPosition::BlockStart(block);
-        self.call_return_effect_applied = false;
+        self.success_effect_applied = false;
     }
 
     /// Advances the cursor to hold all effects up to and including to the "before" effect of the
     /// statement (or terminator) at the given location.
     ///
     /// If you wish to observe the full effect of a statement or terminator, not just the "before"
-    /// effect, use `seek_after` or `seek_after_assume_call_returns`.
+    /// effect, use `seek_after` or `seek_after_assume_success`.
     pub fn seek_before(&mut self, target: Location) {
         assert!(target <= self.body.terminator_loc(target.block));
         self.seek_(target, false);
@@ -93,7 +93,7 @@ where
     /// terminators) up to and including the `target`.
     ///
     /// If the `target` is a `Call` terminator, any call return effect for that terminator will
-    /// **not** be observed. Use `seek_after_assume_call_returns` if you wish to observe the call
+    /// **not** be observed. Use `seek_after_assume_success` if you wish to observe the call
     /// return effect.
     pub fn seek_after(&mut self, target: Location) {
         assert!(target <= self.body.terminator_loc(target.block));
@@ -101,7 +101,7 @@ where
         // If we have already applied the call return effect, we are currently pointing at a `Call`
         // terminator. Unconditionally reset the dataflow cursor, since there is no way to "undo"
         // the call return effect.
-        if self.call_return_effect_applied {
+        if self.success_effect_applied {
             self.seek_to_block_start(target.block);
         }
 
@@ -111,25 +111,25 @@ where
     /// Advances the cursor to hold all effects up to and including of the statement (or
     /// terminator) at the given location.
     ///
-    /// If the `target` is a `Call` terminator, any call return effect for that terminator will
+    /// If the `target` is a `Call` or `Yield` terminator, any call return or resume effect for that
-    /// be observed. Use `seek_after` if you do **not** wish to observe the call return effect.
+    /// terminator will be observed. Use `seek_after` if you do **not** wish to observe the
-    pub fn seek_after_assume_call_returns(&mut self, target: Location) {
+    /// "success" effect.
+    pub fn seek_after_assume_success(&mut self, target: Location) {
         let terminator_loc = self.body.terminator_loc(target.block);
         assert!(target.statement_index <= terminator_loc.statement_index);
 
         self.seek_(target, true);
 
-        if target != terminator_loc {
+        if target != terminator_loc || self.success_effect_applied {
             return;
         }
 
+        // Apply the effect of the "success" path of the terminator.
+
+        self.success_effect_applied = true;
         let terminator = self.body.basic_blocks()[target.block].terminator();
-        if let mir::TerminatorKind::Call {
+        match &terminator.kind {
-            destination: Some((return_place, _)), func, args, ..
+            TerminatorKind::Call { destination: Some((return_place, _)), func, args, .. } => {
-        } = &terminator.kind
-        {
-            if !self.call_return_effect_applied {
-                self.call_return_effect_applied = true;
                 self.results.borrow().analysis.apply_call_return_effect(
                     &mut self.state,
                     target.block,
@@ -138,6 +138,14 @@ where
                     return_place,
                 );
             }
+            TerminatorKind::Yield { resume, resume_arg, .. } => {
+                self.results.borrow().analysis.apply_yield_resume_effect(
+                    &mut self.state,
+                    *resume,
+                    resume_arg,
+                );
+            }
+            _ => {}
         }
     }
 
@@ -172,7 +180,7 @@ where
                 self.seek_to_block_start(target.block)
             }
 
-            // N.B., `call_return_effect_applied` is checked in `seek_after`, not here.
+            // N.B., `success_effect_applied` is checked in `seek_after`, not here.
             _ => (),
         }
 

src/librustc_mir/dataflow/generic/engine.rs

@@ -218,15 +218,18 @@ where
 
             Goto { target }
             | Assert { target, cleanup: None, .. }
-            | Yield { resume: target, drop: None, .. }
             | Drop { target, location: _, unwind: None }
             | DropAndReplace { target, value: _, location: _, unwind: None } => {
                 self.propagate_bits_into_entry_set_for(in_out, target, dirty_list)
             }
 
-            Yield { resume: target, drop: Some(drop), .. } => {
+            Yield { resume: target, drop, resume_arg, .. } => {
+                if let Some(drop) = drop {
+                    self.propagate_bits_into_entry_set_for(in_out, drop, dirty_list);
+                }
+
+                self.analysis.apply_yield_resume_effect(in_out, target, &resume_arg);
                 self.propagate_bits_into_entry_set_for(in_out, target, dirty_list);
-                self.propagate_bits_into_entry_set_for(in_out, drop, dirty_list);
             }
 
             Assert { target, cleanup: Some(unwind), .. }

src/librustc_mir/dataflow/generic/graphviz.rs

@@ -241,7 +241,7 @@ where
                     )?;
 
                     let state_on_unwind = this.results.get().clone();
-                    this.results.seek_after_assume_call_returns(terminator_loc);
+                    this.results.seek_after_assume_success(terminator_loc);
                     write_diff(w, this.results.analysis(), &state_on_unwind, this.results.get())?;
 
                     write!(w, "</td>")

src/librustc_mir/dataflow/generic/mod.rs

@@ -191,6 +191,20 @@ pub trait Analysis<'tcx>: AnalysisDomain<'tcx> {
         return_place: &mir::Place<'tcx>,
     );
 
+    /// Updates the current dataflow state with the effect of resuming from a `Yield` terminator.
+    ///
+    /// This is similar to `apply_call_return_effect` in that it only takes place after the
+    /// generator is resumed, not when it is dropped.
+    ///
+    /// By default, no effects happen.
+    fn apply_yield_resume_effect(
+        &self,
+        _state: &mut BitSet<Self::Idx>,
+        _resume_block: BasicBlock,
+        _resume_place: &mir::Place<'tcx>,
+    ) {
+    }
+
     /// Updates the current dataflow state with the effect of taking a particular branch in a
     /// `SwitchInt` terminator.
     ///
@@ -284,6 +298,15 @@ pub trait GenKillAnalysis<'tcx>: Analysis<'tcx> {
         return_place: &mir::Place<'tcx>,
     );
 
+    /// See `Analysis::apply_yield_resume_effect`.
+    fn yield_resume_effect(
+        &self,
+        _trans: &mut BitSet<Self::Idx>,
+        _resume_block: BasicBlock,
+        _resume_place: &mir::Place<'tcx>,
+    ) {
+    }
+
     /// See `Analysis::apply_discriminant_switch_effect`.
     fn discriminant_switch_effect(
         &self,
@@ -347,6 +370,15 @@ where
         self.call_return_effect(state, block, func, args, return_place);
     }
 
+    fn apply_yield_resume_effect(
+        &self,
+        state: &mut BitSet<Self::Idx>,
+        resume_block: BasicBlock,
+        resume_place: &mir::Place<'tcx>,
+    ) {
+        self.yield_resume_effect(state, resume_block, resume_place);
+    }
+
     fn apply_discriminant_switch_effect(
         &self,
         state: &mut BitSet<Self::Idx>,

src/librustc_mir/dataflow/generic/tests.rs

@@ -294,7 +294,7 @@ fn cursor_seek() {
 
     cursor.seek_after(call_terminator_loc);
     assert!(!cursor.get().contains(call_return_effect));
-    cursor.seek_after_assume_call_returns(call_terminator_loc);
+    cursor.seek_after_assume_success(call_terminator_loc);
     assert!(cursor.get().contains(call_return_effect));
 
     let every_target = || {
@@ -310,7 +310,7 @@ fn cursor_seek() {
             BlockStart(block) => cursor.seek_to_block_start(block),
             Before(loc) => cursor.seek_before(loc),
             After(loc) => cursor.seek_after(loc),
-            AfterAssumeCallReturns(loc) => cursor.seek_after_assume_call_returns(loc),
+            AfterAssumeCallReturns(loc) => cursor.seek_after_assume_success(loc),
         }
 
         assert_eq!(cursor.get(), &cursor.analysis().expected_state_at_target(targ));

src/librustc_mir/dataflow/impls/storage_liveness.rs

@@ -161,11 +161,16 @@ impl<'mir, 'tcx> dataflow::GenKillAnalysis<'tcx> for MaybeRequiresStorage<'mir,
         self.borrowed_locals.borrow().analysis().terminator_effect(trans, terminator, loc);
 
         match &terminator.kind {
-            TerminatorKind::Call { destination: Some((place, _)), .. }
+            TerminatorKind::Call { destination: Some((place, _)), .. } => {
-            | TerminatorKind::Yield { resume_arg: place, .. } => {
                 trans.gen(place.local);
             }
 
+            // Note that we do *not* gen the `resume_arg` of `Yield` terminators. The reason for
+            // that is that a `yield` will return from the function, and `resume_arg` is written
+            // only when the generator is later resumed. Unlike `Call`, this doesn't require the
+            // place to have storage *before* the yield, only after.
+            TerminatorKind::Yield { .. } => {}
+
             // Nothing to do for these. Match exhaustively so this fails to compile when new
             // variants are added.
             TerminatorKind::Call { destination: None, .. }
@@ -230,6 +235,15 @@ impl<'mir, 'tcx> dataflow::GenKillAnalysis<'tcx> for MaybeRequiresStorage<'mir,
     ) {
         trans.gen(return_place.local);
     }
+
+    fn yield_resume_effect(
+        &self,
+        trans: &mut BitSet<Self::Idx>,
+        _resume_block: BasicBlock,
+        resume_place: &mir::Place<'tcx>,
+    ) {
+        trans.gen(resume_place.local);
+    }
 }
 
 impl<'mir, 'tcx> MaybeRequiresStorage<'mir, 'tcx> {

src/librustc_mir/transform/generator.rs

@@ -539,7 +539,7 @@ fn locals_live_across_suspend_points(
             let mut live_locals_here = storage_required;
             live_locals_here.intersect(&liveness.outs[block]);
 
-            // The generator argument is ignored
+            // The generator argument is ignored.
             live_locals_here.remove(self_arg());
 
             debug!("loc = {:?}, live_locals_here = {:?}", loc, live_locals_here);

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

@@ -4,11 +4,15 @@
 
 use std::ops::{Generator, GeneratorState};
 
+fn mkstr(my_name: String, my_mood: String) -> String {
+    format!("{} is {}", my_name.trim(), my_mood.trim())
+}
+
 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())
+        mkstr(my_name, my_mood)
     }
 }
 

src/test/ui/generator/resume-arg-size.rs

@@ -0,0 +1,28 @@
+#![feature(generators)]
+
+// run-pass
+
+use std::mem::size_of_val;
+
+fn main() {
+    // Generator taking a `Copy`able resume arg.
+    let gen_copy = |mut x: usize| {
+        loop {
+            drop(x);
+            x = yield;
+        }
+    };
+
+    // Generator taking a non-`Copy` resume arg.
+    let gen_move = |mut x: Box<usize>| {
+        loop {
+            drop(x);
+            x = yield;
+        }
+    };
+
+    // Neither of these generators have the resume arg live across the `yield`, so they should be
+    // 4 Bytes in size (only storing the discriminant)
+    assert_eq!(size_of_val(&gen_copy), 4);
+    assert_eq!(size_of_val(&gen_move), 4);
+}