Remove uninit drops after building MIR

Cargo.lock

@@ -3764,6 +3764,7 @@ dependencies = [
  "rustc_span",
  "rustc_trait_selection",
  "smallvec",
+ "thin-vec",
  "tracing",
 ]
 

compiler/rustc_borrowck/Cargo.toml

@@ -23,5 +23,6 @@ rustc_session = { path = "../rustc_session" }
 rustc_span = { path = "../rustc_span" }
 rustc_trait_selection = { path = "../rustc_trait_selection" }
 smallvec = { version = "1.8.1", features = ["union", "may_dangle"] }
+thin-vec = "0.2.18"
 tracing = "0.1"
 # tidy-alphabetical-end

compiler/rustc_borrowck/src/renumber.rs

@@ -1,9 +1,10 @@
-use rustc_index::IndexSlice;
+use rustc_index::{IndexSlice, IndexVec};
 use rustc_infer::infer::NllRegionVariableOrigin;
 use rustc_middle::mir::visit::{MutVisitor, TyContext};
-use rustc_middle::mir::{Body, ConstOperand, Location, Promoted};
+use rustc_middle::mir::*;
 use rustc_middle::ty::{self, GenericArgsRef, Ty, TyCtxt, TypeFoldable, fold_regions};
 use rustc_span::Symbol;
+use thin_vec::ThinVec;
 use tracing::{debug, instrument};
 
 use crate::BorrowckInferCtxt;
@@ -21,12 +22,51 @@ pub(crate) fn renumber_mir<'tcx>(
     let mut renumberer = RegionRenumberer { infcx };
 
     for body in promoted.iter_mut() {
+        split_critical_unwind_edges(body);
         renumberer.visit_body_preserves_cfg(body);
     }
 
+    split_critical_unwind_edges(body);
     renumberer.visit_body_preserves_cfg(body);
 }
 
+#[instrument(skip(body), level = "debug")]
+fn split_critical_unwind_edges(body: &mut Body<'_>) {
+    let predecessors: IndexVec<BasicBlock, _> =
+        body.basic_blocks.predecessors().iter().map(|preds| preds.len()).collect();
+    debug!(?predecessors);
+
+    let mut new_blocks = vec![];
+    for bb in predecessors.indices() {
+        let term = body.basic_blocks[bb].terminator();
+        let Some(&UnwindAction::Cleanup(unwind)) = term.unwind() else { continue };
+        if predecessors[unwind] <= 1 {
+            continue;
+        }
+
+        debug!("{bb:?} has critical unwind edge: {unwind:?}");
+        new_blocks.push((bb, unwind));
+    }
+
+    if new_blocks.is_empty() {
+        return;
+    }
+
+    debug!(?new_blocks);
+    let basic_blocks = body.basic_blocks.as_mut();
+    for (bb, target) in new_blocks {
+        let source_info = basic_blocks[bb].terminator().source_info;
+        let terminator = Terminator {
+            source_info,
+            kind: TerminatorKind::Goto { target },
+            attributes: ThinVec::new(),
+        };
+        let new_target = basic_blocks.push(BasicBlockData::new(Some(terminator), true));
+        *basic_blocks[bb].terminator_mut().unwind_mut().unwrap() =
+            UnwindAction::Cleanup(new_target);
+    }
+}
+
 // The fields are used only for debugging output in `sccs_info`.
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
 pub(crate) enum RegionCtxt {

compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -3,6 +3,7 @@ use std::fmt;
 use rustc_abi::{FieldIdx, VariantIdx};
 use rustc_index::IndexVec;
 use rustc_index::bit_set::DenseBitSet;
+use rustc_middle::mir::traversal::reachable_as_bitset;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, TyCtxt};
 use rustc_mir_dataflow::impls::{MaybeInitializedPlaces, MaybeUninitializedPlaces};
