add pre-statement-effect to dataflow

4 files changed, 77 insertions, 5 deletions

diff --git a/src/librustc_mir/dataflow/at_location.rs b/src/librustc_mir/dataflow/at_location.rs
index 7f243ad6e26..b1f73bfbe22 100644
--- a/src/librustc_mir/dataflow/at_location.rs
+++ b/src/librustc_mir/dataflow/at_location.rs
@@ -149,6 +149,18 @@ impl<BD> FlowsAtLocation for FlowAtLocation<BD>
     fn reconstruct_statement_effect(&mut self, loc: Location) {
         self.stmt_gen.reset_to_empty();
         self.stmt_kill.reset_to_empty();
+        {
+            let mut sets = BlockSets {
+                on_entry: &mut self.curr_state,
+                gen_set: &mut self.stmt_gen,
+                kill_set: &mut self.stmt_kill,
+            };
+            self.base_results
+                .operator()
+                .before_statement_effect(&mut sets, loc);
+        }
+        self.apply_local_effect(loc);
+
         let mut sets = BlockSets {
             on_entry: &mut self.curr_state,
             gen_set: &mut self.stmt_gen,
@@ -162,6 +174,18 @@ impl<BD> FlowsAtLocation for FlowAtLocation<BD>
     fn reconstruct_terminator_effect(&mut self, loc: Location) {
         self.stmt_gen.reset_to_empty();
         self.stmt_kill.reset_to_empty();
+        {
+            let mut sets = BlockSets {
+                on_entry: &mut self.curr_state,
+                gen_set: &mut self.stmt_gen,
+                kill_set: &mut self.stmt_kill,
+            };
+            self.base_results
+                .operator()
+                .before_terminator_effect(&mut sets, loc);
+        }
+        self.apply_local_effect(loc);
+
         let mut sets = BlockSets {
             on_entry: &mut self.curr_state,
             gen_set: &mut self.stmt_gen,
diff --git a/src/librustc_mir/dataflow/mod.rs b/src/librustc_mir/dataflow/mod.rs
index 83c46e0199e..19333dec3bc 100644
--- a/src/librustc_mir/dataflow/mod.rs
+++ b/src/librustc_mir/dataflow/mod.rs
@@ -214,6 +214,7 @@ impl<'a, 'tcx: 'a, BD> DataflowAnalysis<'a, 'tcx, BD> where BD: BitDenotation
             }
             for j_stmt in 0..statements.len() {
                 let location = Location { block: bb, statement_index: j_stmt };
+                self.flow_state.operator.before_statement_effect(sets, location);
                 self.flow_state.operator.statement_effect(sets, location);
                 if track_intrablock {
                     sets.apply_local_effect();
@@ -222,6 +223,7 @@ impl<'a, 'tcx: 'a, BD> DataflowAnalysis<'a, 'tcx, BD> where BD: BitDenotation
 
             if terminator.is_some() {
                 let location = Location { block: bb, statement_index: statements.len() };
+                self.flow_state.operator.before_terminator_effect(sets, location);
                 self.flow_state.operator.terminator_effect(sets, location);
                 if track_intrablock {
                     sets.apply_local_effect();
@@ -365,9 +367,10 @@ pub(crate) trait DataflowResultsConsumer<'a, 'tcx: 'a> {
     fn mir(&self) -> &'a Mir<'tcx>;
 }
 
-pub fn state_for_location<T: BitDenotation>(loc: Location,
-                                            analysis: &T,
-                                            result: &DataflowResults<T>)
+pub fn state_for_location<'tcx, T: BitDenotation>(loc: Location,
+                                                  analysis: &T,
+                                                  result: &DataflowResults<T>,
+                                                  mir: &Mir<'tcx>)
     -> IdxSetBuf<T::Idx> {
     let mut entry = result.sets().on_entry_set_for(loc.block.index()).to_owned();
 
@@ -381,8 +384,16 @@ pub fn state_for_location<T: BitDenotation>(loc: Location,
         for stmt in 0..loc.statement_index {
             let mut stmt_loc = loc;
             stmt_loc.statement_index = stmt;
+            analysis.before_statement_effect(&mut sets, stmt_loc);
             analysis.statement_effect(&mut sets, stmt_loc);
         }
+
+        // Apply the pre-statement effect of the statement we're evaluating.
+        if loc.statement_index == mir[loc.block].statements.len() {
+            analysis.before_terminator_effect(&mut sets, loc);
+        } else {
+            analysis.before_statement_effect(&mut sets, loc);
+        }
     }
 
     entry
@@ -637,6 +648,21 @@ pub trait BitDenotation: BitwiseOperator {
     /// (For example, establishing the call arguments.)
     fn start_block_effect(&self, entry_set: &mut IdxSet<Self::Idx>);
 
+    /// Similar to `statement_effect`, except it applies
+    /// *just before* the statement rather than *just after* it.
+    ///
+    /// This matters for "dataflow at location" APIs, because the
+    /// before-statement effect is visible while visiting the
+    /// statement, while the after-statement effect only becomes
+    /// visible at the next statement.
+    ///
+    /// Both the before-statement and after-statement effects are
+    /// applied, in that order, before moving for the next
+    /// statement.
+    fn before_statement_effect(&self,
+                               _sets: &mut BlockSets<Self::Idx>,
+                               _location: Location) {}
+
     /// Mutates the block-sets (the flow sets for the given
     /// basic block) according to the effects of evaluating statement.
     ///
@@ -651,6 +677,21 @@ pub trait BitDenotation: BitwiseOperator {
                         sets: &mut BlockSets<Self::Idx>,
                         location: Location);
 
+    /// Similar to `terminator_effect`, except it applies
+    /// *just before* the terminator rather than *just after* it.
+    ///
+    /// This matters for "dataflow at location" APIs, because the
+    /// before-terminator effect is visible while visiting the
+    /// terminator, while the after-terminator effect only becomes
+    /// visible at the terminator's successors.
+    ///
+    /// Both the before-terminator and after-terminator effects are
+    /// applied, in that order, before moving for the next
+    /// terminator.
+    fn before_terminator_effect(&self,
+                                _sets: &mut BlockSets<Self::Idx>,
+                                _location: Location) {}
+
     /// Mutates the block-sets (the flow sets for the given
     /// basic block) according to the effects of evaluating
     /// the terminator.
diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index 455a07c04cf..b226bb592cb 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -363,7 +363,7 @@ fn locals_live_across_suspend_points<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                 statement_index: data.statements.len(),
             };
 
-            let storage_liveness = state_for_location(loc, &analysis, &storage_live);
+            let storage_liveness = state_for_location(loc, &analysis, &storage_live, mir);
 
             storage_liveness_map.insert(block, storage_liveness.clone());
 
diff --git a/src/librustc_mir/transform/rustc_peek.rs b/src/librustc_mir/transform/rustc_peek.rs
index 6b8e2b073cc..b6153ea1fdb 100644
--- a/src/librustc_mir/transform/rustc_peek.rs
+++ b/src/librustc_mir/transform/rustc_peek.rs
@@ -203,11 +203,18 @@ fn each_block<'a, 'tcx, O>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         // reset GEN and KILL sets before emulating their effect.
         for e in sets.gen_set.words_mut() { *e = 0; }
         for e in sets.kill_set.words_mut() { *e = 0; }
-        results.0.operator.statement_effect(&mut sets, Location { block: bb, statement_index: j });
+        results.0.operator.before_statement_effect(
+            &mut sets, Location { block: bb, statement_index: j });
+        results.0.operator.statement_effect(
+            &mut sets, Location { block: bb, statement_index: j });
         sets.on_entry.union(sets.gen_set);
         sets.on_entry.subtract(sets.kill_set);
     }
 
+    results.0.operator.before_terminator_effect(
+        &mut sets,
+        Location { block: bb, statement_index: statements.len() });
+
     tcx.sess.span_err(span, &format!("rustc_peek: MIR did not match \
                                       anticipated pattern; note that \
                                       rustc_peek expects input of \