Only run MaybeInitializedPlaces once for drop elaboration.

Author: cjgillot

compiler/rustc_mir_dataflow/src/impls/initialized.rs

@@ -10,8 +10,8 @@ use crate::framework::SwitchIntEdgeEffects;
 use crate::move_paths::{HasMoveData, InitIndex, InitKind, LookupResult, MoveData, MovePathIndex};
 use crate::on_lookup_result_bits;
 use crate::MoveDataParamEnv;
-use crate::{drop_flag_effects, on_all_children_bits};
-use crate::{lattice, AnalysisDomain, GenKill, GenKillAnalysis};
+use crate::{drop_flag_effects, on_all_children_bits, on_all_drop_children_bits};
+use crate::{lattice, AnalysisDomain, GenKill, GenKillAnalysis, MaybeUnreachable};
 
 /// `MaybeInitializedPlaces` tracks all places that might be
 /// initialized upon reaching a particular point in the control flow
@@ -52,11 +52,33 @@ pub struct MaybeInitializedPlaces<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     body: &'a Body<'tcx>,
     mdpe: &'a MoveDataParamEnv<'tcx>,
+    skip_unreachable_unwind: bool,
 }
 
 impl<'a, 'tcx> MaybeInitializedPlaces<'a, 'tcx> {
     pub fn new(tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>, mdpe: &'a MoveDataParamEnv<'tcx>) -> Self {
-        MaybeInitializedPlaces { tcx, body, mdpe }
+        MaybeInitializedPlaces { tcx, body, mdpe, skip_unreachable_unwind: false }
+    }
+
+    pub fn skipping_unreachable_unwind(mut self) -> Self {
+        self.skip_unreachable_unwind = true;
+        self
+    }
+
+    pub fn is_unwind_dead(
+        &self,
+        place: mir::Place<'tcx>,
+        state: &MaybeUnreachable<ChunkedBitSet<MovePathIndex>>,
+    ) -> bool {
+        if let LookupResult::Exact(path) = self.move_data().rev_lookup.find(place.as_ref()) {
+            let mut maybe_live = false;
+            on_all_drop_children_bits(self.tcx, self.body, self.mdpe, path, |child| {
+                maybe_live |= state.contains(child);
+            });
+            !maybe_live
+        } else {
+            false
+        }
     }
 }
 
@@ -107,11 +129,18 @@ pub struct MaybeUninitializedPlaces<'a, 'tcx> {
     mdpe: &'a MoveDataParamEnv<'tcx>,
 
     mark_inactive_variants_as_uninit: bool,
+    skip_unreachable_unwind: BitSet<mir::BasicBlock>,
 }
 
 impl<'a, 'tcx> MaybeUninitializedPlaces<'a, 'tcx> {
     pub fn new(tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>, mdpe: &'a MoveDataParamEnv<'tcx>) -> Self {
-        MaybeUninitializedPlaces { tcx, body, mdpe, mark_inactive_variants_as_uninit: false }
+        MaybeUninitializedPlaces {
+            tcx,
+            body,
+            mdpe,
+            mark_inactive_variants_as_uninit: false,
+            skip_unreachable_unwind: BitSet::new_empty(body.basic_blocks.len()),
+        }
     }
 
     /// Causes inactive enum variants to be marked as "maybe uninitialized" after a switch on an
@@ -123,6 +152,14 @@ impl<'a, 'tcx> MaybeUninitializedPlaces<'a, 'tcx> {
         self.mark_inactive_variants_as_uninit = true;
         self
     }
+
+    pub fn skipping_unreachable_unwind(
+        mut self,
+        unreachable_unwind: BitSet<mir::BasicBlock>,
+    ) -> Self {
+        self.skip_unreachable_unwind = unreachable_unwind;
+        self
+    }
 }
 
 impl<'a, 'tcx> HasMoveData<'tcx> for MaybeUninitializedPlaces<'a, 'tcx> {
@@ -271,18 +308,21 @@ impl<'a, 'tcx> DefinitelyInitializedPlaces<'a, 'tcx> {
 }
 
 impl<'tcx> AnalysisDomain<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
-    type Domain = ChunkedBitSet<MovePathIndex>;
+    type Domain = MaybeUnreachable<ChunkedBitSet<MovePathIndex>>;
     const NAME: &'static str = "maybe_init";
 
     fn bottom_value(&self, _: &mir::Body<'tcx>) -> Self::Domain {
         // bottom = uninitialized
-        ChunkedBitSet::new_empty(self.move_data().move_paths.len())
+        MaybeUnreachable::Unreachable
     }
 
     fn initialize_start_block(&self, _: &mir::Body<'tcx>, state: &mut Self::Domain) {
+        *state = MaybeUnreachable::Reachable(ChunkedBitSet::new_empty(
+            self.move_data().move_paths.len(),
+        ));
         drop_flag_effects_for_function_entry(self.tcx, self.body, self.mdpe, |path, s| {
             assert!(s == DropFlagState::Present);
-            state.insert(path);
+            state.gen(path);
         });
     }
 }