@@ -58,6 +59,7 @@ impl<'tcx> crate::MirPass<'tcx> for ElaborateDrops {
         // For types that do not need dropping, the behaviour is trivial. So we only need to track
         // init/uninit for types that do need dropping.
         let move_data = MoveData::gather_moves(body, tcx, |ty| ty.needs_drop(tcx, typing_env));
+        let reachable = reachable_as_bitset(body);
         let elaborate_patch = {
             let env = MoveDataTypingEnv { move_data, typing_env };
 
@@ -66,7 +68,7 @@ impl<'tcx> crate::MirPass<'tcx> for ElaborateDrops {
                 .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 dead_unwinds = compute_dead_unwinds(body, &reachable, &mut inits);
 
             let uninits = MaybeUninitializedPlaces::new(tcx, body, &env.move_data)
                 .mark_inactive_variants_as_uninit()
@@ -79,6 +81,7 @@ impl<'tcx> crate::MirPass<'tcx> for ElaborateDrops {
                 tcx,
                 body,
                 env: &env,
+                reachable: &reachable,
                 init_data: InitializationData { inits, uninits },
                 drop_flags,
                 patch: MirPatch::new(body),
@@ -99,12 +102,14 @@ impl<'tcx> crate::MirPass<'tcx> for ElaborateDrops {
 #[instrument(level = "trace", skip(body, flow_inits), ret)]
 fn compute_dead_unwinds<'a, 'tcx>(
     body: &'a Body<'tcx>,
+    reachable: &DenseBitSet<BasicBlock>,
     flow_inits: &mut ResultsCursor<'a, 'tcx, MaybeInitializedPlaces<'a, 'tcx>>,
 ) -> DenseBitSet<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 = DenseBitSet::new_empty(body.basic_blocks.len());
-    for (bb, bb_data) in body.basic_blocks.iter_enumerated() {
+    for bb in reachable.iter() {
+        let bb_data = &body[bb];
         let TerminatorKind::Drop { place, unwind: UnwindAction::Cleanup(_), .. } =
             bb_data.terminator().kind
         else {
@@ -248,6 +253,7 @@ struct ElaborateDropsCtxt<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     body: &'a Body<'tcx>,
     env: &'a MoveDataTypingEnv<'tcx>,
+    reachable: &'a DenseBitSet<BasicBlock>,
     init_data: InitializationData<'a, 'tcx>,
     drop_flags: IndexVec<MovePathIndex, Option<Local>>,
     patch: MirPatch<'tcx>,
@@ -290,7 +296,8 @@ impl<'a, 'tcx> ElaborateDropsCtxt<'a, 'tcx> {
     }
 
     fn collect_drop_flags(&mut self) {
-        for (bb, data) in self.body.basic_blocks.iter_enumerated() {
+        for bb in self.reachable.iter() {
+            let data = &self.body[bb];
             let terminator = data.terminator();
             let TerminatorKind::Drop { ref place, .. } = terminator.kind else { continue };
 
@@ -337,7 +344,8 @@ impl<'a, 'tcx> ElaborateDropsCtxt<'a, 'tcx> {
 
     fn elaborate_drops(&mut self) {
         // This function should mirror what `collect_drop_flags` does.
-        for (bb, data) in self.body.basic_blocks.iter_enumerated() {
+        for bb in self.reachable.iter() {
+            let data = &self.body[bb];
             let terminator = data.terminator();
             let TerminatorKind::Drop { place, target, unwind, replace, drop } = terminator.kind
             else {

compiler/rustc_mir_transform/src/lib.rs

@@ -412,6 +412,7 @@ fn mir_built(tcx: TyCtxt<'_>, def: LocalDefId) -> &Steal<Body<'_>> {
             // This used to be part of MIR building,
             // now done separately to separate concerns.
             &lint_and_remove_uninhabited::LintAndRemoveUninhabited,
+            &remove_uninit_drops::RemoveUninitDrops,
             // MIR-level lints.
             &Lint(check_inline::CheckForceInline),
             &Lint(check_call_recursion::CheckCallRecursion),

compiler/rustc_mir_transform/src/remove_uninit_drops.rs

@@ -1,5 +1,6 @@
 use rustc_abi::FieldIdx;
 use rustc_index::bit_set::MixedBitSet;
+use rustc_middle::mir::traversal::reachable;
 use rustc_middle::mir::{Body, TerminatorKind};
 use rustc_middle::ty::{self, GenericArgsRef, Ty, TyCtxt, VariantDef};
 use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
@@ -29,7 +30,7 @@ impl<'tcx> crate::MirPass<'tcx> for RemoveUninitDrops {
             .into_results_cursor(body);
 
         let mut to_remove = vec![];
-        for (bb, block) in body.basic_blocks.iter_enumerated() {
+        for (bb, block) in reachable(body) {
             let terminator = block.terminator();
             let TerminatorKind::Drop { place, .. } = &terminator.kind else { continue };
 

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

@@ -35,53 +35,23 @@
           StorageLive(_5);
           StorageLive(_6);
           _6 = move _4;
--         drop(_5) -> [return: bb1, unwind: bb2];
-+         goto -> bb1;
-      }
-
-      bb1: {
           _5 = move _6;
--         drop(_6) -> [return: bb3, unwind: bb6];
-+         goto -> bb3;
-      }
-
-      bb2 (cleanup): {
-          _5 = move _6;
--         drop(_6) -> [return: bb6, unwind terminate(cleanup)];
-+         goto -> bb6;
-      }
-
-      bb3: {
           StorageDead(_6);
           _0 = const ();
-          drop(_5) -> [return: bb4, unwind: bb7];
+-         drop(_5) -> [return: bb1, unwind continue];
++         drop(_5) -> [return: bb1, unwind: bb2];
       }
 
-      bb4: {
+      bb1: {
           StorageDead(_5);
--         drop(_4) -> [return: bb5, unwind continue];
-+         goto -> bb5;
-      }
-
-      bb5: {
           StorageDead(_4);
           StorageDead(_2);
           StorageDead(_1);
           return;
-      }
-
-      bb6 (cleanup): {
--         drop(_5) -> [return: bb7, unwind terminate(cleanup)];
-+         goto -> bb7;
-      }
-
-      bb7 (cleanup): {
--         drop(_4) -> [return: bb8, unwind terminate(cleanup)];
-+         goto -> bb8;
-      }
-
-      bb8 (cleanup): {
-          resume;
++     }
++ 
++     bb2 (cleanup): {
++         resume;
       }
   }
 

tests/mir-opt/basic_assignment.main.SimplifyCfg-initial.after.mir

@@ -41,46 +41,21 @@ fn main() -> () {
         StorageLive(_5);
         StorageLive(_6);
         _6 = move _4;
-        drop(_5) -> [return: bb1, unwind: bb2];
-    }
-
-    bb1: {
-        _5 = move _6;
-        drop(_6) -> [return: bb3, unwind: bb6];
-    }
-
-    bb2 (cleanup): {
         _5 = move _6;
-        drop(_6) -> [return: bb6, unwind terminate(cleanup)];
-    }
-
-    bb3: {
         StorageDead(_6);
         _0 = const ();
-        drop(_5) -> [return: bb4, unwind: bb7];
+        drop(_5) -> [return: bb1, unwind: bb2];
     }
 
-    bb4: {
+    bb1: {
         StorageDead(_5);
-        drop(_4) -> [return: bb5, unwind: bb8];
-    }
-
-    bb5: {
         StorageDead(_4);
         StorageDead(_2);
         StorageDead(_1);
         return;
     }
 
-    bb6 (cleanup): {
-        drop(_5) -> [return: bb7, unwind terminate(cleanup)];
-    }
-
-    bb7 (cleanup): {
-        drop(_4) -> [return: bb8, unwind terminate(cleanup)];
-    }
-
-    bb8 (cleanup): {
+    bb2 (cleanup): {
         resume;
     }
 }