Implement dataflow state joins through the analysis. Make JoinSemiLattice the default implementation when the join does not require extra context.

Jarcho · Jarcho · commit 0a29306bc067 · 2024-02-24T23:55:41.000-05:00
diff --git a/compiler/rustc_mir_dataflow/src/framework/direction.rs b/compiler/rustc_mir_dataflow/src/framework/direction.rs
@@ -56,7 +56,7 @@ pub trait Direction {
         exit_state: &mut A::Domain,
         block: BasicBlock,
         edges: TerminatorEdges<'_, 'tcx>,
-        propagate: impl FnMut(BasicBlock, &A::Domain),
+        propagate: impl FnMut(&mut A, BasicBlock, &A::Domain),
     ) where
         A: Analysis<'tcx>;
 }
@@ -223,7 +223,7 @@ impl Direction for Backward {
         exit_state: &mut A::Domain,
         bb: BasicBlock,
         _edges: TerminatorEdges<'_, 'tcx>,
-        mut propagate: impl FnMut(BasicBlock, &A::Domain),
+        mut propagate: impl FnMut(&mut A, BasicBlock, &A::Domain),
     ) where
         A: Analysis<'tcx>,
     {
@@ -239,7 +239,7 @@ impl Direction for Backward {
                         pred,
                         CallReturnPlaces::Call(destination),
                     );
-                    propagate(pred, &tmp);
+                    propagate(analysis, pred, &tmp);
                 }
 
                 mir::TerminatorKind::InlineAsm {
@@ -251,7 +251,7 @@ impl Direction for Backward {
                         pred,
                         CallReturnPlaces::InlineAsm(operands),
                     );
-                    propagate(pred, &tmp);
+                    propagate(analysis, pred, &tmp);
                 }
 
                 mir::TerminatorKind::Yield { resume, resume_arg, .. } if resume == bb => {
@@ -261,7 +261,7 @@ impl Direction for Backward {
                         resume,
                         CallReturnPlaces::Yield(resume_arg),
                     );
-                    propagate(pred, &tmp);
+                    propagate(analysis, pred, &tmp);
                 }
 
                 mir::TerminatorKind::SwitchInt { targets: _, ref discr } => {
@@ -277,11 +277,11 @@ impl Direction for Backward {
                     analysis.apply_switch_int_edge_effects(pred, discr, &mut applier);
 
                     if !applier.effects_applied {
-                        propagate(pred, exit_state)
+                        propagate(analysis, pred, exit_state)
                     }
                 }
 
-                _ => propagate(pred, exit_state),
+                _ => propagate(analysis, pred, exit_state),
             }
         }
     }
@@ -296,12 +296,17 @@ struct BackwardSwitchIntEdgeEffectsApplier<'mir, 'tcx, D, F> {
     effects_applied: bool,
 }
 
-impl<D, F> super::SwitchIntEdgeEffects<D> for BackwardSwitchIntEdgeEffectsApplier<'_, '_, D, F>
+impl<'tcx, A, F> super::SwitchIntEdgeEffects<'tcx, A>
+    for BackwardSwitchIntEdgeEffectsApplier<'_, '_, A::Domain, F>
 where
