Model generator resumption in dataflow

We now have a way to apply an effect only *after* a `yield` resumes,
similar to calls (which can either return or unwind).

Author: jonas-schievink

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

@@ -506,7 +506,7 @@ fn locals_live_across_suspend_points(
 
     for (block, data) in body.basic_blocks().iter_enumerated() {
         if let TerminatorKind::Yield { .. } = data.terminator().kind {
-            let loc = Location { block: block, statement_index: data.statements.len() };
+            let loc = Location { block, statement_index: data.statements.len() };
 
             if !movable {
                 // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
@@ -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);