Auto merge of #142707 - ashivaram23:drop_wildcard, r=<try>

Update `MaybeUninitializedPlaces` and `MaybeInitializedPlaces` for `otherwise`

<!-- homu-ignore:start -->
<!--
If this PR is related to an unstable feature or an otherwise tracked effort,
please link to the relevant tracking issue here. If you don't know of a related
tracking issue or there are none, feel free to ignore this.

This PR will get automatically assigned to a reviewer. In case you would like
a specific user to review your work, you can assign it to them by using

    r? <reviewer name>
-->
<!-- homu-ignore:end -->

This allows `ElaborateDrops` to remove drops when a `match` wildcard arm covers multiple no-Drop enum variants. It modifies dataflow analysis to update the `MaybeUninitializedPlaces` and `MaybeInitializedPlaces` data for a block reached through an `otherwise` edge.

This appears to fix #142705, but I don't know for sure if it's actually correct (are there cases where this would be wrong and break things?). I also haven't tested that it actually improves compile times, or machine code output outside of the examples in the issue.

Author: rust-bors[bot]

compiler/rustc_mir_dataflow/src/drop_flag_effects.rs

@@ -155,16 +155,15 @@ where
     }
 }
 
-/// Calls `handle_inactive_variant` for each descendant move path of `enum_place` that contains a
-/// `Downcast` to a variant besides the `active_variant`.
-///
-/// NOTE: If there are no move paths corresponding to an inactive variant,
-/// `handle_inactive_variant` will not be called for that variant.
-pub(crate) fn on_all_inactive_variants<'tcx>(
+/// 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>(
     move_data: &MoveData<'tcx>,
     enum_place: mir::Place<'tcx>,
     active_variant: VariantIdx,
     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;
@@ -184,6 +183,8 @@ pub(crate) fn on_all_inactive_variants<'tcx>(
 
         if variant_idx != active_variant {
             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));
         }
     }
 }

compiler/rustc_mir_dataflow/src/framework/direction.rs

@@ -1,8 +1,6 @@
 use std::ops::RangeInclusive;
 
-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};
@@ -117,7 +115,7 @@ impl Direction for Backward {
                         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);
