Collect VariantIdx in MaybePlacesSwitchIntData

ashivaram23 · ashivaram23 · commit 04c859285b9a · 2025-07-12T13:27:52.000-05:00
diff --git a/compiler/rustc_middle/src/mir/basic_blocks.rs b/compiler/rustc_middle/src/mir/basic_blocks.rs
@@ -1,5 +1,6 @@
 use std::sync::OnceLock;
 
+use rustc_abi::VariantIdx;
 use rustc_data_structures::graph;
 use rustc_data_structures::graph::dominators::{Dominators, dominators};
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
@@ -23,7 +24,7 @@ type Predecessors = IndexVec<BasicBlock, SmallVec<[BasicBlock; 4]>>;
 #[derive(Debug, Clone, Copy)]
 pub enum SwitchTargetValue {
     // A normal switch value.
-    Normal(u128),
+    Normal(VariantIdx),
     // The final "otherwise" fallback value.
     Otherwise,
 }
diff --git a/compiler/rustc_mir_dataflow/src/drop_flag_effects.rs b/compiler/rustc_mir_dataflow/src/drop_flag_effects.rs
@@ -1,6 +1,5 @@
 use rustc_abi::VariantIdx;
 use rustc_middle::mir::{self, Body, Location, Terminator, TerminatorKind};
-use smallvec::SmallVec;
 use tracing::debug;
 
 use super::move_paths::{InitKind, LookupResult, MoveData, MovePathIndex};
@@ -158,15 +157,17 @@ where
 
 /// Indicates which variants are inactive at a `SwitchInt` edge by listing their `VariantIdx`s or
 /// specifying the single active variant's `VariantIdx`.
