Auto merge of #71006 - ecstatic-morse:dataflow-bidi, r=ecstatic-morse

Use existing framework for backward dataflow analyses

This PR adds support for backward analyses to the dataflow framework and adds a new live variable analysis (based on the existing one in `librustc_mir/util/liveness.rs`). By adding these to the framework instead of having a separate API, all newly implemented backward dataflow analyses get cursors/visitors, `rustc_peek` tests, and graphviz visualizations for free. In the near-term, this makes it much easier to implement global dead-store elimination, and I believe that this will enable even more MIR optimizations in the future.

This PR makes many changes to the dataflow API, since some concepts and terminology only make sense in forward dataflow. Below is a list of the important changes.
- ~~`entry_set` -> `fixpoint` (the fixpoint for backward dataflow problems is after the block's terminator)~~
- `seek_{before,after}` -> `seek_{before,after}_primary_effect` (the unprefixed dataflow effect is now referred to as the "primary" effect instead of the "after" effect. The "before" effect remains the same, although I considered changing it to the "antecedent" effect. In both backward and forward dataflow, the "before" effect is applied prior to the "primary" effect. I feel very strongly that this is the correct choice, as it means consumers don't have to switch between `seek_before` and `seek_after` based on the direction of their analysis.
- `seek_after_assume_call_returns` is now gone. Users can use `ResultsCursor::apply_custom_effect` to emulate it.
- `visit_{statement,terminator}_exit` -> `visit_{statement,terminator}_after_primary_effect`
- `visit_{statement,terminator}` -> `visit_{statement,terminator}_before_primary_effect`

Implementing this also required refactoring the dataflow cursor implementation so it could work in both directions. This is a large percentage of the diff, since the cursor code is rather complex. The fact that the cursor is exhaustively tested in both directions should reassure whomever is unlucky enough to review this :rofl:.

In order to avoid computing the reverse CFG for forward dataflow analyses, I've added some hacks to the existing `mir::BodyAndCache` interface. I've requested changes to this interface that would let me implement this more efficiently.

r? @eddyb (feel free to reassign)
cc @rust-lang/wg-mir-opt

4 files changed, 93 insertions, 59 deletions

diff --git a/src/librustc_mir/transform/check_consts/validation.rs b/src/librustc_mir/transform/check_consts/validation.rs
index 78cb60ea484..691903c9069 100644
--- a/src/librustc_mir/transform/check_consts/validation.rs
+++ b/src/librustc_mir/transform/check_consts/validation.rs
@@ -61,7 +61,7 @@ impl Qualifs<'mir, 'tcx> {
                 .into_results_cursor(&body)
         });
 
-        indirectly_mutable.seek_before(location);
+        indirectly_mutable.seek_before_primary_effect(location);
         indirectly_mutable.get().contains(local)
     }
 
@@ -88,7 +88,7 @@ impl Qualifs<'mir, 'tcx> {
                 .into_results_cursor(&body)
         });
 
-        needs_drop.seek_before(location);
+        needs_drop.seek_before_primary_effect(location);
         needs_drop.get().contains(local) || self.indirectly_mutable(ccx, local, location)
     }
 
@@ -115,7 +115,7 @@ impl Qualifs<'mir, 'tcx> {
                 .into_results_cursor(&body)
         });
 
-        has_mut_interior.seek_before(location);
+        has_mut_interior.seek_before_primary_effect(location);
         has_mut_interior.get().contains(local) || self.indirectly_mutable(ccx, local, location)
     }
 
@@ -161,7 +161,7 @@ impl Qualifs<'mir, 'tcx> {
                     .iterate_to_fixpoint()
                     .into_results_cursor(&ccx.body);
 
-                cursor.seek_after(return_loc);
+                cursor.seek_after_primary_effect(return_loc);
                 cursor.contains(RETURN_PLACE)
             }
         };
