Remove redundant drop terminators ahead of mir_coroutine_witnesses

compiler/rustc_middle/src/queries.rs

@@ -672,6 +672,11 @@ rustc_queries! {
         desc { "promoting constants in MIR for `{}`", tcx.def_path_str(key) }
     }
 
+    query mir_post_borrowck_cleanup(key: LocalDefId) -> &'tcx Steal<mir::Body<'tcx>> {
+        no_hash
+        desc { "post borrowck cleanup of MIR for `{}`", tcx.def_path_str(key) }
+    }
+
     query closure_typeinfo(key: LocalDefId) -> ty::ClosureTypeInfo<'tcx> {
         desc {
             "finding symbols for captures of closure `{}`",

compiler/rustc_mir_transform/src/lib.rs

@@ -171,6 +171,7 @@ declare_passes! {
     mod promote_consts : PromoteTemps;
     mod ref_prop : ReferencePropagation;
     mod remove_noop_landing_pads : RemoveNoopLandingPads;
+    mod remove_dead_drops : RemoveDeadDrops;
     mod remove_place_mention : RemovePlaceMention;
     mod remove_storage_markers : RemoveStorageMarkers;
     mod remove_uninit_drops : RemoveUninitDrops;
@@ -221,6 +222,7 @@ pub fn provide(providers: &mut Providers) {
         mir_built,
         mir_const_qualif,
         mir_promoted,
+        mir_post_borrowck_cleanup,
         mir_drops_elaborated_and_const_checked,
         mir_for_ctfe,
         mir_coroutine_witnesses: coroutine::mir_coroutine_witnesses,
@@ -474,7 +476,7 @@ fn mir_promoted(
     pm::run_passes(
         tcx,
         &mut body,
-        &[&promote_pass, &simplify::SimplifyCfg::PromoteConsts, &coverage::InstrumentCoverage],
+        &[&promote_pass, &simplify::SimplifyCfg::PromoteConsts, &coverage::InstrumentCoverage, &remove_dead_drops::RemoveDeadDrops],
         Some(MirPhase::Analysis(AnalysisPhase::Initial)),
         pm::Optimizations::Allowed,
     );
@@ -485,6 +487,19 @@ fn mir_promoted(
     (tcx.alloc_steal_mir(body), tcx.alloc_steal_promoted(promoted))
 }
 
+fn mir_post_borrowck_cleanup(tcx: TyCtxt<'_>, def: LocalDefId) -> &Steal<Body<'_>> {
+    let (body, _) = tcx.mir_promoted(def);
+    let mut body = body.steal();
+    pm::run_passes(
+        tcx,
+        &mut body,
+        &[&remove_dead_drops::RemoveDeadDrops],
+        None,
+        pm::Optimizations::Allowed,
+    );
+    tcx.alloc_steal_mir(body)
+}
+
 /// Compute the MIR that is used during CTFE (and thus has no optimizations run on it)
 fn mir_for_ctfe(tcx: TyCtxt<'_>, def_id: LocalDefId) -> &Body<'_> {
     debug_assert!(!tcx.is_trivial_const(def_id), "Tried to get mir_for_ctfe of a trivial const");

compiler/rustc_mir_transform/src/remove_dead_drops.rs

@@ -0,0 +1,117 @@
+use rustc_index::bit_set::DenseBitSet;
+use rustc_middle::mir::*;
+use rustc_middle::ty::TyCtxt;
+use rustc_mir_dataflow::{Analysis, GenKill};
+
+use super::simplify::simplify_cfg;
+
+pub(crate) struct RemoveDeadDrops;
+
+use rustc_middle::mir::visit::*;
+
+struct MaybeInitializedLocals;
+
+/// This is conservative: If any part of a local is initialized, we mark it
+/// initialized, while we only mark uninitialized if the whole local is moved
+/// from or StorageDead.
+struct TransferFunction<'a, T> {
+    trans: &'a mut T,
+}
+
+impl<'tcx, T: GenKill<Local>> Visitor<'tcx> for TransferFunction<'_, T> {
+    fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, location: Location) {
+        self.super_place(place, context, location);
+
+        if context.is_place_assignment() {
+            self.trans.gen_(place.local);
+        } else if matches!(
+            context,
+            PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead)
+        ) {
+            self.trans.kill(place.local);
+        }
+    }
+
+    fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
+        self.super_operand(operand, location);
+
+        if let Operand::Move(place) = operand {
+            if place.projection.is_empty() {
+                self.trans.kill(place.local);
+            }
+        }
+    }
+}
+
+impl<'tcx> Analysis<'tcx> for MaybeInitializedLocals {
+    type Domain = DenseBitSet<Local>;
+    const NAME: &'static str = "maybe_initialized_locals";
+
+    fn bottom_value(&self, body: &Body<'tcx>) -> Self::Domain {
+        DenseBitSet::new_empty(body.local_decls.len())
+    }
+
+    fn initialize_start_block(&self, body: &Body<'tcx>, state: &mut Self::Domain) {
+        for arg in 1..=body.arg_count {
+            state.insert(Local::from_usize(arg));
+        }
+    }
+
+    fn apply_primary_statement_effect(
+        &self,
+        state: &mut Self::Domain,
+        stmt: &Statement<'tcx>,
+        location: Location,
+    ) {
+        TransferFunction { trans: state }.visit_statement(stmt, location);
+    }
+
+    fn apply_primary_terminator_effect<'mir>(
+        &self,
+        state: &mut Self::Domain,
+        terminator: &'mir Terminator<'tcx>,
+        location: Location,
+    ) -> TerminatorEdges<'mir, 'tcx> {
+        TransferFunction { trans: state }.visit_terminator(terminator, location);
+        terminator.edges()
+    }
+}
+
+impl<'tcx> crate::MirPass<'tcx> for RemoveDeadDrops {
+    fn is_required(&self) -> bool {
+        true
+    }
+
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let mut maybe_init_cursor = MaybeInitializedLocals
+            .iterate_to_fixpoint(tcx, body, None)
+            .into_results_cursor(body);
+
+        let mut dead_drops = Vec::new();
+
+        for (block, data) in body.basic_blocks.iter_enumerated() {
+            if let Some(terminator) = &data.terminator
+                && let TerminatorKind::Drop { place, target, .. } = &terminator.kind
+            {
+                let term_location = Location { block, statement_index: data.statements.len() };
+                maybe_init_cursor.seek_before_primary_effect(term_location);
+
+                let is_dead = !maybe_init_cursor.get().contains(place.local);
+                if is_dead {
+                    dead_drops.push((block, *target));
+                }
+            }
+        }
+
+        if !dead_drops.is_empty() {
+            for (block, target) in dead_drops {
+                if let Some(terminator) = &mut body.basic_blocks.as_mut()[block].terminator {
+                    terminator.kind = TerminatorKind::Goto { target };
+                }
+            }
+
+            // Removing drop terminators may simplify the CFG, so run cleanup.
+            simplify_cfg(tcx, body);
+        }
+    }
+}

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