address review comments

ashivaram23 · ashivaram23 · commit 881c62f9587f · 2025-07-08T01:23:56.000-05:00
diff --git a/compiler/rustc_mir_dataflow/src/drop_flag_effects.rs b/compiler/rustc_mir_dataflow/src/drop_flag_effects.rs
@@ -1,5 +1,6 @@
 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};
@@ -155,15 +156,29 @@ where
     }
 }
 
-/// Handles each variant that corresponds to one of the child move paths of `enum_place`. If the
-/// variant is `active_variant`, `handle_active_variant` will be called. Otherwise,
-/// `handle_inactive_variant` will be called.
-pub(crate) fn on_all_variants<'tcx>(
+/// 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]>),
+    Active(VariantIdx),
+}
+
+impl InactiveVariants {
+    fn contains(&self, variant_idx: VariantIdx) -> bool {
+        match self {
+            InactiveVariants::Inactives(inactives) => inactives.contains(&variant_idx),
+            InactiveVariants::Active(active) => variant_idx != *active,
+        }
+    }
+}
+
+/// Calls `handle_inactive_variant` for each child move path of `enum_place` corresponding to an
+/// inactive variant at a particular `SwitchInt` edge.
+pub(crate) fn on_all_inactive_variants<'tcx>(
     move_data: &MoveData<'tcx>,
     enum_place: mir::Place<'tcx>,
-    active_variant: VariantIdx,
+    inactive_variants: &InactiveVariants,
     mut handle_inactive_variant: impl FnMut(MovePathIndex),
-    mut handle_active_variant: impl FnMut(MovePathIndex),
 ) {
     let LookupResult::Exact(enum_mpi) = move_data.rev_lookup.find(enum_place.as_ref()) else {
         return;
@@ -181,10 +196,8 @@ pub(crate) fn on_all_variants<'tcx>(
             unreachable!();
         };
 
-        if variant_idx != active_variant {
+        if inactive_variants.contains(variant_idx) {
             on_all_children_bits(move_data, variant_mpi, |mpi| handle_inactive_variant(mpi));
-        } else {
-            on_all_children_bits(move_data, variant_mpi, |mpi| handle_active_variant(mpi));
         }
     }
 }
diff --git a/compiler/rustc_mir_dataflow/src/framework/direction.rs b/compiler/rustc_mir_dataflow/src/framework/direction.rs
@@ -1,6 +1,8 @@
 use std::ops::RangeInclusive;
 
-use rustc_middle::mir::{self, BasicBlock, CallReturnPlaces, Location, TerminatorEdges};
+use rustc_middle::mir::{
+    self, BasicBlock, CallReturnPlaces, Location, SwitchTargetValue, TerminatorEdges,
+};
 
 use super::visitor::ResultsVisitor;
 use super::{Analysis, Effect, EffectIndex};
@@ -110,12 +112,13 @@ impl Direction for Backward {
                     propagate(pred, &tmp);
                 }
 
