Use TerminatorEdge for dataflow-const-prop.

2 files changed, 49 insertions, 63 deletions

diff --git a/compiler/rustc_mir_dataflow/src/value_analysis.rs b/compiler/rustc_mir_dataflow/src/value_analysis.rs
index bda5b352041..f9ab95f6162 100644
--- a/compiler/rustc_mir_dataflow/src/value_analysis.rs
+++ b/compiler/rustc_mir_dataflow/src/value_analysis.rs
@@ -241,11 +241,19 @@ pub trait ValueAnalysis<'tcx> {
 
     /// The effect of a successful function call return should not be
     /// applied here, see [`Analysis::apply_terminator_effect`].
-    fn handle_terminator(&self, terminator: &Terminator<'tcx>, state: &mut State<Self::Value>) {
+    fn handle_terminator<'mir>(
+        &self,
+        terminator: &'mir Terminator<'tcx>,
+        state: &mut State<Self::Value>,
+    ) -> TerminatorEdge<'mir, 'tcx> {
         self.super_terminator(terminator, state)
     }
 
-    fn super_terminator(&self, terminator: &Terminator<'tcx>, state: &mut State<Self::Value>) {
+    fn super_terminator<'mir>(
+        &self,
+        terminator: &'mir Terminator<'tcx>,
+        state: &mut State<Self::Value>,
+    ) -> TerminatorEdge<'mir, 'tcx> {
         match &terminator.kind {
             TerminatorKind::Call { .. } | TerminatorKind::InlineAsm { .. } => {
                 // Effect is applied by `handle_call_return`.
@@ -257,8 +265,10 @@ pub trait ValueAnalysis<'tcx> {
                 // They would have an effect, but are not allowed in this phase.
                 bug!("encountered disallowed terminator");
             }
+            TerminatorKind::SwitchInt { discr, targets } => {
+                return self.handle_switch_int(discr, targets, state);
+            }
             TerminatorKind::Goto { .. }
-            | TerminatorKind::SwitchInt { .. }
             | TerminatorKind::Resume
             | TerminatorKind::Terminate
             | TerminatorKind::Return
@@ -270,6 +280,7 @@ pub trait ValueAnalysis<'tcx> {
                 // These terminators have no effect on the analysis.
             }
         }
+        terminator.edges()
     }
 
     fn handle_call_return(
@@ -290,19 +301,22 @@ pub trait ValueAnalysis<'tcx> {
         })
     }
 
-    fn handle_switch_int(
+    fn handle_switch_int<'mir>(
         &self,
-        discr: &Operand<'tcx>,
-        apply_edge_effects: &mut impl SwitchIntEdgeEffects<State<Self::Value>>,
-    ) {
-        self.super_switch_int(discr, apply_edge_effects)
+        discr: &'mir Operand<'tcx>,
+        targets: &'mir SwitchTargets,
+        state: &mut State<Self::Value>,
+    ) -> TerminatorEdge<'mir, 'tcx> {
+        self.super_switch_int(discr, targets, state)
     }
 
-    fn super_switch_int(
+    fn super_switch_int<'mir>(
         &self,
-        _discr: &Operand<'tcx>,
-        _apply_edge_effects: &mut impl SwitchIntEdgeEffects<State<Self::Value>>,
-    ) {
+        discr: &'mir Operand<'tcx>,
+        targets: &'mir SwitchTargets,
+        _state: &mut State<Self::Value>,
+    ) -> TerminatorEdge<'mir, 'tcx> {
+        TerminatorEdge::SwitchInt { discr, targets }
     }
 
     fn wrap(self) -> ValueAnalysisWrapper<Self>
