Auto merge of #142707 - ashivaram23:drop_wildcard, r=dianqk

Apply effects to `otherwise` edge in dataflow analysis

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.

Fixes rust-lang/rust#142705.

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,28 @@ where
     }
 }
 
-/// Calls `handle_inactive_variant` for each descendant move path of `enum_place` that contains a
+/// Indicates which variants are inactive at a `SwitchInt` edge by listing their `VariantIdx`s or
-/// `Downcast` to a variant besides the `active_variant`.
+/// specifying the single active variant's `VariantIdx`.
-///
+pub(crate) enum InactiveVariants {
-/// NOTE: If there are no move paths corresponding to an inactive variant,
+    Inactives(SmallVec<[VariantIdx; 4]>),
-/// `handle_inactive_variant` will not be called for that variant.
+    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),
 ) {
     let LookupResult::Exact(enum_mpi) = move_data.rev_lookup.find(enum_place.as_ref()) else {
@@ -182,7 +196,7 @@ pub(crate) fn on_all_inactive_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));
         }
     }

compiler/rustc_mir_dataflow/src/framework/direction.rs

@@ -112,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);
+                            analysis
+                                .apply_switch_int_edge_effect(&mut data, &mut tmp, value, targets);
                             propagate(pred, &tmp);
                         }
                     } else {
@@ -290,20 +291,20 @@ impl Direction for Forward {
                     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);
+                        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.
-                    let otherwise = targets.otherwise();
                     analysis.apply_switch_int_edge_effect(
                         &mut data,
                         exit_state,
                         SwitchTargetValue::Otherwise,
+                        targets,
                     );
-                    propagate(otherwise, exit_state);
+                    propagate(targets.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,7 @@ pub trait Analysis<'tcx> {
         _data: &mut Self::SwitchIntData,
         _state: &mut Self::Domain,
         _value: SwitchTargetValue,
+        _targets: &mir::SwitchTargets,
     ) {
         unreachable!();
     }

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.
@@ -131,12 +138,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 +229,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 +240,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 +255,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 +461,24 @@ impl<'tcx> Analysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
         data: &mut Self::SwitchIntData,
         state: &mut Self::Domain,
         value: SwitchTargetValue,
+        targets: &mir::SwitchTargets,
     ) {
-        if let SwitchTargetValue::Normal(value) = value {
+        let inactive_variants = match value {
-            // Kill all move paths that correspond to variants we know to be inactive along this
+            SwitchTargetValue::Normal(value) => InactiveVariants::Active(data.next_discr(value)),
-            // particular outgoing edge of a `SwitchInt`.
+            SwitchTargetValue::Otherwise if self.exclude_inactive_in_otherwise => {
-            drop_flag_effects::on_all_inactive_variants(
+                InactiveVariants::Inactives(data.variants(targets))
-                self.move_data,
+            }
-                data.enum_place,
+            _ => return,
-                data.next_discr(value),
+        };
-                |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`.
+        drop_flag_effects::on_all_inactive_variants(
+            self.move_data,
+            data.enum_place,
+            &inactive_variants,
+            |mpi| state.kill(mpi),
+        );
     }
 }
 
@@ -544,17 +581,24 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
         data: &mut Self::SwitchIntData,
         state: &mut Self::Domain,
         value: SwitchTargetValue,
+        targets: &mir::SwitchTargets,
     ) {
-        if let SwitchTargetValue::Normal(value) = value {
+        let inactive_variants = match value {
-            // Mark all move paths that correspond to variants other than this one as maybe
+            SwitchTargetValue::Normal(value) => InactiveVariants::Active(data.next_discr(value)),
-            // uninitialized (in reality, they are *definitely* uninitialized).
+            SwitchTargetValue::Otherwise if self.include_inactive_in_otherwise => {
-            drop_flag_effects::on_all_inactive_variants(
+                InactiveVariants::Inactives(data.variants(targets))
-                self.move_data,
+            }
-                data.enum_place,
+            _ => return,
-                data.next_discr(value),
+        };
-                |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).
+        drop_flag_effects::on_all_inactive_variants(
+            self.move_data,
+            data.enum_place,
+            &inactive_variants,
+            |mpi| state.gen_(mpi),
+        );
     }
 }
 

compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -62,12 +62,14 @@ 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);
             let dead_unwinds = compute_dead_unwinds(body, &mut inits);
 
             let uninits = MaybeUninitializedPlaces::new(tcx, body, &env.move_data)
+                .include_inactive_in_otherwise()
                 .mark_inactive_variants_as_uninit()
                 .skipping_unreachable_unwind(dead_unwinds)
                 .iterate_to_fixpoint(tcx, body, Some("elaborate_drops"))

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);
 

tests/mir-opt/otherwise_drops.result_ok.ElaborateDrops.diff

@@ -0,0 +1,108 @@
+- // MIR for `result_ok` before ElaborateDrops
++ // MIR for `result_ok` after ElaborateDrops
+  
+  fn result_ok(_1: Result<String, ()>) -> Option<String> {
+      debug result => _1;
+      let mut _0: std::option::Option<std::string::String>;
+      let mut _2: isize;
+      let _3: std::string::String;
+      let mut _4: std::string::String;
++     let mut _5: bool;
++     let mut _6: isize;
++     let mut _7: isize;
+      scope 1 {
+          debug s => _3;
+      }
+  
+      bb0: {
++         _5 = const false;
++         _5 = const true;
+          PlaceMention(_1);
+          _2 = discriminant(_1);
+          switchInt(move _2) -> [0: bb2, otherwise: bb1];
+      }
+  
+      bb1: {
+          _0 = Option::<String>::None;
+          goto -> bb5;
+      }
+  
+      bb2: {
+          StorageLive(_3);
+          _3 = move ((_1 as Ok).0: std::string::String);
+          StorageLive(_4);
+          _4 = move _3;
+          _0 = Option::<String>::Some(move _4);
+-         drop(_4) -> [return: bb3, unwind: bb7];
++         goto -> bb3;
+      }
+  
+      bb3: {
+          StorageDead(_4);
+-         drop(_3) -> [return: bb4, unwind: bb8];
++         goto -> bb4;
+      }
+  
+      bb4: {
+          StorageDead(_3);
+          goto -> bb5;
+      }
+  
+      bb5: {
+-         drop(_1) -> [return: bb6, unwind: bb9];
++         goto -> bb16;
+      }
+  
+      bb6: {
+          return;
+      }
+  
+      bb7 (cleanup): {
+-         drop(_3) -> [return: bb8, unwind terminate(cleanup)];
++         goto -> bb8;
+      }
+  
+      bb8 (cleanup): {
+-         drop(_1) -> [return: bb9, unwind terminate(cleanup)];
++         goto -> bb9;
+      }
+  
+      bb9 (cleanup): {
+          resume;
++     }
++ 
++     bb10: {
++         goto -> bb6;
++     }
++ 
++     bb11 (cleanup): {
++         goto -> bb9;
++     }
++ 
++     bb12 (cleanup): {
++         goto -> bb9;
++     }
++ 
++     bb13: {
++         goto -> bb10;
++     }
++ 
++     bb14: {
++         goto -> bb10;
++     }
++ 
++     bb15 (cleanup): {
++         goto -> bb9;
++     }
++ 
++     bb16: {
++         _6 = discriminant(_1);
++         switchInt(move _6) -> [0: bb13, otherwise: bb14];
++     }
++ 
++     bb17 (cleanup): {
++         _7 = discriminant(_1);
++         switchInt(move _7) -> [0: bb11, otherwise: bb15];
+      }
+  }
+  

tests/mir-opt/otherwise_drops.rs

@@ -0,0 +1,13 @@
+//@ compile-flags: -C panic=abort
+//@ test-mir-pass: ElaborateDrops
+
+// Ensures there are no drops for the wildcard match arm.
+
+// EMIT_MIR otherwise_drops.result_ok.ElaborateDrops.diff
+fn result_ok(result: Result<String, ()>) -> Option<String> {
+    // CHECK-NOT: drop
+    match result {
+        Ok(s) => Some(s),
+        _ => None,
+    }
+}

tests/mir-opt/sroa/lifetimes.foo.ScalarReplacementOfAggregates.diff

@@ -96,7 +96,6 @@
   
       bb2: {
           StorageLive(_8);
-          _28 = const false;
           _8 = move ((_5 as Ok).0: std::boxed::Box<dyn std::fmt::Display>);
           StorageLive(_9);
           StorageLive(_10);
@@ -186,7 +185,7 @@
       bb9: {
           StorageDead(_6);
           _29 = discriminant(_5);
-          switchInt(move _29) -> [0: bb11, otherwise: bb13];
+          switchInt(move _29) -> [0: bb10, otherwise: bb11];
       }
   
       bb10: {
@@ -200,14 +199,6 @@
       }
   
       bb11: {
-          switchInt(copy _28) -> [0: bb10, otherwise: bb12];
-      }
-  
-      bb12: {
-          drop(((_5 as Ok).0: std::boxed::Box<dyn std::fmt::Display>)) -> [return: bb10, unwind unreachable];
-      }
-  
-      bb13: {
           drop(_5) -> [return: bb10, unwind unreachable];
       }
   }