-    D: Clone,
-    F: FnMut(BasicBlock, &D),
+    A: Analysis<'tcx>,
+    F: FnMut(&mut A, BasicBlock, &A::Domain),
 {
-    fn apply(&mut self, mut apply_edge_effect: impl FnMut(&mut D, SwitchIntTarget)) {
+    fn apply(
+        &mut self,
+        analysis: &mut A,
+        mut apply_edge_effect: impl FnMut(&mut A, &mut A::Domain, SwitchIntTarget),
+    ) {
         assert!(!self.effects_applied);
 
         let values = &self.body.basic_blocks.switch_sources()[&(self.bb, self.pred)];
@@ -310,8 +315,8 @@ where
         let mut tmp = None;
         for target in targets {
             let tmp = opt_clone_from_or_clone(&mut tmp, self.exit_state);
-            apply_edge_effect(tmp, target);
-            (self.propagate)(self.pred, tmp);
+            apply_edge_effect(analysis, tmp, target);
+            (self.propagate)(analysis, self.pred, tmp);
         }
 
         self.effects_applied = true;
@@ -475,25 +480,25 @@ impl Direction for Forward {
         exit_state: &mut A::Domain,
         bb: BasicBlock,
         edges: TerminatorEdges<'_, 'tcx>,
-        mut propagate: impl FnMut(BasicBlock, &A::Domain),
+        mut propagate: impl FnMut(&mut A, BasicBlock, &A::Domain),
     ) where
         A: Analysis<'tcx>,
     {
         match edges {
             TerminatorEdges::None => {}
-            TerminatorEdges::Single(target) => propagate(target, exit_state),
+            TerminatorEdges::Single(target) => propagate(analysis, target, exit_state),
             TerminatorEdges::Double(target, unwind) => {
-                propagate(target, exit_state);
-                propagate(unwind, exit_state);
+                propagate(analysis, target, exit_state);
+                propagate(analysis, unwind, exit_state);
             }
             TerminatorEdges::AssignOnReturn { return_, cleanup, place } => {
                 // This must be done *first*, otherwise the unwind path will see the assignments.
                 if let Some(cleanup) = cleanup {
-                    propagate(cleanup, exit_state);
+                    propagate(analysis, cleanup, exit_state);
                 }
                 if let Some(return_) = return_ {
                     analysis.apply_call_return_effect(exit_state, bb, place);
-                    propagate(return_, exit_state);
+                    propagate(analysis, return_, exit_state);
                 }
             }
             TerminatorEdges::SwitchInt { targets, discr } => {
@@ -515,7 +520,7 @@ impl Direction for Forward {
 
                 if !effects_applied {
                     for target in targets.all_targets() {
-                        propagate(*target, exit_state);
+                        propagate(analysis, *target, exit_state);
                     }
                 }
             }
@@ -531,26 +536,35 @@ struct ForwardSwitchIntEdgeEffectsApplier<'mir, D, F> {
     effects_applied: bool,
 }
 
-impl<D, F> super::SwitchIntEdgeEffects<D> for ForwardSwitchIntEdgeEffectsApplier<'_, D, F>
+impl<'tcx, A, F> super::SwitchIntEdgeEffects<'tcx, A>
+    for ForwardSwitchIntEdgeEffectsApplier<'_, A::Domain, F>
 where
-    D: Clone,
-    F: FnMut(BasicBlock, &D),
+    A: Analysis<'tcx>,
+    F: FnMut(&mut A, BasicBlock, &A::Domain),
 {
-    fn apply(&mut self, mut apply_edge_effect: impl FnMut(&mut D, SwitchIntTarget)) {
+    fn apply(
+        &mut self,
+        analysis: &mut A,
+        mut apply_edge_effect: impl FnMut(&mut A, &mut A::Domain, SwitchIntTarget),
+    ) {
         assert!(!self.effects_applied);
 
         let mut tmp = None;
         for (value, target) in self.targets.iter() {
             let tmp = opt_clone_from_or_clone(&mut tmp, self.exit_state);
-            apply_edge_effect(tmp, SwitchIntTarget { value: Some(value), target });
-            (self.propagate)(target, tmp);
+            apply_edge_effect(analysis, tmp, SwitchIntTarget { value: Some(value), target });
+            (self.propagate)(analysis, target, tmp);
         }
 
         // Once we get to the final, "otherwise" branch, there is no need to preserve `exit_state`,
         // so pass it directly to `apply_edge_effect` to save a clone of the dataflow state.
         let otherwise = self.targets.otherwise();
-        apply_edge_effect(self.exit_state, SwitchIntTarget { value: None, target: otherwise });
-        (self.propagate)(otherwise, self.exit_state);
+        apply_edge_effect(
+            analysis,
+            self.exit_state,
+            SwitchIntTarget { value: None, target: otherwise },
+        );
+        (self.propagate)(analysis, otherwise, self.exit_state);
 
         self.effects_applied = true;
     }
diff --git a/compiler/rustc_mir_dataflow/src/framework/engine.rs b/compiler/rustc_mir_dataflow/src/framework/engine.rs
@@ -23,7 +23,7 @@ use super::fmt::DebugWithContext;
 use super::graphviz;
 use super::{
     visit_results, Analysis, AnalysisDomain, Direction, GenKill, GenKillAnalysis, GenKillSet,
-    JoinSemiLattice, ResultsCursor, ResultsVisitor,
+    ResultsCursor, ResultsVisitor,
 };
 
 pub type EntrySets<'tcx, A> = IndexVec<BasicBlock, <A as AnalysisDomain<'tcx>>::Domain>;
@@ -97,7 +97,7 @@ where
 impl<'mir, 'tcx, A, D, T> Engine<'mir, 'tcx, A>
 where
     A: GenKillAnalysis<'tcx, Idx = T, Domain = D>,
-    D: Clone + JoinSemiLattice + GenKill<T> + BitSetExt<T>,
+    D: Clone + Eq + GenKill<T> + BitSetExt<T>,
     T: Idx,
 {
     /// Creates a new `Engine` to solve a gen-kill dataflow problem.
@@ -136,7 +136,7 @@ where
 impl<'mir, 'tcx, A, D> Engine<'mir, 'tcx, A>
 where
     A: Analysis<'tcx, Domain = D>,
-    D: Clone + JoinSemiLattice,
+    D: Clone + Eq,
 {
     /// Creates a new `Engine` to solve a dataflow problem with an arbitrary transfer
     /// function.
@@ -229,8 +229,8 @@ where
                 &mut state,
                 bb,
                 edges,
-                |target: BasicBlock, state: &A::Domain| {
-                    let set_changed = entry_sets[target].join(state);
+                |analysis: &mut A, target: BasicBlock, state: &A::Domain| {
+                    let set_changed = analysis.join(&mut entry_sets[target], state, target);
                     if set_changed {
                         dirty_queue.insert(target);
                     }
diff --git a/compiler/rustc_mir_dataflow/src/framework/mod.rs b/compiler/rustc_mir_dataflow/src/framework/mod.rs
@@ -93,7 +93,7 @@ impl<T: Idx> BitSetExt<T> for ChunkedBitSet<T> {
 /// initial value at the entry point of each basic block.
 pub trait AnalysisDomain<'tcx> {
     /// The type that holds the dataflow state at any given point in the program.
-    type Domain: Clone + JoinSemiLattice;
+    type Domain: Clone + Eq;
 
     /// The direction of this analysis. Either `Forward` or `Backward`.
     type Direction: Direction = Forward;
@@ -118,6 +118,24 @@ pub trait AnalysisDomain<'tcx> {
     fn initialize_start_block(&self, body: &mir::Body<'tcx>, state: &mut Self::Domain);
 }
 
+/// The join operation of a dataflow analysis. A default implementation is used when the
+/// `Domain` type implements `JoinSemiLattice`.
+pub trait AnalysisJoin<'tcx>: AnalysisDomain<'tcx> {
+    /// Computes the least upper bound of two elements, storing the result in `state` and returning
+    /// `true` if `state` has changed.
+    fn join(&mut self, state: &mut Self::Domain, other: &Self::Domain, loc: BasicBlock) -> bool;
+}
+
+impl<'tcx, T> AnalysisJoin<'tcx> for T
+where
+    T: ?Sized + AnalysisDomain<'tcx>,
+    <T as AnalysisDomain<'tcx>>::Domain: JoinSemiLattice,
+{
+    fn join(&mut self, state: &mut Self::Domain, other: &Self::Domain, _: BasicBlock) -> bool {
+        state.join(other)
+    }
+}
+
 /// A dataflow problem with an arbitrarily complex transfer function.
 ///
 /// # Convergence
@@ -134,7 +152,7 @@ pub trait AnalysisDomain<'tcx> {
 /// monotonically until fixpoint is reached. Note that this monotonicity requirement only applies
 /// to the same point in the program at different points in time. The dataflow state at a given
 /// point in the program may or may not be greater than the state at any preceding point.
-pub trait Analysis<'tcx>: AnalysisDomain<'tcx> {
+pub trait Analysis<'tcx>: AnalysisJoin<'tcx> + Sized {
     /// Updates the current dataflow state with the effect of evaluating a statement.
     fn apply_statement_effect(
         &mut self,
@@ -215,7 +233,7 @@ pub trait Analysis<'tcx>: AnalysisDomain<'tcx> {
         &mut self,
         _block: BasicBlock,
         _discr: &mir::Operand<'tcx>,
-        _apply_edge_effects: &mut impl SwitchIntEdgeEffects<Self::Domain>,
+        _apply_edge_effects: &mut impl SwitchIntEdgeEffects<'tcx, Self>,
     ) {
     }
 
@@ -238,10 +256,7 @@ pub trait Analysis<'tcx>: AnalysisDomain<'tcx> {
         self,
         tcx: TyCtxt<'tcx>,
         body: &'mir mir::Body<'tcx>,
-    ) -> Engine<'mir, 'tcx, Self>
-    where
-        Self: Sized,
-    {
+    ) -> Engine<'mir, 'tcx, Self> {
         Engine::new_generic(tcx, body, self)
     }
 }
@@ -306,11 +321,11 @@ pub trait GenKillAnalysis<'tcx>: Analysis<'tcx> {
     );
 
     /// See `Analysis::apply_switch_int_edge_effects`.
-    fn switch_int_edge_effects<G: GenKill<Self::Idx>>(
+    fn switch_int_edge_effects(
         &mut self,
         _block: BasicBlock,
         _discr: &mir::Operand<'tcx>,
-        _edge_effects: &mut impl SwitchIntEdgeEffects<G>,
+        _edge_effects: &mut impl SwitchIntEdgeEffects<'tcx, Self>,
     ) {
     }
 }
@@ -372,7 +387,7 @@ where
         &mut self,
         block: BasicBlock,
         discr: &mir::Operand<'tcx>,
-        edge_effects: &mut impl SwitchIntEdgeEffects<A::Domain>,
+        edge_effects: &mut impl SwitchIntEdgeEffects<'tcx, A>,
     ) {
         self.switch_int_edge_effects(block, discr, edge_effects);
     }
@@ -383,10 +398,7 @@ where
         self,
         tcx: TyCtxt<'tcx>,
         body: &'mir mir::Body<'tcx>,
-    ) -> Engine<'mir, 'tcx, Self>
-    where
-        Self: Sized,
-    {
+    ) -> Engine<'mir, 'tcx, Self> {
         Engine::new_gen_kill(tcx, body, self)
     }
 }
@@ -573,10 +585,14 @@ pub struct SwitchIntTarget {
 }
 
 /// A type that records the edge-specific effects for a `SwitchInt` terminator.
-pub trait SwitchIntEdgeEffects<D> {
+pub trait SwitchIntEdgeEffects<'tcx, A: AnalysisDomain<'tcx>> {
     /// Calls `apply_edge_effect` for each outgoing edge from a `SwitchInt` terminator and
     /// records the results.
-    fn apply(&mut self, apply_edge_effect: impl FnMut(&mut D, SwitchIntTarget));
+    fn apply(
+        &mut self,
+        analysis: &mut A,
+        apply_edge_effect: impl FnMut(&mut A, &mut A::Domain, SwitchIntTarget),
+    );
 }
 
 #[cfg(test)]
diff --git a/compiler/rustc_mir_dataflow/src/impls/initialized.rs b/compiler/rustc_mir_dataflow/src/impls/initialized.rs
@@ -396,11 +396,11 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
         });
     }
 
-    fn switch_int_edge_effects<G: GenKill<Self::Idx>>(
+    fn switch_int_edge_effects(
         &mut self,
         block: mir::BasicBlock,
         discr: &mir::Operand<'tcx>,
-        edge_effects: &mut impl SwitchIntEdgeEffects<G>,
+        edge_effects: &mut impl SwitchIntEdgeEffects<'tcx, Self>,
     ) {
         if !self.tcx.sess.opts.unstable_opts.precise_enum_drop_elaboration {
             return;
@@ -415,7 +415,7 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
         };
 
         let mut discriminants = enum_def.discriminants(self.tcx);
-        edge_effects.apply(|trans, edge| {
+        edge_effects.apply(self, |self_, trans, edge| {
             let Some(value) = edge.value else {
                 return;
             };
@@ -430,7 +430,7 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
             // Kill all move paths that correspond to variants we know to be inactive along this
             // particular outgoing edge of a `SwitchInt`.
             drop_flag_effects::on_all_inactive_variants(
-                self.move_data(),
+                self_.move_data(),
                 enum_place,
                 variant,
                 |mpi| trans.kill(mpi),
@@ -521,11 +521,11 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
         });
     }
 
-    fn switch_int_edge_effects<G: GenKill<Self::Idx>>(
+    fn switch_int_edge_effects(
         &mut self,
         block: mir::BasicBlock,
         discr: &mir::Operand<'tcx>,
-        edge_effects: &mut impl SwitchIntEdgeEffects<G>,
+        edge_effects: &mut impl SwitchIntEdgeEffects<'tcx, Self>,
     ) {
         if !self.tcx.sess.opts.unstable_opts.precise_enum_drop_elaboration {
             return;
@@ -544,7 +544,7 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
         };
 
         let mut discriminants = enum_def.discriminants(self.tcx);
-        edge_effects.apply(|trans, edge| {
+        edge_effects.apply(self, |self_, trans, edge| {
             let Some(value) = edge.value else {
                 return;
             };
@@ -559,7 +559,7 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
             // Mark all move paths that correspond to variants other than this one as maybe
             // uninitialized (in reality, they are *definitely* uninitialized).
             drop_flag_effects::on_all_inactive_variants(
-                self.move_data(),
+                self_.move_data(),
                 enum_place,
                 variant,
                 |mpi| trans.gen(mpi),
diff --git a/compiler/rustc_mir_dataflow/src/value_analysis.rs b/compiler/rustc_mir_dataflow/src/value_analysis.rs
