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 #69672 and makes #69033 pass the async fn size tests again (in other words, there will be no size regression of async fn if both this and #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)

Author: bors

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
-    /// effect has been applied to `state`.
+    /// When this flag is set, the cursor is pointing at a `Call` or `Yield` terminator whose call
+    /// 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,15 +93,15 @@ 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));
 
         // 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
-    /// be observed. Use `seek_after` if you do **not** wish to observe the call return effect.
-    pub fn seek_after_assume_call_returns(&mut self, target: Location) {
+    /// If the `target` is a `Call` or `Yield` terminator, any call return or resume effect for that
+    /// terminator will be observed. Use `seek_after` if you do **not** wish to observe the
+    /// "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 {
-            destination: Some((return_place, _)), func, args, ..
-        } = &terminator.kind
-        {
-            if !self.call_return_effect_applied {
-                self.call_return_effect_applied = true;
+        match &terminator.kind {
+            TerminatorKind::Call { destination: Some((return_place, _)), func, args, .. } => {
                 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::Yield { resume_arg: place, .. } => {
+            TerminatorKind::Call { destination: Some((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-yielding-or-returning-itself.rs renamed to src/test/ui/generator/generator-yielding-or-returning-itself.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)
     }
 }