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/coroutine.rs

@@ -1392,7 +1392,7 @@ pub(crate) fn mir_coroutine_witnesses<'tcx>(
     tcx: TyCtxt<'tcx>,
     def_id: LocalDefId,
 ) -> Option<CoroutineLayout<'tcx>> {
-    let (body, _) = tcx.mir_promoted(def_id);
+    let body = tcx.mir_post_borrowck_cleanup(def_id);
     let body = body.borrow();
     let body = &*body;
 

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,
@@ -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.borrow().clone();
+    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");
@@ -543,7 +558,7 @@ fn mir_drops_elaborated_and_const_checked(tcx: TyCtxt<'_>, def: LocalDefId) -> &
 
     tcx.ensure_done().check_liveness(def);
 
-    let (body, _) = tcx.mir_promoted(def);
+    let body = tcx.mir_post_borrowck_cleanup(def);
     let mut body = body.steal();
 
     if let Some(error_reported) = tainted_by_errors {

compiler/rustc_mir_transform/src/remove_dead_drops.rs

@@ -0,0 +1,62 @@
+use rustc_middle::mir::*;
+use rustc_middle::ty::TyCtxt;
+use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
+use rustc_mir_dataflow::move_paths::{LookupResult, MoveData};
+use rustc_mir_dataflow::{Analysis, MaybeReachable};
+
+use super::simplify::simplify_cfg;
+
+pub(crate) struct RemoveDeadDrops;
+
+impl<'tcx> crate::MirPass<'tcx> for RemoveDeadDrops {
+    fn is_required(&self) -> bool {
+        true
+    }
+
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        if body.coroutine.is_none() {
+            return;
+        }
+
+        let move_data = MoveData::gather_moves(body, tcx, |_| true);
+
+        let mut maybe_init_cursor = MaybeInitializedPlaces::new(tcx, body, &move_data)
+            .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 LookupResult::Exact(path) = move_data.rev_lookup.find(place.as_ref()) else {
+                    continue;
+                };
+
+                let term_location = Location { block, statement_index: data.statements.len() };
+                maybe_init_cursor.seek_before_primary_effect(term_location);
+
+                let is_dead = match maybe_init_cursor.get() {
+                    MaybeReachable::Unreachable => true,
+                    MaybeReachable::Reachable(maybe_init) => !maybe_init.contains(path),
+                };
+
+                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/building/async_await.b-{closure#0}.coroutine_resume.0.mir

@@ -105,7 +105,7 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
     bb0: {
         _39 = copy (_1.0: &mut {async fn body of b()});
         _38 = discriminant((*_39));
-        switchInt(move _38) -> [0: bb1, 1: bb29, 3: bb27, 4: bb28, otherwise: bb8];
+        switchInt(move _38) -> [0: bb1, 1: bb27, 3: bb25, 4: bb26, otherwise: bb8];
     }
 
     bb1: {
@@ -206,65 +206,61 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
 
     bb12: {
         nop;
-        goto -> bb13;
-    }
-
-    bb13: {
         StorageDead(_4);
         StorageDead(_3);
         StorageLive(_21);
         StorageLive(_22);
-        _22 = a() -> [return: bb14, unwind unreachable];
+        _22 = a() -> [return: bb13, unwind unreachable];
     }
 
-    bb14: {
-        _21 = <{async fn body of a()} as IntoFuture>::into_future(move _22) -> [return: bb15, unwind unreachable];
+    bb13: {
+        _21 = <{async fn body of a()} as IntoFuture>::into_future(move _22) -> [return: bb14, unwind unreachable];
     }
 
-    bb15: {
+    bb14: {
         StorageDead(_22);
         PlaceMention(_21);
         nop;
         (((*_39) as variant#4).0: {async fn body of a()}) = move _21;
-        goto -> bb16;
+        goto -> bb15;
     }
 
-    bb16: {
+    bb15: {
         StorageLive(_24);
         StorageLive(_25);
         StorageLive(_26);
         StorageLive(_27);
         StorageLive(_28);
         _28 = &mut (((*_39) as variant#4).0: {async fn body of a()});
         _27 = &mut (*_28);
-        _26 = Pin::<&mut {async fn body of a()}>::new_unchecked(move _27) -> [return: bb17, unwind unreachable];
+        _26 = Pin::<&mut {async fn body of a()}>::new_unchecked(move _27) -> [return: bb16, unwind unreachable];
     }
 
-    bb17: {
+    bb16: {
         StorageDead(_27);
         StorageLive(_29);
         StorageLive(_30);
         StorageLive(_31);
         _31 = copy _2;
         _30 = move _31;
-        goto -> bb18;
+        goto -> bb17;
     }
 
-    bb18: {
+    bb17: {
         _29 = &mut (*_30);
         StorageDead(_31);
-        _25 = <{async fn body of a()} as Future>::poll(move _26, move _29) -> [return: bb19, unwind unreachable];
+        _25 = <{async fn body of a()} as Future>::poll(move _26, move _29) -> [return: bb18, unwind unreachable];
     }
 
-    bb19: {
+    bb18: {
         StorageDead(_29);
         StorageDead(_26);
         PlaceMention(_25);
         _32 = discriminant(_25);
-        switchInt(move _32) -> [0: bb21, 1: bb20, otherwise: bb8];
+        switchInt(move _32) -> [0: bb20, 1: bb19, otherwise: bb8];
     }
 
-    bb20: {
+    bb19: {
         _24 = const ();
         StorageDead(_30);
         StorageDead(_28);
@@ -281,7 +277,7 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
         return;
     }
 
-    bb21: {
+    bb20: {
         StorageLive(_33);
         _33 = copy ((_25 as Ready).0: ());
         _37 = copy _33;
@@ -290,38 +286,34 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
         StorageDead(_28);
         StorageDead(_25);
         StorageDead(_24);
-        drop((((*_39) as variant#4).0: {async fn body of a()})) -> [return: bb23, unwind unreachable];
+        drop((((*_39) as variant#4).0: {async fn body of a()})) -> [return: bb22, unwind unreachable];
     }
 
-    bb22: {
+    bb21: {
         StorageDead(_36);
         _2 = move _35;
         StorageDead(_35);
         _7 = const ();
-        goto -> bb16;
+        goto -> bb15;
     }
 
-    bb23: {
+    bb22: {
         nop;
-        goto -> bb24;
-    }
-
-    bb24: {
         StorageDead(_21);
-        goto -> bb26;
+        goto -> bb24;
     }
 
-    bb25: {
+    bb23: {
         _0 = Poll::<()>::Ready(move _37);
         discriminant((*_39)) = 1;
         return;
     }
 
-    bb26: {
-        goto -> bb25;
+    bb24: {
+        goto -> bb23;
     }
 
-    bb27: {
+    bb25: {
         StorageLive(_3);
         StorageLive(_4);
         StorageLive(_19);
@@ -330,15 +322,15 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
         goto -> bb11;
     }
 
-    bb28: {
+    bb26: {
         StorageLive(_21);
         StorageLive(_35);
         StorageLive(_36);
         _35 = move _2;
-        goto -> bb22;
+        goto -> bb21;
     }
 
-    bb29: {
-        assert(const false, "`async fn` resumed after completion") -> [success: bb29, unwind unreachable];
+    bb27: {
+        assert(const false, "`async fn` resumed after completion") -> [success: bb27, unwind unreachable];
     }
 }