Extend dataflow framework to support arbitrary lattices

Author: ecstatic-morse

compiler/rustc_mir/src/dataflow/framework/cursor.rs

@@ -4,6 +4,7 @@ use std::borrow::Borrow;
 use std::cmp::Ordering;
 
 use rustc_index::bit_set::BitSet;
+use rustc_index::vec::Idx;
 use rustc_middle::mir::{self, BasicBlock, Location};
 
 use super::{Analysis, Direction, Effect, EffectIndex, Results};
@@ -26,7 +27,7 @@ where
 {
     body: &'mir mir::Body<'tcx>,
     results: R,
-    state: BitSet<A::Idx>,
+    state: A::Domain,
 
     pos: CursorPosition,
 
@@ -46,17 +47,16 @@ where
 {
     /// Returns a new cursor that can inspect `results`.
     pub fn new(body: &'mir mir::Body<'tcx>, results: R) -> Self {
-        let bits_per_block = results.borrow().entry_set_for_block(mir::START_BLOCK).domain_size();
-
+        let bottom_value = results.borrow().analysis.bottom_value(body);
         ResultsCursor {
             body,
             results,
 
-            // Initialize to an empty `BitSet` and set `state_needs_reset` to tell the cursor that
+            // Initialize to the `bottom_value` and set `state_needs_reset` to tell the cursor that
             // it needs to reset to block entry before the first seek. The cursor position is
             // immaterial.
             state_needs_reset: true,
-            state: BitSet::new_empty(bits_per_block),
+            state: bottom_value,
             pos: CursorPosition::block_entry(mir::START_BLOCK),
 
             #[cfg(debug_assertions)]
@@ -79,17 +79,10 @@ where
     }
 
     /// Returns the dataflow state at the current location.
-    pub fn get(&self) -> &BitSet<A::Idx> {
+    pub fn get(&self) -> &A::Domain {
         &self.state
     }
 
-    /// Returns `true` if the dataflow state at the current location contains the given element.
-    ///
-    /// Shorthand for `self.get().contains(elem)`
-    pub fn contains(&self, elem: A::Idx) -> bool {
-        self.state.contains(elem)
-    }
-
     /// Resets the cursor to hold the entry set for the given basic block.
     ///
     /// For forward dataflow analyses, this is the dataflow state prior to the first statement.
@@ -99,7 +92,7 @@ where
         #[cfg(debug_assertions)]
         assert!(self.reachable_blocks.contains(block));
 
-        self.state.overwrite(&self.results.borrow().entry_set_for_block(block));
+        self.state.clone_from(&self.results.borrow().entry_set_for_block(block));
         self.pos = CursorPosition::block_entry(block);
         self.state_needs_reset = false;
     }
@@ -207,12 +200,23 @@ where
     ///
     /// This can be used, e.g., to apply the call return effect directly to the cursor without
     /// creating an extra copy of the dataflow state.
-    pub fn apply_custom_effect(&mut self, f: impl FnOnce(&A, &mut BitSet<A::Idx>)) {
+    pub fn apply_custom_effect(&mut self, f: impl FnOnce(&A, &mut A::Domain)) {
         f(&self.results.borrow().analysis, &mut self.state);
         self.state_needs_reset = true;
     }
 }
 
+impl<'mir, 'tcx, A, R, T> ResultsCursor<'mir, 'tcx, A, R>
+where
+    A: Analysis<'tcx, Domain = BitSet<T>>,
+    T: Idx,
+    R: Borrow<Results<'tcx, A>>,
+{
+    pub fn contains(&self, elem: T) -> bool {
+        self.get().contains(elem)
+    }
+}
+
 #[derive(Clone, Copy, Debug)]
 struct CursorPosition {
     block: BasicBlock,

compiler/rustc_mir/src/dataflow/framework/direction.rs

@@ -18,7 +18,7 @@ pub trait Direction {
     /// `effects.start()` must precede or equal `effects.end()` in this direction.
     fn apply_effects_in_range<A>(
         analysis: &A,
-        state: &mut BitSet<A::Idx>,
+        state: &mut A::Domain,
         block: BasicBlock,
         block_data: &mir::BasicBlockData<'tcx>,
         effects: RangeInclusive<EffectIndex>,
@@ -27,7 +27,7 @@ pub trait Direction {
 
     fn apply_effects_in_block<A>(
         analysis: &A,
-        state: &mut BitSet<A::Idx>,
+        state: &mut A::Domain,
         block: BasicBlock,
         block_data: &mir::BasicBlockData<'tcx>,
     ) where
@@ -55,9 +55,9 @@ pub trait Direction {
         tcx: TyCtxt<'tcx>,
         body: &mir::Body<'tcx>,
         dead_unwinds: Option<&BitSet<BasicBlock>>,
-        exit_state: &mut BitSet<A::Idx>,
+        exit_state: &mut A::Domain,
         block: (BasicBlock, &'_ mir::BasicBlockData<'tcx>),
-        propagate: impl FnMut(BasicBlock, &BitSet<A::Idx>),
+        propagate: impl FnMut(BasicBlock, &A::Domain),
     ) where
         A: Analysis<'tcx>;
 }
@@ -72,7 +72,7 @@ impl Direction for Backward {
 
     fn apply_effects_in_block<A>(
         analysis: &A,
-        state: &mut BitSet<A::Idx>,
+        state: &mut A::Domain,
         block: BasicBlock,
         block_data: &mir::BasicBlockData<'tcx>,
     ) where
@@ -112,7 +112,7 @@ impl Direction for Backward {
 
     fn apply_effects_in_range<A>(
         analysis: &A,
-        state: &mut BitSet<A::Idx>,
+        state: &mut A::Domain,
         block: BasicBlock,
         block_data: &mir::BasicBlockData<'tcx>,
         effects: RangeInclusive<EffectIndex>,
@@ -224,9 +224,9 @@ impl Direction for Backward {
         _tcx: TyCtxt<'tcx>,
         body: &mir::Body<'tcx>,
         dead_unwinds: Option<&BitSet<BasicBlock>>,
-        exit_state: &mut BitSet<A::Idx>,
+        exit_state: &mut A::Domain,
         (bb, _bb_data): (BasicBlock, &'_ mir::BasicBlockData<'tcx>),
-        mut propagate: impl FnMut(BasicBlock, &BitSet<A::Idx>),
+        mut propagate: impl FnMut(BasicBlock, &A::Domain),
     ) where
         A: Analysis<'tcx>,
     {
@@ -281,7 +281,7 @@ impl Direction for Forward {
 
     fn apply_effects_in_block<A>(
         analysis: &A,
-        state: &mut BitSet<A::Idx>,
+        state: &mut A::Domain,
         block: BasicBlock,
         block_data: &mir::BasicBlockData<'tcx>,
     ) where
@@ -321,7 +321,7 @@ impl Direction for Forward {
 
     fn apply_effects_in_range<A>(
         analysis: &A,
-        state: &mut BitSet<A::Idx>,
+        state: &mut A::Domain,
         block: BasicBlock,
         block_data: &mir::BasicBlockData<'tcx>,
         effects: RangeInclusive<EffectIndex>,
@@ -428,9 +428,9 @@ impl Direction for Forward {
         tcx: TyCtxt<'tcx>,
         body: &mir::Body<'tcx>,
         dead_unwinds: Option<&BitSet<BasicBlock>>,
-        exit_state: &mut BitSet<A::Idx>,
+        exit_state: &mut A::Domain,
         (bb, bb_data): (BasicBlock, &'_ mir::BasicBlockData<'tcx>),
-        mut propagate: impl FnMut(BasicBlock, &BitSet<A::Idx>),
+        mut propagate: impl FnMut(BasicBlock, &A::Domain),
     ) where
         A: Analysis<'tcx>,
     {
@@ -499,7 +499,7 @@ impl Direction for Forward {
                         // MIR building adds discriminants to the `values` array in the same order as they
                         // are yielded by `AdtDef::discriminants`. We rely on this to match each
                         // discriminant in `values` to its corresponding variant in linear time.
-                        let mut tmp = BitSet::new_empty(exit_state.domain_size());
+                        let mut tmp = analysis.bottom_value(body);
                         let mut discriminants = enum_def.discriminants(tcx);
                         for (value, target) in values.iter().zip(targets.iter().copied()) {
                             let (variant_idx, _) =
@@ -508,7 +508,7 @@ impl Direction for Forward {
                                          from that of `SwitchInt::values`",
                                 );
 
-                            tmp.overwrite(exit_state);
+                            tmp.clone_from(exit_state);
                             analysis.apply_discriminant_switch_effect(
                                 &mut tmp,
                                 bb,