+                            analysis.apply_switch_int_edge_effect(&mut data, &mut tmp, value, None);
                             propagate(pred, &tmp);
                         }
                     } else {
@@ -245,7 +243,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>,
@@ -285,25 +283,10 @@ impl Direction for Forward {
                 }
             }
             TerminatorEdges::SwitchInt { targets, discr } => {
-                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);
-                        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.
-                    let otherwise = targets.otherwise();
-                    analysis.apply_switch_int_edge_effect(
-                        &mut data,
-                        exit_state,
-                        SwitchTargetValue::Otherwise,
+                if let Some(data) = analysis.get_switch_int_data(block, discr) {
+                    analysis.apply_switch_int_edge_effect_for_targets(
+                        targets, data, exit_state, propagate,
                     );
-                    propagate(otherwise, exit_state);
                 } else {
                     for target in targets.all_targets() {
                         propagate(*target, exit_state);

compiler/rustc_mir_dataflow/src/framework/mod.rs

@@ -224,6 +224,20 @@ 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!();
     }

compiler/rustc_mir_dataflow/src/impls/initialized.rs

@@ -131,12 +131,26 @@ pub struct MaybeInitializedPlaces<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     body: &'a Body<'tcx>,
     move_data: &'a MoveData<'tcx>,
+    exclude_inactive_in_otherwise: bool,
     skip_unreachable_unwind: bool,
 }
 
 impl<'a, 'tcx> MaybeInitializedPlaces<'a, 'tcx> {
     pub fn new(tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>, move_data: &'a MoveData<'tcx>) -> Self {
-        MaybeInitializedPlaces { tcx, body, move_data, skip_unreachable_unwind: false }
+        MaybeInitializedPlaces {
+            tcx,
+            body,
+            move_data,
+            exclude_inactive_in_otherwise: false,
+            skip_unreachable_unwind: false,
+        }
+    }
+
+    /// Ensures definitely inactive variants are excluded from the set of initialized places for
+    /// blocks reached through an `otherwise` edge.
+    pub fn exclude_inactive_in_otherwise(mut self) -> Self {
+        self.exclude_inactive_in_otherwise = true;
+        self
     }
 
     pub fn skipping_unreachable_unwind(mut self) -> Self {
@@ -208,6 +222,7 @@ pub struct MaybeUninitializedPlaces<'a, 'tcx> {
     move_data: &'a MoveData<'tcx>,
 
     mark_inactive_variants_as_uninit: bool,
+    include_inactive_in_otherwise: bool,
     skip_unreachable_unwind: DenseBitSet<mir::BasicBlock>,
 }
 
@@ -218,6 +233,7 @@ impl<'a, 'tcx> MaybeUninitializedPlaces<'a, 'tcx> {
             body,
             move_data,
             mark_inactive_variants_as_uninit: false,
+            include_inactive_in_otherwise: false,
             skip_unreachable_unwind: DenseBitSet::new_empty(body.basic_blocks.len()),
         }
     }
@@ -232,6 +248,13 @@ impl<'a, 'tcx> MaybeUninitializedPlaces<'a, 'tcx> {
         self
     }
 
+    /// Ensures definitely inactive variants are included in the set of uninitialized places for
+    /// blocks reached through an `otherwise` edge.
+    pub fn include_inactive_in_otherwise(mut self) -> Self {
+        self.include_inactive_in_otherwise = true;
+        self
+    }
+
     pub fn skipping_unreachable_unwind(
         mut self,
         unreachable_unwind: DenseBitSet<mir::BasicBlock>,
@@ -431,17 +454,63 @@ impl<'tcx> Analysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
         data: &mut Self::SwitchIntData,
         state: &mut Self::Domain,
         value: SwitchTargetValue,
+        otherwise_state: Option<&mut Self::Domain>,
+    ) {
+        let SwitchTargetValue::Normal(value) = value else {
+            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),
     ) {
-        if let SwitchTargetValue::Normal(value) = value {
-            // 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,
-                data.enum_place,
-                data.next_discr(value),
-                |mpi| state.kill(mpi),
+        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));
     }
 }
 
@@ -544,17 +613,63 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
         data: &mut Self::SwitchIntData,
         state: &mut Self::Domain,
         value: SwitchTargetValue,
+        otherwise_state: Option<&mut Self::Domain>,
+    ) {
+        let SwitchTargetValue::Normal(value) = value else {
+            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),
     ) {
-        if let SwitchTargetValue::Normal(value) = value {
-            // 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,
-                data.enum_place,
-                data.next_discr(value),
-                |mpi| state.gen_(mpi),
+        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));
     }
 }
 

compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -62,6 +62,7 @@ impl<'tcx> crate::MirPass<'tcx> for ElaborateDrops {
             let env = MoveDataTypingEnv { move_data, typing_env };
 
             let mut inits = MaybeInitializedPlaces::new(tcx, body, &env.move_data)
+                .exclude_inactive_in_otherwise()
                 .skipping_unreachable_unwind()
                 .iterate_to_fixpoint(tcx, body, Some("elaborate_drops"))
                 .into_results_cursor(body);

compiler/rustc_mir_transform/src/remove_uninit_drops.rs

@@ -22,6 +22,7 @@ impl<'tcx> crate::MirPass<'tcx> for RemoveUninitDrops {
         let move_data = MoveData::gather_moves(body, tcx, |ty| ty.needs_drop(tcx, typing_env));
 
         let mut maybe_inits = MaybeInitializedPlaces::new(tcx, body, &move_data)
+            .exclude_inactive_in_otherwise()
             .iterate_to_fixpoint(tcx, body, Some("remove_uninit_drops"))
             .into_results_cursor(body);