@@ -324,10 +364,18 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
         terminator: &'mir mir::Terminator<'tcx>,
         location: Location,
     ) -> TerminatorEdge<'mir, 'tcx> {
+        let mut edges = terminator.edges();
+        if self.skip_unreachable_unwind
+            && let mir::TerminatorKind::Drop { target, unwind, place, replace: _ } = terminator.kind
+            && matches!(unwind, mir::UnwindAction::Cleanup(_))
+            && self.is_unwind_dead(place, state)
+        {
+            edges = TerminatorEdge::Single(target);
+        }
         drop_flag_effects_for_location(self.tcx, self.body, self.mdpe, location, |path, s| {
             Self::update_bits(state, path, s)
         });
-        terminator.edges()
+        edges
     }
 
     fn call_return_effect(
@@ -448,7 +496,13 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
         drop_flag_effects_for_location(self.tcx, self.body, self.mdpe, location, |path, s| {
             Self::update_bits(trans, path, s)
         });
-        terminator.edges()
+        if self.skip_unreachable_unwind.contains(location.block) {
+            let mir::TerminatorKind::Drop { target, unwind, .. } = terminator.kind else { bug!() };
+            assert!(matches!(unwind, mir::UnwindAction::Cleanup(_)));
+            TerminatorEdge::Single(target)
+        } else {
+            terminator.edges()
+        }
     }
 
     fn call_return_effect(

compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -48,6 +48,7 @@ use std::fmt;
 pub struct ElaborateDrops;
 
 impl<'tcx> MirPass<'tcx> for ElaborateDrops {
+    #[instrument(level = "trace", skip(self, tcx, body))]
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         debug!("elaborate_drops({:?} @ {:?})", body.source, body.span);
 
@@ -65,23 +66,29 @@ impl<'tcx> MirPass<'tcx> for ElaborateDrops {
         };
         let elaborate_patch = {
             let env = MoveDataParamEnv { move_data, param_env };
-            remove_dead_unwinds(tcx, body, &env);
 
-            let inits = MaybeInitializedPlaces::new(tcx, body, &env)
+            let mut inits = MaybeInitializedPlaces::new(tcx, body, &env)
+                .skipping_unreachable_unwind()
                 .into_engine(tcx, body)
                 .pass_name("elaborate_drops")
                 .iterate_to_fixpoint()
                 .into_results_cursor(body);
+            let dead_unwinds = compute_dead_unwinds(&body, &mut inits);
+            let mut reachable = BitSet::new_empty(body.basic_blocks.len());
+            for bb in body.basic_blocks.indices() {
+                if inits.results().entry_set_for_block(bb).is_reachable() {
+                    reachable.insert(bb);
+                }
+            }
 
             let uninits = MaybeUninitializedPlaces::new(tcx, body, &env)
                 .mark_inactive_variants_as_uninit()
+                .skipping_unreachable_unwind(dead_unwinds)
                 .into_engine(tcx, body)
                 .pass_name("elaborate_drops")
                 .iterate_to_fixpoint()
                 .into_results_cursor(body);
 
-            let reachable = traversal::reachable_as_bitset(body);
-
             let drop_flags = IndexVec::from_elem(None, &env.move_data.move_paths);
             ElaborateDropsCtxt {
                 tcx,
@@ -101,63 +108,28 @@ impl<'tcx> MirPass<'tcx> for ElaborateDrops {
 
 /// Removes unwind edges which are known to be unreachable, because they are in `drop` terminators
 /// that can't drop anything.
-fn remove_dead_unwinds<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    body: &mut Body<'tcx>,
-    env: &MoveDataParamEnv<'tcx>,
-) {
-    debug!("remove_dead_unwinds({:?})", body.span);
+#[instrument(level = "trace", skip(body, flow_inits), ret)]
+fn compute_dead_unwinds<'mir, 'tcx>(
+    body: &'mir Body<'tcx>,
+    flow_inits: &mut ResultsCursor<'mir, 'tcx, MaybeInitializedPlaces<'mir, 'tcx>>,
+) -> BitSet<BasicBlock> {
     // We only need to do this pass once, because unwind edges can only
     // reach cleanup blocks, which can't have unwind edges themselves.
-    let mut dead_unwinds = Vec::new();
-    let mut flow_inits = MaybeInitializedPlaces::new(tcx, body, &env)
-        .into_engine(tcx, body)
-        .pass_name("remove_dead_unwinds")
-        .iterate_to_fixpoint()
-        .into_results_cursor(body);
+    let mut dead_unwinds = BitSet::new_empty(body.basic_blocks.len());
     for (bb, bb_data) in body.basic_blocks.iter_enumerated() {
-        let place = match bb_data.terminator().kind {
-            TerminatorKind::Drop { place, unwind: UnwindAction::Cleanup(_), .. } => place,
-            _ => continue,
-        };
-
-        debug!("remove_dead_unwinds @ {:?}: {:?}", bb, bb_data);
-
-        let LookupResult::Exact(path) = env.move_data.rev_lookup.find(place.as_ref()) else {
-            debug!("remove_dead_unwinds: has parent; skipping");
+        let TerminatorKind::Drop { place, unwind: UnwindAction::Cleanup(_), .. } =
+            bb_data.terminator().kind
+        else {
             continue;
         };
 
         flow_inits.seek_before_primary_effect(body.terminator_loc(bb));
-        debug!(
-            "remove_dead_unwinds @ {:?}: path({:?})={:?}; init_data={:?}",
-            bb,
-            place,
-            path,
-            flow_inits.get()
-        );
-
-        let mut maybe_live = false;
-        on_all_drop_children_bits(tcx, body, &env, path, |child| {
-            maybe_live |= flow_inits.contains(child);
-        });
-
-        debug!("remove_dead_unwinds @ {:?}: maybe_live={}", bb, maybe_live);
-        if !maybe_live {
-            dead_unwinds.push(bb);
+        if flow_inits.analysis().is_unwind_dead(place, flow_inits.get()) {
+            dead_unwinds.insert(bb);
         }
     }
 
-    if dead_unwinds.is_empty() {
-        return;
-    }
-
-    let basic_blocks = body.basic_blocks.as_mut();
-    for &bb in dead_unwinds.iter() {
-        if let Some(unwind) = basic_blocks[bb].terminator_mut().unwind_mut() {
-            *unwind = UnwindAction::Unreachable;
-        }
-    }
+    dead_unwinds
 }
 
 struct InitializationData<'mir, 'tcx> {

compiler/rustc_mir_transform/src/remove_uninit_drops.rs

@@ -4,7 +4,9 @@ use rustc_middle::ty::GenericArgsRef;
 use rustc_middle::ty::{self, ParamEnv, Ty, TyCtxt, VariantDef};
 use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
 use rustc_mir_dataflow::move_paths::{LookupResult, MoveData, MovePathIndex};
-use rustc_mir_dataflow::{self, move_path_children_matching, Analysis, MoveDataParamEnv};
+use rustc_mir_dataflow::{
+    self, move_path_children_matching, Analysis, MaybeUnreachable, MoveDataParamEnv,
+};
 use rustc_target::abi::FieldIdx;
 
 use crate::MirPass;
@@ -41,6 +43,7 @@ impl<'tcx> MirPass<'tcx> for RemoveUninitDrops {
             let TerminatorKind::Drop { place, .. } = &terminator.kind else { continue };
 
             maybe_inits.seek_before_primary_effect(body.terminator_loc(bb));
+            let MaybeUnreachable::Reachable(maybe_inits) = maybe_inits.get() else { continue };
 
             // If there's no move path for the dropped place, it's probably a `Deref`. Let it alone.
             let LookupResult::Exact(mpi) = mdpe.move_data.rev_lookup.find(place.as_ref()) else {
@@ -50,7 +53,7 @@ impl<'tcx> MirPass<'tcx> for RemoveUninitDrops {
             let should_keep = is_needs_drop_and_init(
                 tcx,
                 param_env,
-                maybe_inits.get(),
+                maybe_inits,
                 &mdpe.move_data,
                 place.ty(body, tcx).ty,
                 mpi,

tests/mir-opt/basic_assignment.main.ElaborateDrops.diff

@@ -80,10 +80,6 @@
 
       bb8 (cleanup): {
           resume;
-+     }
-+ 
-+     bb9 (cleanup): {
-+         unreachable;
       }
   }
 

tests/mir-opt/issue_41110.test.ElaborateDrops.panic-abort.diff

@@ -80,23 +80,19 @@
 
       bb9 (cleanup): {
 -         drop(_1) -> [return: bb10, unwind terminate];
-+         goto -> bb13;
++         goto -> bb12;
       }
 
       bb10 (cleanup): {
           resume;
 +     }
 + 
 +     bb11 (cleanup): {
-+         unreachable;
-+     }
-+ 
-+     bb12 (cleanup): {
 +         drop(_1) -> [return: bb10, unwind terminate];
 +     }
 + 
-+     bb13 (cleanup): {
-+         switchInt(_6) -> [0: bb10, otherwise: bb12];
++     bb12 (cleanup): {
++         switchInt(_6) -> [0: bb10, otherwise: bb11];
       }
   }
 

tests/mir-opt/issue_41110.test.ElaborateDrops.panic-unwind.diff

@@ -80,23 +80,19 @@
 
       bb9 (cleanup): {
 -         drop(_1) -> [return: bb10, unwind terminate];
-+         goto -> bb13;
++         goto -> bb12;
       }
 
       bb10 (cleanup): {
           resume;
 +     }
 + 
 +     bb11 (cleanup): {
-+         unreachable;
-+     }
-+ 
-+     bb12 (cleanup): {
 +         drop(_1) -> [return: bb10, unwind terminate];
 +     }
 + 
-+     bb13 (cleanup): {
-+         switchInt(_6) -> [0: bb10, otherwise: bb12];
++     bb12 (cleanup): {
++         switchInt(_6) -> [0: bb10, otherwise: bb11];
       }
   }