-pub(crate) enum InactiveVariants {
-    Inactives(SmallVec<[VariantIdx; 4]>),
+pub(crate) enum InactiveVariants<'a> {
+    Inactives(&'a [(VariantIdx, mir::BasicBlock)]),
     Active(VariantIdx),
 }
 
-impl InactiveVariants {
+impl InactiveVariants<'_> {
     fn contains(&self, variant_idx: VariantIdx) -> bool {
         match self {
-            InactiveVariants::Inactives(inactives) => inactives.contains(&variant_idx),
+            InactiveVariants::Inactives(inactives) => {
+                inactives.iter().any(|(idx, _)| *idx == variant_idx)
+            }
             InactiveVariants::Active(active) => variant_idx != *active,
         }
     }
@@ -177,7 +178,7 @@ impl InactiveVariants {
 pub(crate) fn on_all_inactive_variants<'tcx>(
     move_data: &MoveData<'tcx>,
     enum_place: mir::Place<'tcx>,
-    inactive_variants: &InactiveVariants,
+    inactive_variants: &InactiveVariants<'_>,
     mut handle_inactive_variant: impl FnMut(MovePathIndex),
 ) {
     let LookupResult::Exact(enum_mpi) = move_data.rev_lookup.find(enum_place.as_ref()) else {
diff --git a/compiler/rustc_mir_dataflow/src/framework/direction.rs b/compiler/rustc_mir_dataflow/src/framework/direction.rs
@@ -113,8 +113,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, discr, targets) {
                         bug!(
                             "SwitchInt edge effects are unsupported in backward dataflow analyses"
                         );
@@ -283,23 +283,22 @@ impl Direction for Forward {
                 }
             }
             TerminatorEdges::SwitchInt { targets, discr } => {
-                if let Some(mut data) = analysis.get_switch_int_data(block, discr) {
+                if let Some(data) = analysis.get_switch_int_data(block, discr, targets) {
                     let mut tmp = analysis.bottom_value(body);
-                    for (value, target) in targets.iter() {
+                    for (variant_idx, target) in A::switch_int_target_variants(&data) {
                         tmp.clone_from(exit_state);
-                        let value = SwitchTargetValue::Normal(value);
-                        analysis.apply_switch_int_edge_effect(&mut data, &mut tmp, value, targets);
-                        propagate(target, &tmp);
+                        let value = SwitchTargetValue::Normal(*variant_idx);
+                        analysis.apply_switch_int_edge_effect(&data, &mut tmp, value);
+                        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,
+                        &data,
                         exit_state,
                         SwitchTargetValue::Otherwise,
-                        targets,
                     );
                     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
@@ -34,6 +34,7 @@
 
 use std::cmp::Ordering;
 
+use rustc_abi::VariantIdx;
 use rustc_data_structures::work_queue::WorkQueue;
 use rustc_index::bit_set::{DenseBitSet, MixedBitSet};
 use rustc_index::{Idx, IndexVec};
@@ -214,17 +215,24 @@ pub trait Analysis<'tcx> {
         &mut self,
         _block: mir::BasicBlock,
         _discr: &mir::Operand<'tcx>,
+        _targets: &mir::SwitchTargets,
     ) -> Option<Self::SwitchIntData> {
         None
     }
 
+    /// Returns an iterator over `SwitchInt` target variants stored in `Self::SwitchIntData`.
+    fn switch_int_target_variants(
+        _data: &Self::SwitchIntData,
+    ) -> impl Iterator<Item = &(VariantIdx, mir::BasicBlock)> {
+        [].iter()
+    }
+
     /// See comments on `get_switch_int_data`.
     fn apply_switch_int_edge_effect(
         &mut self,
-        _data: &mut Self::SwitchIntData,
+        _data: &Self::SwitchIntData,
         _state: &mut Self::Domain,
         _value: SwitchTargetValue,
-        _targets: &mir::SwitchTargets,
     ) {
         unreachable!();
     }
diff --git a/compiler/rustc_mir_dataflow/src/impls/initialized.rs b/compiler/rustc_mir_dataflow/src/impls/initialized.rs
@@ -7,9 +7,7 @@ use rustc_middle::bug;
 use rustc_middle::mir::{
     self, Body, CallReturnPlaces, Location, SwitchTargetValue, TerminatorEdges,
 };
-use rustc_middle::ty::util::Discr;
-use rustc_middle::ty::{self, TyCtxt};
-use smallvec::SmallVec;
+use rustc_middle::ty::{self, AdtDef, TyCtxt};
 use tracing::{debug, instrument};
 
 use crate::drop_flag_effects::{DropFlagState, InactiveVariants};
@@ -22,30 +20,25 @@ use crate::{
 // Used by both `MaybeInitializedPlaces` and `MaybeUninitializedPlaces`.
 pub struct MaybePlacesSwitchIntData<'tcx> {
     enum_place: mir::Place<'tcx>,
-    discriminants: Vec<(VariantIdx, Discr<'tcx>)>,
-    index: usize,
+    targets: Vec<(VariantIdx, mir::BasicBlock)>,
 }
 
-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;
-            }
-        }
-    }
+/// Maps values of targets in `SwitchTargets` to `(VariantIdx, BasicBlock).` Panics if the variants
+/// in `targets` aren't in the same order as `AdtDef::discriminants`.
+fn collect_switch_targets<'tcx>(
+    enum_def: AdtDef<'tcx>,
+    targets: &mir::SwitchTargets,
+    tcx: TyCtxt<'tcx>,
+) -> Vec<(VariantIdx, mir::BasicBlock)> {
+    let mut discriminants = enum_def.discriminants(tcx);
+
+    Vec::from_iter(targets.iter().map(|(value, bb)| {
+        let Some((variant_idx, _)) = discriminants.find(|(_, discr)| discr.val == value) else {
+            bug!("ran out of discriminants before matching all switch targets");
+        };
+
+        (variant_idx, bb)
+    }))
 }
 
 impl<'tcx> MaybePlacesSwitchIntData<'tcx> {
@@ -54,6 +47,7 @@ impl<'tcx> MaybePlacesSwitchIntData<'tcx> {
         body: &Body<'tcx>,
         block: mir::BasicBlock,
         discr: &mir::Operand<'tcx>,
+        targets: &mir::SwitchTargets,
     ) -> Option<Self> {
         let Some(discr) = discr.place() else { return None };
 
@@ -78,8 +72,7 @@ impl<'tcx> MaybePlacesSwitchIntData<'tcx> {
                         ty::Adt(enum_def, _) => {
                             return Some(MaybePlacesSwitchIntData {
                                 enum_place,
-                                discriminants: enum_def.discriminants(tcx).collect(),
-                                index: 0,
+                                targets: collect_switch_targets(*enum_def, targets, tcx),
                             });
                         }
 
@@ -448,25 +441,32 @@ impl<'tcx> Analysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
         &mut self,
         block: mir::BasicBlock,
         discr: &mir::Operand<'tcx>,
+        targets: &mir::SwitchTargets,
     ) -> 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, discr, targets)
+    }
+
+    #[inline]
+    fn switch_int_target_variants<'a>(
+        data: &'a Self::SwitchIntData,
+    ) -> impl Iterator<Item = &'a (VariantIdx, mir::BasicBlock)> {
+        data.targets.iter()
     }
 
     fn apply_switch_int_edge_effect(
         &mut self,
-        data: &mut Self::SwitchIntData,
+        data: &Self::SwitchIntData,
         state: &mut Self::Domain,
         value: SwitchTargetValue,
-        targets: &mir::SwitchTargets,
     ) {
         let inactive_variants = match value {
-            SwitchTargetValue::Normal(value) => InactiveVariants::Active(data.next_discr(value)),
+            SwitchTargetValue::Normal(variant_idx) => InactiveVariants::Active(variant_idx),
             SwitchTargetValue::Otherwise if self.exclude_inactive_in_otherwise => {
-                InactiveVariants::Inactives(data.variants(targets))
+                InactiveVariants::Inactives(&data.targets)
             }
             _ => return,
         };
@@ -564,6 +564,7 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
         &mut self,
         block: mir::BasicBlock,
         discr: &mir::Operand<'tcx>,
+        targets: &mir::SwitchTargets,
     ) -> Option<Self::SwitchIntData> {
         if !self.tcx.sess.opts.unstable_opts.precise_enum_drop_elaboration {
             return None;
@@ -573,20 +574,26 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
             return None;
         }
 
-        MaybePlacesSwitchIntData::new(self.tcx, self.body, block, discr)
+        MaybePlacesSwitchIntData::new(self.tcx, self.body, block, discr, targets)
+    }
+
+    #[inline]
+    fn switch_int_target_variants<'a>(
+        data: &'a Self::SwitchIntData,
+    ) -> impl Iterator<Item = &'a (VariantIdx, mir::BasicBlock)> {
+        data.targets.iter()
     }
 
     fn apply_switch_int_edge_effect(
         &mut self,
-        data: &mut Self::SwitchIntData,
+        data: &Self::SwitchIntData,
         state: &mut Self::Domain,
         value: SwitchTargetValue,
-        targets: &mir::SwitchTargets,
     ) {
         let inactive_variants = match value {
-            SwitchTargetValue::Normal(value) => InactiveVariants::Active(data.next_discr(value)),
+            SwitchTargetValue::Normal(variant_idx) => InactiveVariants::Active(variant_idx),
             SwitchTargetValue::Otherwise if self.include_inactive_in_otherwise => {
-                InactiveVariants::Inactives(data.variants(targets))
+                InactiveVariants::Inactives(&data.targets)
             }
             _ => return,
         };