diff --git a/src/librustc_mir/transform/elaborate_drops.rs b/src/librustc_mir/transform/elaborate_drops.rs
index 6074619dd15..a1becf062ee 100644
--- a/src/librustc_mir/transform/elaborate_drops.rs
+++ b/src/librustc_mir/transform/elaborate_drops.rs
@@ -101,7 +101,7 @@ fn find_dead_unwinds<'tcx>(
             }
         };
 
-        flow_inits.seek_before(body.terminator_loc(bb));
+        flow_inits.seek_before_primary_effect(body.terminator_loc(bb));
         debug!(
             "find_dead_unwinds @ {:?}: path({:?})={:?}; init_data={:?}",
             bb,
@@ -131,8 +131,8 @@ struct InitializationData<'mir, 'tcx> {
 
 impl InitializationData<'_, '_> {
     fn seek_before(&mut self, loc: Location) {
-        self.inits.seek_before(loc);
-        self.uninits.seek_before(loc);
+        self.inits.seek_before_primary_effect(loc);
+        self.uninits.seek_before_primary_effect(loc);
     }
 
     fn maybe_live_dead(&self, path: MovePathIndex) -> (bool, bool) {
diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index 00caa06a6c0..25804c2a62c 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -50,12 +50,13 @@
 //! Otherwise it drops all the values in scope at the last suspension point.
 
 use crate::dataflow::{self, Analysis};
-use crate::dataflow::{MaybeBorrowedLocals, MaybeRequiresStorage, MaybeStorageLive};
+use crate::dataflow::{
+    MaybeBorrowedLocals, MaybeLiveLocals, MaybeRequiresStorage, MaybeStorageLive,
+};
 use crate::transform::no_landing_pads::no_landing_pads;
 use crate::transform::simplify;
 use crate::transform::{MirPass, MirSource};
 use crate::util::dump_mir;
-use crate::util::liveness;
 use crate::util::storage;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
@@ -195,7 +196,7 @@ struct SuspensionPoint<'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,
+    storage_liveness: BitSet<Local>,
 }
 
 struct TransformVisitor<'tcx> {
@@ -211,7 +212,7 @@ struct TransformVisitor<'tcx> {
     remap: FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
 
     // A map from a suspension point in a block to the locals which have live storage at that point
-    storage_liveness: IndexVec<BasicBlock, Option<liveness::LiveVarSet>>,
+    storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>,
 
     // A list of suspension points, generated during the transform
     suspension_points: Vec<SuspensionPoint<'tcx>>,
@@ -418,7 +419,7 @@ struct LivenessInfo {
     /// GeneratorSavedLocal is indexed in terms of the elements in this set;
     /// i.e. GeneratorSavedLocal::new(1) corresponds to the second local
     /// included in this set.
-    live_locals: liveness::LiveVarSet,
+    live_locals: BitSet<Local>,
 
     /// The set of saved locals live at each suspension point.
     live_locals_at_suspension_points: Vec<BitSet<GeneratorSavedLocal>>,
@@ -430,7 +431,7 @@ struct LivenessInfo {
 
     /// For every suspending block, the locals which are storage-live across
     /// that suspension point.
-    storage_liveness: IndexVec<BasicBlock, Option<liveness::LiveVarSet>>,
+    storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>,
 }
 
 fn locals_live_across_suspend_points(
@@ -467,17 +468,22 @@ fn locals_live_across_suspend_points(
         dataflow::ResultsCursor::new(body_ref, &requires_storage_results);
 
     // Calculate the liveness of MIR locals ignoring borrows.
-    let mut live_locals = liveness::LiveVarSet::new_empty(body.local_decls.len());
-    let mut liveness = liveness::liveness_of_locals(body);
-    liveness::dump_mir(tcx, "generator_liveness", source, body_ref, &liveness);
+    let mut liveness = MaybeLiveLocals
+        .into_engine(tcx, body_ref, def_id)
+        .iterate_to_fixpoint()
+        .into_results_cursor(body_ref);
 
     let mut storage_liveness_map = IndexVec::from_elem(None, body.basic_blocks());
     let mut live_locals_at_suspension_points = Vec::new();
+    let mut live_locals_at_any_suspension_point = BitSet::new_empty(body.local_decls.len());
 
     for (block, data) in body.basic_blocks().iter_enumerated() {
         if let TerminatorKind::Yield { .. } = data.terminator().kind {
             let loc = Location { block, statement_index: data.statements.len() };
 
+            liveness.seek_to_block_end(block);
+            let mut live_locals = liveness.get().clone();
+
             if !movable {
                 // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
                 // This is correct for movable generators since borrows cannot live across
@@ -489,59 +495,51 @@ fn locals_live_across_suspend_points(
                 // If a borrow is converted to a raw reference, we must also assume that it lives
                 // forever. Note that the final liveness is still bounded by the storage liveness
                 // of the local, which happens using the `intersect` operation below.
-                borrowed_locals_cursor.seek_before(loc);
-                liveness.outs[block].union(borrowed_locals_cursor.get());
+                borrowed_locals_cursor.seek_before_primary_effect(loc);
+                live_locals.union(borrowed_locals_cursor.get());
             }
 
-            storage_live.seek_before(loc);
-            let mut storage_liveness = storage_live.get().clone();
-
-            // Later passes handle the generator's `self` argument separately.
-            storage_liveness.remove(SELF_ARG);
-
             // Store the storage liveness for later use so we can restore the state
             // after a suspension point
-            storage_liveness_map[block] = Some(storage_liveness);
-
-            requires_storage_cursor.seek_before(loc);
-            let storage_required = requires_storage_cursor.get().clone();
+            storage_live.seek_before_primary_effect(loc);
+            storage_liveness_map[block] = Some(storage_live.get().clone());
 
             // Locals live are live at this point only if they are used across
             // suspension points (the `liveness` variable)
             // and their storage is required (the `storage_required` variable)
-            let mut live_locals_here = storage_required;
-            live_locals_here.intersect(&liveness.outs[block]);
+            requires_storage_cursor.seek_before_primary_effect(loc);
+            live_locals.intersect(requires_storage_cursor.get());
 
             // The generator argument is ignored.
-            live_locals_here.remove(SELF_ARG);
+            live_locals.remove(SELF_ARG);
 
-            debug!("loc = {:?}, live_locals_here = {:?}", loc, live_locals_here);
+            debug!("loc = {:?}, live_locals = {:?}", loc, live_locals);
 
             // Add the locals live at this suspension point to the set of locals which live across
             // any suspension points
-            live_locals.union(&live_locals_here);
+            live_locals_at_any_suspension_point.union(&live_locals);
 
-            live_locals_at_suspension_points.push(live_locals_here);
+            live_locals_at_suspension_points.push(live_locals);
         }
     }
-    debug!("live_locals = {:?}", live_locals);
+    debug!("live_locals_anywhere = {:?}", live_locals_at_any_suspension_point);
 
     // Renumber our liveness_map bitsets to include only the locals we are
     // saving.
     let live_locals_at_suspension_points = live_locals_at_suspension_points
         .iter()
-        .map(|live_here| renumber_bitset(&live_here, &live_locals))
+        .map(|live_here| renumber_bitset(&live_here, &live_locals_at_any_suspension_point))
         .collect();
 
     let storage_conflicts = compute_storage_conflicts(
         body_ref,
-        &live_locals,
+        &live_locals_at_any_suspension_point,
         always_live_locals.clone(),
         requires_storage_results,
     );
 
     LivenessInfo {
-        live_locals,
+        live_locals: live_locals_at_any_suspension_point,
         live_locals_at_suspension_points,
         storage_conflicts,
         storage_liveness: storage_liveness_map,
@@ -555,7 +553,7 @@ fn locals_live_across_suspend_points(
 /// `[0, 1, 2]`. Thus, if `input = [3, 5]` we would return `[1, 2]`.
 fn renumber_bitset(
     input: &BitSet<Local>,
-    stored_locals: &liveness::LiveVarSet,
+    stored_locals: &BitSet<Local>,
 ) -> BitSet<GeneratorSavedLocal> {
     assert!(stored_locals.superset(&input), "{:?} not a superset of {:?}", stored_locals, input);
     let mut out = BitSet::new_empty(stored_locals.count());
@@ -575,7 +573,7 @@ fn renumber_bitset(
 /// computation; see `GeneratorLayout` for more.
 fn compute_storage_conflicts(
     body: &'mir Body<'tcx>,
-    stored_locals: &liveness::LiveVarSet,
+    stored_locals: &BitSet<Local>,
     always_live_locals: storage::AlwaysLiveLocals,
     requires_storage: dataflow::Results<'tcx, MaybeRequiresStorage<'mir, 'tcx>>,
 ) -> BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal> {
@@ -630,7 +628,7 @@ fn compute_storage_conflicts(
 
 struct StorageConflictVisitor<'mir, 'tcx, 's> {
     body: &'mir Body<'tcx>,
-    stored_locals: &'s liveness::LiveVarSet,
+    stored_locals: &'s BitSet<Local>,
     // FIXME(tmandry): Consider using sparse bitsets here once we have good
     // benchmarks for generators.
     local_conflicts: BitMatrix<Local, Local>,
@@ -639,7 +637,7 @@ struct StorageConflictVisitor<'mir, 'tcx, 's> {
 impl dataflow::ResultsVisitor<'mir, 'tcx> for StorageConflictVisitor<'mir, 'tcx, '_> {
     type FlowState = BitSet<Local>;
 
-    fn visit_statement(
+    fn visit_statement_before_primary_effect(
         &mut self,
         state: &Self::FlowState,
         _statement: &'mir Statement<'tcx>,
@@ -648,7 +646,7 @@ impl dataflow::ResultsVisitor<'mir, 'tcx> for StorageConflictVisitor<'mir, 'tcx,
         self.apply_state(state, loc);
     }
 
-    fn visit_terminator(
+    fn visit_terminator_before_primary_effect(
         &mut self,
         state: &Self::FlowState,
         _terminator: &'mir Terminator<'tcx>,
@@ -689,7 +687,7 @@ fn compute_layout<'tcx>(
 ) -> (
     FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
     GeneratorLayout<'tcx>,
-    IndexVec<BasicBlock, Option<liveness::LiveVarSet>>,
+    IndexVec<BasicBlock, Option<BitSet<Local>>>,
 ) {
     // Use a liveness analysis to compute locals which are live across a suspension point
     let LivenessInfo {
diff --git a/src/librustc_mir/transform/rustc_peek.rs b/src/librustc_mir/transform/rustc_peek.rs
index b9302d58cd5..43ddc0c914c 100644
--- a/src/librustc_mir/transform/rustc_peek.rs
+++ b/src/librustc_mir/transform/rustc_peek.rs
@@ -15,7 +15,7 @@ use crate::dataflow::MaybeMutBorrowedLocals;
 use crate::dataflow::MoveDataParamEnv;
 use crate::dataflow::{Analysis, Results, ResultsCursor};
 use crate::dataflow::{
-    DefinitelyInitializedPlaces, MaybeInitializedPlaces, MaybeUninitializedPlaces,
+    DefinitelyInitializedPlaces, MaybeInitializedPlaces, MaybeLiveLocals, MaybeUninitializedPlaces,
 };
 
 pub struct SanityCheck;
@@ -36,31 +36,45 @@ impl<'tcx> MirPass<'tcx> for SanityCheck {
         let move_data = MoveData::gather_moves(body, tcx, param_env).unwrap();
         let mdpe = MoveDataParamEnv { move_data, param_env };
 
-        let flow_inits = MaybeInitializedPlaces::new(tcx, body, &mdpe)
-            .into_engine(tcx, body, def_id)
-            .iterate_to_fixpoint();
-        let flow_uninits = MaybeUninitializedPlaces::new(tcx, body, &mdpe)
-            .into_engine(tcx, body, def_id)
-            .iterate_to_fixpoint();
-        let flow_def_inits = DefinitelyInitializedPlaces::new(tcx, body, &mdpe)
-            .into_engine(tcx, body, def_id)
-            .iterate_to_fixpoint();
-        let flow_mut_borrowed = MaybeMutBorrowedLocals::mut_borrows_only(tcx, body, param_env)
-            .into_engine(tcx, body, def_id)
-            .iterate_to_fixpoint();
-
         if has_rustc_mir_with(&attributes, sym::rustc_peek_maybe_init).is_some() {
+            let flow_inits = MaybeInitializedPlaces::new(tcx, body, &mdpe)
+                .into_engine(tcx, body, def_id)
+                .iterate_to_fixpoint();
+
             sanity_check_via_rustc_peek(tcx, body, def_id, &attributes, &flow_inits);
         }
+
         if has_rustc_mir_with(&attributes, sym::rustc_peek_maybe_uninit).is_some() {
+            let flow_uninits = MaybeUninitializedPlaces::new(tcx, body, &mdpe)
+                .into_engine(tcx, body, def_id)
+                .iterate_to_fixpoint();
+
             sanity_check_via_rustc_peek(tcx, body, def_id, &attributes, &flow_uninits);
         }
+
         if has_rustc_mir_with(&attributes, sym::rustc_peek_definite_init).is_some() {
+            let flow_def_inits = DefinitelyInitializedPlaces::new(tcx, body, &mdpe)
+                .into_engine(tcx, body, def_id)
+                .iterate_to_fixpoint();
+
             sanity_check_via_rustc_peek(tcx, body, def_id, &attributes, &flow_def_inits);
         }
+
         if has_rustc_mir_with(&attributes, sym::rustc_peek_indirectly_mutable).is_some() {
+            let flow_mut_borrowed = MaybeMutBorrowedLocals::mut_borrows_only(tcx, body, param_env)
+                .into_engine(tcx, body, def_id)
+                .iterate_to_fixpoint();
+
             sanity_check_via_rustc_peek(tcx, body, def_id, &attributes, &flow_mut_borrowed);
         }
+
+        if has_rustc_mir_with(&attributes, sym::rustc_peek_liveness).is_some() {
+            let flow_liveness =
+                MaybeLiveLocals.into_engine(tcx, body, def_id).iterate_to_fixpoint();
+
+            sanity_check_via_rustc_peek(tcx, body, def_id, &attributes, &flow_liveness);
+        }
+
         if has_rustc_mir_with(&attributes, sym::stop_after_dataflow).is_some() {
             tcx.sess.fatal("stop_after_dataflow ended compilation");
         }
@@ -126,7 +140,7 @@ pub fn sanity_check_via_rustc_peek<'tcx, A>(
                 mir::Rvalue::Use(mir::Operand::Move(place) | mir::Operand::Copy(place)),
             ) => {
                 let loc = Location { block: bb, statement_index };
-                cursor.seek_before(loc);
+                cursor.seek_before_primary_effect(loc);
                 let state = cursor.get();
                 results.analysis.peek_at(tcx, *place, state, call);
             }
@@ -286,3 +300,25 @@ impl<'tcx> RustcPeekAt<'tcx> for MaybeMutBorrowedLocals<'_, 'tcx> {
         }
     }
 }
+
+impl<'tcx> RustcPeekAt<'tcx> for MaybeLiveLocals {
+    fn peek_at(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        place: mir::Place<'tcx>,
+        flow_state: &BitSet<Local>,
+        call: PeekCall,
+    ) {
+        warn!("peek_at: place={:?}", place);
+        let local = if let Some(l) = place.as_local() {
+            l
+        } else {
+            tcx.sess.span_err(call.span, "rustc_peek: argument was not a local");
+            return;
+        };
+
+        if !flow_state.contains(local) {
+            tcx.sess.span_err(call.span, "rustc_peek: bit not set");
+        }
+    }
+}