@@ -359,9 +373,10 @@ where
         _location: Location,
     ) -> TerminatorEdge<'mir, 'tcx> {
         if state.is_reachable() {
-            self.0.handle_terminator(terminator, state);
+            self.0.handle_terminator(terminator, state)
+        } else {
+            TerminatorEdge::None
         }
-        terminator.edges()
     }
 
     fn apply_call_return_effect(
@@ -378,11 +393,9 @@ where
     fn apply_switch_int_edge_effects(
         &mut self,
         _block: BasicBlock,
-        discr: &Operand<'tcx>,
-        apply_edge_effects: &mut impl SwitchIntEdgeEffects<Self::Domain>,
+        _discr: &Operand<'tcx>,
+        _apply_edge_effects: &mut impl SwitchIntEdgeEffects<Self::Domain>,
     ) {
-        // FIXME: Dataflow framework provides no access to current state here.
-        self.0.handle_switch_int(discr, apply_edge_effects)
     }
 }
 
diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
index 7d7588fcaec..08a4619591a 100644
--- a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
+++ b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
@@ -13,9 +13,7 @@ use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_mir_dataflow::value_analysis::{
     Map, State, TrackElem, ValueAnalysis, ValueAnalysisWrapper, ValueOrPlace,
 };
-use rustc_mir_dataflow::{
-    lattice::FlatSet, Analysis, Results, ResultsVisitor, SwitchIntEdgeEffects,
-};
+use rustc_mir_dataflow::{lattice::FlatSet, Analysis, Results, ResultsVisitor};
 use rustc_span::DUMMY_SP;
 use rustc_target::abi::{Align, FieldIdx, VariantIdx};
 
@@ -249,49 +247,24 @@ impl<'tcx> ValueAnalysis<'tcx> for ConstAnalysis<'_, 'tcx> {
             .unwrap_or(FlatSet::Top)
     }
 
-    fn handle_switch_int(
+    fn handle_switch_int<'mir>(
         &self,
-        discr: &Operand<'tcx>,
-        apply_edge_effects: &mut impl SwitchIntEdgeEffects<State<Self::Value>>,
-    ) {
-        // FIXME: The dataflow framework only provides the state if we call `apply()`, which makes
-        // this more inefficient than it has to be.
-        let mut discr_value = None;
-        let mut handled = false;
-        apply_edge_effects.apply(|state, target| {
-            let discr_value = match discr_value {
-                Some(value) => value,
-                None => {
-                    let value = match self.handle_operand(discr, state) {
-                        ValueOrPlace::Value(value) => value,
-                        ValueOrPlace::Place(place) => state.get_idx(place, self.map()),
-                    };
-                    let result = match value {
-                        FlatSet::Top => FlatSet::Top,
-                        FlatSet::Elem(ScalarTy(scalar, _)) => {
-                            let int = scalar.assert_int();
-                            FlatSet::Elem(int.assert_bits(int.size()))
-                        }
-                        FlatSet::Bottom => FlatSet::Bottom,
-                    };
-                    discr_value = Some(result);
-                    result
-                }
-            };
-
-            let FlatSet::Elem(choice) = discr_value else {
-                // Do nothing if we don't know which branch will be taken.
-                return;
-            };
-
-            if target.value.map(|n| n == choice).unwrap_or(!handled) {
-                // Branch is taken. Has no effect on state.
-                handled = true;
-            } else {
-                // Branch is not taken.
-                state.mark_unreachable();
-            }
-        })
+        discr: &'mir Operand<'tcx>,
+        targets: &'mir SwitchTargets,
+        state: &mut State<Self::Value>,
+    ) -> TerminatorEdge<'mir, 'tcx> {
+        let value = match self.handle_operand(discr, state) {
+            ValueOrPlace::Value(value) => value,
+            ValueOrPlace::Place(place) => state.get_idx(place, self.map()),
+        };
+        let FlatSet::Elem(ScalarTy(scalar, _)) = value else {
+            // Do nothing if we don't know which branch will be taken.
+            return TerminatorEdge::SwitchInt { discr, targets };
+        };
+
+        let int = scalar.assert_int();
+        let choice = int.assert_bits(int.size());
+        TerminatorEdge::Single(targets.target_for_value(choice))
     }
 }