-                mir::TerminatorKind::SwitchInt { targets: _, ref discr } => {
+                mir::TerminatorKind::SwitchInt { ref targets, ref discr } => {
                     if let Some(mut data) = analysis.get_switch_int_data(block, discr) {
                         let mut tmp = analysis.bottom_value(body);
                         for &value in &body.basic_blocks.switch_sources()[&(block, pred)] {
                             tmp.clone_from(exit_state);
-                            analysis.apply_switch_int_edge_effect(&mut data, &mut tmp, value, None);
+                            analysis
+                                .apply_switch_int_edge_effect(&mut data, &mut tmp, value, targets);
                             propagate(pred, &tmp);
                         }
                     } else {
@@ -243,7 +246,7 @@ impl Direction for Forward {
 
     fn apply_effects_in_block<'mir, 'tcx, A>(
         analysis: &mut A,
-        _body: &mir::Body<'tcx>,
+        body: &mir::Body<'tcx>,
         state: &mut A::Domain,
         block: BasicBlock,
         block_data: &'mir mir::BasicBlockData<'tcx>,
@@ -283,10 +286,25 @@ impl Direction for Forward {
                 }
             }
             TerminatorEdges::SwitchInt { targets, discr } => {
-                if let Some(data) = analysis.get_switch_int_data(block, discr) {
-                    analysis.apply_switch_int_edge_effect_for_targets(
-                        targets, data, exit_state, propagate,
+                if let Some(mut data) = analysis.get_switch_int_data(block, discr) {
+                    let mut tmp = analysis.bottom_value(body);
+                    for (value, target) in targets.iter() {
+                        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);
+                    }
+
+                    // 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,
                     );
+                    propagate(targets.otherwise(), exit_state);
                 } else {
                     for target in targets.all_targets() {
                         propagate(*target, exit_state);
diff --git a/compiler/rustc_mir_dataflow/src/framework/mod.rs b/compiler/rustc_mir_dataflow/src/framework/mod.rs
@@ -224,20 +224,7 @@ pub trait Analysis<'tcx> {
         _data: &mut Self::SwitchIntData,
         _state: &mut Self::Domain,
         _value: SwitchTargetValue,
-        _otherwise_state: Option<&mut Self::Domain>,
-    ) {
-        unreachable!();
-    }
-
-    /// Calls `apply_switch_int_edge_effect` for each target in `targets` and calls `propagate` with
-    /// the new state. This is used in forward analysis for `MaybeUninitializedPlaces` and
-    /// `MaybeInitializedPlaces`.
-    fn apply_switch_int_edge_effect_for_targets(
-        &mut self,
         _targets: &mir::SwitchTargets,
-        mut _data: Self::SwitchIntData,
-        _state: &mut Self::Domain,
-        mut _propagate: impl FnMut(mir::BasicBlock, &Self::Domain),
     ) {
         unreachable!();
     }
diff --git a/compiler/rustc_mir_dataflow/src/impls/initialized.rs b/compiler/rustc_mir_dataflow/src/impls/initialized.rs
@@ -9,9 +9,10 @@ use rustc_middle::mir::{
 };
 use rustc_middle::ty::util::Discr;
 use rustc_middle::ty::{self, TyCtxt};
+use smallvec::SmallVec;
 use tracing::{debug, instrument};
 
-use crate::drop_flag_effects::DropFlagState;
+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,
@@ -26,6 +27,12 @@ pub struct MaybePlacesSwitchIntData<'tcx> {
 }
 
 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.
@@ -454,63 +461,24 @@ impl<'tcx> Analysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
         data: &mut Self::SwitchIntData,
         state: &mut Self::Domain,
         value: SwitchTargetValue,
-        otherwise_state: Option<&mut Self::Domain>,
+        targets: &mir::SwitchTargets,
     ) {
-        let SwitchTargetValue::Normal(value) = value else {
-            return;
+        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))
+            }
+            _ => return,
         };
 
-        let handle_inactive_variant = |mpi| state.kill(mpi);
-
         // Kill all move paths that correspond to variants we know to be inactive along this
         // particular outgoing edge of a `SwitchInt`.
-        match otherwise_state {
-            Some(otherwise_state) => {
-                drop_flag_effects::on_all_variants(
-                    self.move_data,
-                    data.enum_place,
-                    data.next_discr(value),
-                    handle_inactive_variant,
-                    |mpi| otherwise_state.kill(mpi),
-                );
-            }
-            None => {
-                drop_flag_effects::on_all_variants(
-                    self.move_data,
-                    data.enum_place,
-                    data.next_discr(value),
-                    handle_inactive_variant,
-                    |_mpi| {},
-                );
-            }
-        }
-    }
-
-    fn apply_switch_int_edge_effect_for_targets(
-        &mut self,
-        targets: &mir::SwitchTargets,
-        mut data: Self::SwitchIntData,
-        state: &mut Self::Domain,
-        mut propagate: impl FnMut(mir::BasicBlock, &Self::Domain),
-    ) {
-        let analyze_otherwise = self.exclude_inactive_in_otherwise
-            && (1..data.discriminants.len()).contains(&targets.all_values().len());
-
-        let mut otherwise_state = if analyze_otherwise { Some(state.clone()) } else { None };
-        let mut target_state = MaybeReachable::Unreachable;
-
-        for (value, target) in targets.iter() {
-            target_state.clone_from(&state);
-            self.apply_switch_int_edge_effect(
-                &mut data,
-                &mut target_state,
-                SwitchTargetValue::Normal(value),
-                otherwise_state.as_mut(),
-            );
-            propagate(target, &target_state);
-        }
-
-        propagate(targets.otherwise(), otherwise_state.as_ref().unwrap_or(state));
+        drop_flag_effects::on_all_inactive_variants(
+            self.move_data,
+            data.enum_place,
+            &inactive_variants,
+            |mpi| state.kill(mpi),
+        );
     }
 }
 
