Rollup merge of #156169 - nnethercote:simplify-SwitchInt-stuff, r=cjgillot

matthiaskrgr · web-flow · commit 02f733bb4882 · 2026-05-10T03:17:00.000+02:00
Change `SwitchInt` handling in dataflow analysis. We call `get_switch_int_data` once for the switch and then pass that data to `apply_switch_int_edge_effect` for each switch target. The only case in practice is `MaybePlacesSwitchIntData` which does an awkward thing, maintaining an index into the discriminants and updating it on each call to `apply_switch_int_edge_effect`. This commit changes things to do more work up front in `get_switch_int_data`, in order to then do less work in `apply_switch_int_edge_effect`. This avoids the need for the `variants` and `next_discr` methods and the discriminants index. Overall it's a little simpler. r? @cjgillot
diff --git a/compiler/rustc_middle/src/mir/basic_blocks.rs b/compiler/rustc_middle/src/mir/basic_blocks.rs
@@ -21,14 +21,6 @@ pub struct BasicBlocks<'tcx> {
 // Typically 95%+ of basic blocks have 4 or fewer predecessors.
 type Predecessors = IndexVec<BasicBlock, SmallVec<[BasicBlock; 4]>>;
 
-#[derive(Debug, Clone, Copy)]
-pub enum SwitchTargetValue {
-    // A normal switch value.
-    Normal(u128),
-    // The final "otherwise" fallback value.
-    Otherwise,
-}
-
 #[derive(Clone, Default, Debug)]
 struct Cache {
     predecessors: OnceLock<Predecessors>,
diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
@@ -7,7 +7,7 @@ use std::fmt::{self, Debug, Formatter};
 use std::iter;
 use std::ops::{Index, IndexMut};
 
-pub use basic_blocks::{BasicBlocks, SwitchTargetValue};
+pub use basic_blocks::BasicBlocks;
 use either::Either;
 use polonius_engine::Atom;
 use rustc_abi::{FieldIdx, VariantIdx};
diff --git a/compiler/rustc_mir_dataflow/src/framework/direction.rs b/compiler/rustc_mir_dataflow/src/framework/direction.rs
@@ -1,12 +1,10 @@
 use std::ops::RangeInclusive;
 
 use rustc_middle::bug;
-use rustc_middle::mir::{
-    self, BasicBlock, CallReturnPlaces, Location, SwitchTargetValue, TerminatorEdges,
-};
+use rustc_middle::mir::{self, BasicBlock, CallReturnPlaces, Location, TerminatorEdges};
 
 use super::visitor::ResultsVisitor;
-use super::{Analysis, Effect, EffectIndex};
+use super::{Analysis, Effect, EffectIndex, SwitchTargetIndex};
 
 pub trait Direction {
     const IS_FORWARD: bool;
@@ -113,8 +111,8 @@ impl Direction for Backward {
                     propagate(pred, &tmp);
                 }
 
-                mir::TerminatorKind::SwitchInt { ref discr, .. } => {
-                    if let Some(_data) = analysis.get_switch_int_data(pred, discr) {
+                mir::TerminatorKind::SwitchInt { ref targets, ref discr } => {
+                    if let Some(_data) = analysis.get_switch_int_data(pred, targets, discr) {
                         bug!(
                             "SwitchInt edge effects are unsupported in backward dataflow analyses"
                         );
@@ -283,23 +281,22 @@ impl Direction for Forward {
                 }
             }
             TerminatorEdges::SwitchInt { targets, discr } => {
-                if let Some(mut data) = analysis.get_switch_int_data(block, discr) {
+                if let Some(mut data) = analysis.get_switch_int_data(block, targets, discr) {
                     let mut tmp = analysis.bottom_value(body);
-                    for (value, target) in targets.iter() {
+                    for (i, (_value, target)) in targets.iter().enumerate() {
                         tmp.clone_from(exit_state);
-                        let value = SwitchTargetValue::Normal(value);
-                        analysis.apply_switch_int_edge_effect(&mut data, &mut tmp, value, targets);
+                        let target_idx = SwitchTargetIndex::Normal(i);
+                        analysis.apply_switch_int_edge_effect(&mut tmp, &mut data, target_idx);
                         propagate(target, &tmp);
                     }
 
                     // Once we get to the final, "otherwise" branch, there is no need to preserve
                     // `exit_state`, so pass it directly to `apply_switch_int_edge_effect` to save
                     // a clone of the dataflow state.
                     analysis.apply_switch_int_edge_effect(
-                        &mut data,
                         exit_state,
-                        SwitchTargetValue::Otherwise,
-                        targets,
+                        &mut data,
+                        SwitchTargetIndex::Otherwise,
                     );
                     propagate(targets.otherwise(), exit_state);
                 } else {
diff --git a/compiler/rustc_mir_dataflow/src/framework/mod.rs b/compiler/rustc_mir_dataflow/src/framework/mod.rs
@@ -38,9 +38,7 @@ use rustc_data_structures::work_queue::WorkQueue;
 use rustc_index::bit_set::{DenseBitSet, MixedBitSet};
 use rustc_index::{Idx, IndexVec};
 use rustc_middle::bug;
-use rustc_middle::mir::{
-    self, BasicBlock, CallReturnPlaces, Location, SwitchTargetValue, TerminatorEdges, traversal,
-};
+use rustc_middle::mir::{self, BasicBlock, CallReturnPlaces, Location, TerminatorEdges, traversal};
 use rustc_middle::ty::TyCtxt;
 use tracing::error;
 
@@ -212,6 +210,7 @@ pub trait Analysis<'tcx> {
     fn get_switch_int_data(
         &self,
         _block: mir::BasicBlock,
+        _targets: &mir::SwitchTargets,
         _discr: &mir::Operand<'tcx>,
     ) -> Option<Self::SwitchIntData> {
         None
@@ -220,10 +219,9 @@ pub trait Analysis<'tcx> {
     /// See comments on `get_switch_int_data`.
     fn apply_switch_int_edge_effect(
         &self,
-        _data: &mut Self::SwitchIntData,
         _state: &mut Self::Domain,
-        _value: SwitchTargetValue,
-        _targets: &mir::SwitchTargets,
+        _data: &mut Self::SwitchIntData,
+        _target_idx: SwitchTargetIndex,
     ) {
         unreachable!();
     }
@@ -313,6 +311,14 @@ pub trait Analysis<'tcx> {
     }
 }
 
+#[derive(Debug, Clone, Copy)]
+pub enum SwitchTargetIndex {
+    // Index of a normal switch target.
+    Normal(usize),
+    // The final "otherwise" fallback target.
+    Otherwise,
+}
+
 /// The legal operations for a transfer function in a gen/kill problem.
 pub trait GenKill<T> {
     /// Inserts `elem` into the state vector.
diff --git a/compiler/rustc_mir_dataflow/src/impls/initialized.rs b/compiler/rustc_mir_dataflow/src/impls/initialized.rs
@@ -4,55 +4,38 @@ use rustc_abi::VariantIdx;
 use rustc_index::Idx;
 use rustc_index::bit_set::{DenseBitSet, MixedBitSet};
 use rustc_middle::bug;
-use rustc_middle::mir::{
-    self, Body, CallReturnPlaces, Location, SwitchTargetValue, TerminatorEdges,
-};
-use rustc_middle::ty::util::Discr;
+use rustc_middle::mir::{self, Body, CallReturnPlaces, Location, TerminatorEdges};
 use rustc_middle::ty::{self, TyCtxt};
 use smallvec::SmallVec;
 use tracing::{debug, instrument};
 
 use crate::drop_flag_effects::{DropFlagState, InactiveVariants};
 use crate::move_paths::{HasMoveData, InitIndex, InitKind, LookupResult, MoveData, MovePathIndex};
 use crate::{
-    Analysis, GenKill, MaybeReachable, drop_flag_effects, drop_flag_effects_for_function_entry,
-    drop_flag_effects_for_location, on_all_children_bits, on_lookup_result_bits,
+    Analysis, GenKill, MaybeReachable, SwitchTargetIndex, drop_flag_effects,
+    drop_flag_effects_for_function_entry, drop_flag_effects_for_location, on_all_children_bits,
+    on_lookup_result_bits,
 };
 
 // Used by both `MaybeInitializedPlaces` and `MaybeUninitializedPlaces`.
 pub struct MaybePlacesSwitchIntData<'tcx> {
     enum_place: mir::Place<'tcx>,
-    discriminants: Vec<(VariantIdx, Discr<'tcx>)>,
-    index: usize,
-}
 
-impl<'tcx> MaybePlacesSwitchIntData<'tcx> {
-    /// Creates a `SmallVec` mapping each target in `targets` to its `VariantIdx`.
-    fn variants(&mut self, targets: &mir::SwitchTargets) -> SmallVec<[VariantIdx; 4]> {
-        self.index = 0;
-        targets.all_values().iter().map(|value| self.next_discr(value.get())).collect()
-    }
-
-    // The discriminant order in the `SwitchInt` targets should match the order yielded by
-    // `AdtDef::discriminants`. We rely on this to match each discriminant in the targets to its
-    // corresponding variant in linear time.
-    fn next_discr(&mut self, value: u128) -> VariantIdx {
-        // An out-of-bounds abort will occur if the discriminant ordering isn't as described above.
-        loop {
-            let (variant, discr) = self.discriminants[self.index];
-            self.index += 1;
-            if discr.val == value {
-                return variant;
-            }
-        }
-    }
+    // Variant indices targeted by the SwitchInt. For example, if you have:
+    // ```
+    // enum E { A = 1, B = 3, C = 5, D = 7 }
+    // ```
+    // and a `SwitchInt(A -> bb1, C -> bb2, _ -> bb3)`, this vec will contain `[0, 2]` because
+    // those are the variant indices for `A` and `C`.
+    variants: SmallVec<[VariantIdx; 4]>,
 }
 
 impl<'tcx> MaybePlacesSwitchIntData<'tcx> {
     fn new(
         tcx: TyCtxt<'tcx>,
         body: &Body<'tcx>,
         block: mir::BasicBlock,
+        targets: &mir::SwitchTargets,
         discr: &mir::Operand<'tcx>,
     ) -> Option<Self> {
         let Some(discr) = discr.place() else { return None };
@@ -76,11 +59,29 @@ impl<'tcx> MaybePlacesSwitchIntData<'tcx> {
                 {
                     match enum_place.ty(body, tcx).ty.kind() {
                         ty::Adt(enum_def, _) => {
-                            return Some(MaybePlacesSwitchIntData {
-                                enum_place,
-                                discriminants: enum_def.discriminants(tcx).collect(),
-                                index: 0,
-                            });
+                            // The value of each discriminant, in AdtDef order.
+                            let discriminant_vals: SmallVec<[u128; 4]> =
+                                enum_def.discriminants(tcx).map(|(_, discr)| discr.val).collect();
+                            let mut i = 0;
+
+                            // For each value in the SwitchInt, find the VariantIdx for the variant
+                            // with that value. This works because `discriminant_vals` and
+                            // `targets.all_values()` are guaranteed to list variants in the same
+                            // order. (If that ever changes we will get out-of-bounds panics here.)
+                            let variants = targets
+                                .all_values()
+                                .iter()
+                                .map(|value| {
+                                    loop {
+                                        if discriminant_vals[i] == value.get() {
+                                            return VariantIdx::new(i);
+                                        }
+                                        i += 1;
+                                    }
+                                })
+                                .collect();
+
+                            return Some(MaybePlacesSwitchIntData { enum_place, variants });
                         }
 
                         // `Rvalue::Discriminant` is also used to get the active yield point for a
@@ -452,26 +453,28 @@ impl<'tcx> Analysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
     fn get_switch_int_data(
         &self,
         block: mir::BasicBlock,
+        targets: &mir::SwitchTargets,
         discr: &mir::Operand<'tcx>,
     ) -> Option<Self::SwitchIntData> {
         if !self.tcx.sess.opts.unstable_opts.precise_enum_drop_elaboration {
             return None;
         }
 
-        MaybePlacesSwitchIntData::new(self.tcx, self.body, block, discr)
+        MaybePlacesSwitchIntData::new(self.tcx, self.body, block, targets, discr)
     }
 
     fn apply_switch_int_edge_effect(
         &self,
-        data: &mut Self::SwitchIntData,
         state: &mut Self::Domain,
-        value: SwitchTargetValue,
-        targets: &mir::SwitchTargets,
+        data: &mut Self::SwitchIntData,
+        target_idx: SwitchTargetIndex,
     ) {
-        let inactive_variants = match value {
-            SwitchTargetValue::Normal(value) => InactiveVariants::Active(data.next_discr(value)),
-            SwitchTargetValue::Otherwise if self.exclude_inactive_in_otherwise => {
-                InactiveVariants::Inactives(data.variants(targets))
+        let inactive_variants = match target_idx {
+            SwitchTargetIndex::Normal(target_idx) => {
+                InactiveVariants::Active(data.variants[target_idx])
+            }
+            SwitchTargetIndex::Otherwise if self.exclude_inactive_in_otherwise => {
+                InactiveVariants::Inactives(data.variants.clone())
             }
             _ => return,
         };
@@ -568,6 +571,7 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
     fn get_switch_int_data(
         &self,
         block: mir::BasicBlock,
+        targets: &mir::SwitchTargets,
         discr: &mir::Operand<'tcx>,
     ) -> Option<Self::SwitchIntData> {
         if !self.tcx.sess.opts.unstable_opts.precise_enum_drop_elaboration {
@@ -578,20 +582,21 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
             return None;
         }
 
-        MaybePlacesSwitchIntData::new(self.tcx, self.body, block, discr)
+        MaybePlacesSwitchIntData::new(self.tcx, self.body, block, targets, discr)
     }
 
     fn apply_switch_int_edge_effect(
         &self,
-        data: &mut Self::SwitchIntData,
         state: &mut Self::Domain,
-        value: SwitchTargetValue,
-        targets: &mir::SwitchTargets,
+        data: &mut Self::SwitchIntData,
+        target_idx: SwitchTargetIndex,
     ) {
-        let inactive_variants = match value {
-            SwitchTargetValue::Normal(value) => InactiveVariants::Active(data.next_discr(value)),
-            SwitchTargetValue::Otherwise if self.include_inactive_in_otherwise => {
-                InactiveVariants::Inactives(data.variants(targets))
+        let inactive_variants = match target_idx {
+            SwitchTargetIndex::Normal(target_idx) => {
+                InactiveVariants::Active(data.variants[target_idx])
+            }
+            SwitchTargetIndex::Otherwise if self.include_inactive_in_otherwise => {
+                InactiveVariants::Inactives(data.variants.clone())
             }
             _ => return,
         };
diff --git a/compiler/rustc_mir_dataflow/src/lib.rs b/compiler/rustc_mir_dataflow/src/lib.rs
@@ -17,8 +17,8 @@ pub use self::drop_flag_effects::{
 };
 pub use self::framework::{
     Analysis, Backward, Direction, EntryStates, Forward, GenKill, JoinSemiLattice, MaybeReachable,
-    Results, ResultsCursor, ResultsVisitor, fmt, graphviz, lattice, visit_reachable_results,
-    visit_results,
+    Results, ResultsCursor, ResultsVisitor, SwitchTargetIndex, fmt, graphviz, lattice,
+    visit_reachable_results, visit_results,
 };
 use self::move_paths::MoveData;
 