@@ -613,63 +581,24 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
         data: &mut Self::SwitchIntData,
         state: &mut Self::Domain,
         value: SwitchTargetValue,
-        otherwise_state: Option<&mut Self::Domain>,
+        targets: &mir::SwitchTargets,
     ) {
-        let SwitchTargetValue::Normal(value) = value else {
-            return;
+        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))
+            }
+            _ => return,
         };
 
-        let handle_inactive_variant = |mpi| state.gen_(mpi);
-
         // Mark all move paths that correspond to variants other than this one as maybe
         // uninitialized (in reality, they are *definitely* uninitialized).
-        match otherwise_state {
-            Some(otherwise_state) => {
-                drop_flag_effects::on_all_variants(
-                    self.move_data,
-                    data.enum_place,
-                    data.next_discr(value),
-                    handle_inactive_variant,
-                    |mpi| otherwise_state.gen_(mpi),
-                );
-            }
-            None => {
-                drop_flag_effects::on_all_variants(
-                    self.move_data,
-                    data.enum_place,
-                    data.next_discr(value),
-                    handle_inactive_variant,
-                    |_mpi| {},
-                );
-            }
-        }
-    }
-
-    fn apply_switch_int_edge_effect_for_targets(
-        &mut self,
-        targets: &mir::SwitchTargets,
-        mut data: Self::SwitchIntData,
-        state: &mut Self::Domain,
-        mut propagate: impl FnMut(mir::BasicBlock, &Self::Domain),
-    ) {
-        let analyze_otherwise = self.include_inactive_in_otherwise
-            && (1..data.discriminants.len()).contains(&targets.all_values().len());
-
-        let mut otherwise_state = if analyze_otherwise { Some(state.clone()) } else { None };
-        let mut target_state = MixedBitSet::new_empty(self.move_data().move_paths.len());
-
-        for (value, target) in targets.iter() {
-            target_state.clone_from(&state);
-            self.apply_switch_int_edge_effect(
-                &mut data,
-                &mut target_state,
-                SwitchTargetValue::Normal(value),
-                otherwise_state.as_mut(),
-            );
-            propagate(target, &target_state);
-        }
-
-        propagate(targets.otherwise(), otherwise_state.as_ref().unwrap_or(state));
+        drop_flag_effects::on_all_inactive_variants(
+            self.move_data,
+            data.enum_place,
+            &inactive_variants,
+            |mpi| state.gen_(mpi),
+        );
     }
 }
 
diff --git a/tests/mir-opt/otherwise_drops.rs b/tests/mir-opt/otherwise_drops.rs
@@ -1,10 +1,12 @@
-// skip-filecheck
 //@ test-mir-pass: ElaborateDrops
 
 // Ensures there are no drops for the wildcard match arm.
 
 // EMIT_MIR otherwise_drops.result_ok.ElaborateDrops.after.mir
 fn result_ok(result: Result<String, ()>) -> Option<String> {
+    // CHECK-LABEL: fn result_ok(
+    // CHECK-NOT: drop
+    // CHECK: return
     match result {
         Ok(s) => Some(s),
         _ => None,
