Async drop - fix for StorageLive/StorageDead codegen for pinned async drop future

10 files changed, 387 insertions, 9 deletions

diff --git a/compiler/rustc_mir_transform/src/coroutine/drop.rs b/compiler/rustc_mir_transform/src/coroutine/drop.rs
index 6021e795d21..dc68629ec0d 100644
--- a/compiler/rustc_mir_transform/src/coroutine/drop.rs
+++ b/compiler/rustc_mir_transform/src/coroutine/drop.rs
@@ -132,6 +132,7 @@ fn build_poll_switch<'tcx>(
     body: &mut Body<'tcx>,
     poll_enum: Ty<'tcx>,
     poll_unit_place: &Place<'tcx>,
+    fut_pin_place: &Place<'tcx>,
     ready_block: BasicBlock,
     yield_block: BasicBlock,
 ) -> BasicBlock {
@@ -162,9 +163,11 @@ fn build_poll_switch<'tcx>(
             Rvalue::Discriminant(*poll_unit_place),
         ))),
     };
+    let storage_dead =
+        Statement { source_info, kind: StatementKind::StorageDead(fut_pin_place.local) };
     let unreachable_block = insert_term_block(body, TerminatorKind::Unreachable);
     body.basic_blocks_mut().push(BasicBlockData {
-        statements: [discr_assign].to_vec(),
+        statements: [storage_dead, discr_assign].to_vec(),
         terminator: Some(Terminator {
             source_info,
             kind: TerminatorKind::SwitchInt {
@@ -332,10 +335,17 @@ pub(super) fn expand_async_drops<'tcx>(
             kind: StatementKind::Assign(Box::new((context_ref_place, arg))),
         });
         let yield_block = insert_term_block(body, TerminatorKind::Unreachable); // `kind` replaced later to yield
-        let switch_block =
-            build_poll_switch(tcx, body, poll_enum, &poll_unit_place, target, yield_block);
         let (pin_bb, fut_pin_place) =
             build_pin_fut(tcx, body, fut_place.clone(), UnwindAction::Continue);
+        let switch_block = build_poll_switch(
+            tcx,
+            body,
+            poll_enum,
+            &poll_unit_place,
+            &fut_pin_place,
+            target,
+            yield_block,
+        );
         let call_bb = build_poll_call(
             tcx,
             body,
@@ -357,16 +367,17 @@ pub(super) fn expand_async_drops<'tcx>(
                 body.local_decls.push(LocalDecl::new(context_mut_ref, source_info.span)),
             );
             let drop_yield_block = insert_term_block(body, TerminatorKind::Unreachable); // `kind` replaced later to yield
+            let (pin_bb2, fut_pin_place2) =
+                build_pin_fut(tcx, body, fut_place, UnwindAction::Continue);
             let drop_switch_block = build_poll_switch(
                 tcx,
                 body,
                 poll_enum,
                 &poll_unit_place,
+                &fut_pin_place2,
                 drop.unwrap(),
                 drop_yield_block,
             );
-            let (pin_bb2, fut_pin_place2) =
-                build_pin_fut(tcx, body, fut_place, UnwindAction::Continue);
             let drop_call_bb = build_poll_call(
                 tcx,
                 body,
diff --git a/compiler/rustc_mir_transform/src/elaborate_drop.rs b/compiler/rustc_mir_transform/src/elaborate_drop.rs
index 3a5e2620b14..c9bc52c6c7e 100644
--- a/compiler/rustc_mir_transform/src/elaborate_drop.rs
+++ b/compiler/rustc_mir_transform/src/elaborate_drop.rs
@@ -390,6 +390,20 @@ where
             Location { block: self.succ, statement_index: 0 },
             StatementKind::StorageDead(fut.local),
         );
+        // StorageDead(fut) in unwind block (at the begin)
+        if let Unwind::To(block) = unwind {
+            self.elaborator.patch().add_statement(
+                Location { block, statement_index: 0 },
+                StatementKind::StorageDead(fut.local),
+            );
+        }
+        // StorageDead(fut) in dropline block (at the begin)
+        if let Some(block) = dropline {
+            self.elaborator.patch().add_statement(
+                Location { block, statement_index: 0 },
+                StatementKind::StorageDead(fut.local),
+            );
+        }
 
         // #1:pin_obj_bb >>> call Pin<ObjTy>::new_unchecked(&mut obj)
         self.elaborator.patch().patch_terminator(
diff --git a/tests/mir-opt/async_drop_live_dead.a-{closure#0}.coroutine_drop_async.0.panic-abort.mir b/tests/mir-opt/async_drop_live_dead.a-{closure#0}.coroutine_drop_async.0.panic-abort.mir
new file mode 100644
index 00000000000..347e4119cd0
--- /dev/null
+++ b/tests/mir-opt/async_drop_live_dead.a-{closure#0}.coroutine_drop_async.0.panic-abort.mir
@@ -0,0 +1,100 @@
+// MIR for `a::{closure#0}` 0 coroutine_drop_async
+
+fn a::{closure#0}(_1: Pin<&mut {async fn body of a<T>()}>, _2: &mut Context<'_>) -> Poll<()> {
+    debug _task_context => _19;
+    debug x => ((*(_1.0: &mut {async fn body of a<T>()})).0: T);
+    let mut _0: std::task::Poll<()>;
+    let _3: T;
+    let mut _4: impl std::future::Future<Output = ()>;
+    let mut _5: &mut T;
+    let mut _6: std::pin::Pin<&mut T>;
+    let mut _7: &mut T;
+    let mut _8: *mut T;
+    let mut _9: ();
+    let mut _10: std::task::Poll<()>;
+    let mut _11: &mut std::task::Context<'_>;
+    let mut _12: &mut impl std::future::Future<Output = ()>;
+    let mut _13: std::pin::Pin<&mut impl std::future::Future<Output = ()>>;
+    let mut _14: isize;
+    let mut _15: &mut std::task::Context<'_>;
+    let mut _16: &mut impl std::future::Future<Output = ()>;
+    let mut _17: std::pin::Pin<&mut impl std::future::Future<Output = ()>>;
+    let mut _18: isize;
+    let mut _19: &mut std::task::Context<'_>;
+    let mut _20: u32;
+    scope 1 {
+        debug x => (((*(_1.0: &mut {async fn body of a<T>()})) as variant#4).0: T);
+    }
+
+    bb0: {
+        _20 = discriminant((*(_1.0: &mut {async fn body of a<T>()})));
+        switchInt(move _20) -> [0: bb9, 3: bb12, 4: bb13, otherwise: bb14];
+    }
+
+    bb1: {
+        nop;
+        nop;
+        goto -> bb2;
+    }
+
+    bb2: {
+        _0 = Poll::<()>::Ready(const ());
+        return;
+    }
+
+    bb3: {
+        _0 = Poll::<()>::Pending;
+        discriminant((*(_1.0: &mut {async fn body of a<T>()}))) = 4;
+        return;
+    }
+
+    bb4: {
+        StorageLive(_17);
+        _16 = &mut (((*(_1.0: &mut {async fn body of a<T>()})) as variant#4).1: impl std::future::Future<Output = ()>);
+        _17 = Pin::<&mut impl Future<Output = ()>>::new_unchecked(move _16) -> [return: bb7, unwind unreachable];
+    }
+
+    bb5: {
+        unreachable;
+    }
+
+    bb6: {
+        StorageDead(_17);
+        _18 = discriminant(_10);
+        switchInt(move _18) -> [0: bb1, 1: bb3, otherwise: bb5];
+    }
+
+    bb7: {
+        _10 = <impl Future<Output = ()> as Future>::poll(move _17, move _15) -> [return: bb6, unwind unreachable];
+    }
+
+    bb8: {
+        _0 = Poll::<()>::Ready(const ());
+        return;
+    }
+
+    bb9: {
+        goto -> bb11;
+    }
+
+    bb10: {
+        goto -> bb8;
+    }
+
+    bb11: {
+        drop(((*(_1.0: &mut {async fn body of a<T>()})).0: T)) -> [return: bb10, unwind unreachable];
+    }
+
+    bb12: {
+        goto -> bb4;
+    }
+
+    bb13: {
+        goto -> bb4;
+    }
+
+    bb14: {
+        _0 = Poll::<()>::Ready(const ());
+        return;
+    }
+}
diff --git a/tests/mir-opt/async_drop_live_dead.a-{closure#0}.coroutine_drop_async.0.panic-unwind.mir b/tests/mir-opt/async_drop_live_dead.a-{closure#0}.coroutine_drop_async.0.panic-unwind.mir
new file mode 100644
index 00000000000..b1cf5373f91
--- /dev/null
+++ b/tests/mir-opt/async_drop_live_dead.a-{closure#0}.coroutine_drop_async.0.panic-unwind.mir
@@ -0,0 +1,123 @@
+// MIR for `a::{closure#0}` 0 coroutine_drop_async
+
+fn a::{closure#0}(_1: Pin<&mut {async fn body of a<T>()}>, _2: &mut Context<'_>) -> Poll<()> {
+    debug _task_context => _19;
+    debug x => ((*(_1.0: &mut {async fn body of a<T>()})).0: T);
+    let mut _0: std::task::Poll<()>;
+    let _3: T;
+    let mut _4: impl std::future::Future<Output = ()>;
+    let mut _5: &mut T;
+    let mut _6: std::pin::Pin<&mut T>;
+    let mut _7: &mut T;
+    let mut _8: *mut T;
+    let mut _9: ();
+    let mut _10: std::task::Poll<()>;
+    let mut _11: &mut std::task::Context<'_>;
+    let mut _12: &mut impl std::future::Future<Output = ()>;
+    let mut _13: std::pin::Pin<&mut impl std::future::Future<Output = ()>>;
+    let mut _14: isize;
+    let mut _15: &mut std::task::Context<'_>;
+    let mut _16: &mut impl std::future::Future<Output = ()>;
+    let mut _17: std::pin::Pin<&mut impl std::future::Future<Output = ()>>;
+    let mut _18: isize;
+    let mut _19: &mut std::task::Context<'_>;
+    let mut _20: u32;
+    scope 1 {
+        debug x => (((*(_1.0: &mut {async fn body of a<T>()})) as variant#4).0: T);
+    }
+
+    bb0: {
+        _20 = discriminant((*(_1.0: &mut {async fn body of a<T>()})));
+        switchInt(move _20) -> [0: bb12, 2: bb18, 3: bb16, 4: bb17, otherwise: bb19];
+    }
+
+    bb1: {
+        nop;
+        nop;
+        goto -> bb2;
+    }
+
+    bb2: {
+        _0 = Poll::<()>::Ready(const ());
+        return;
+    }
+
+    bb3 (cleanup): {
+        nop;
+        nop;
+        goto -> bb5;
+    }
+
+    bb4 (cleanup): {
+        goto -> bb15;
+    }
+
+    bb5 (cleanup): {
+        goto -> bb4;
+    }
+
+    bb6: {
+        _0 = Poll::<()>::Pending;
+        discriminant((*(_1.0: &mut {async fn body of a<T>()}))) = 4;
+        return;
+    }
+
+    bb7: {
+        StorageLive(_17);
+        _16 = &mut (((*(_1.0: &mut {async fn body of a<T>()})) as variant#4).1: impl std::future::Future<Output = ()>);
+        _17 = Pin::<&mut impl Future<Output = ()>>::new_unchecked(move _16) -> [return: bb10, unwind: bb15];
+    }
+
+    bb8: {
+        unreachable;
+    }
+
+    bb9: {
+        StorageDead(_17);
+        _18 = discriminant(_10);
+        switchInt(move _18) -> [0: bb1, 1: bb6, otherwise: bb8];
+    }
+
+    bb10: {
+        _10 = <impl Future<Output = ()> as Future>::poll(move _17, move _15) -> [return: bb9, unwind: bb3];
+    }
+
+    bb11: {
+        _0 = Poll::<()>::Ready(const ());
+        return;
+    }
+
+    bb12: {
+        goto -> bb14;
+    }
+
+    bb13: {
+        goto -> bb11;
+    }
+
+    bb14: {
+        drop(((*(_1.0: &mut {async fn body of a<T>()})).0: T)) -> [return: bb13, unwind: bb4];
+    }
+
+    bb15 (cleanup): {
+        discriminant((*(_1.0: &mut {async fn body of a<T>()}))) = 2;
+        resume;
+    }
+
+    bb16: {
+        goto -> bb7;
+    }
+
+    bb17: {
+        goto -> bb7;
+    }
+
+    bb18: {
+        assert(const false, "`async fn` resumed after panicking") -> [success: bb18, unwind continue];
+    }
+
+    bb19: {
+        _0 = Poll::<()>::Ready(const ());
+        return;
+    }
+}
diff --git a/tests/mir-opt/async_drop_live_dead.rs b/tests/mir-opt/async_drop_live_dead.rs
new file mode 100644
index 00000000000..348866bbb8c
--- /dev/null
+++ b/tests/mir-opt/async_drop_live_dead.rs
@@ -0,0 +1,11 @@
+//@ edition:2024
+// skip-filecheck
+// EMIT_MIR_FOR_EACH_PANIC_STRATEGY
+
+#![feature(async_drop)]
+#![allow(incomplete_features)]
+
+// EMIT_MIR async_drop_live_dead.a-{closure#0}.coroutine_drop_async.0.mir
+async fn a<T>(x: T) {}
+
+fn main() {}
diff --git a/tests/ui/async-await/async-drop/async-drop-initial.rs b/tests/ui/async-await/async-drop/async-drop-initial.rs
index 263b70699f5..cd33c143fba 100644
--- a/tests/ui/async-await/async-drop/async-drop-initial.rs
+++ b/tests/ui/async-await/async-drop/async-drop-initial.rs
@@ -62,7 +62,7 @@ fn main() {
         test_async_drop(&j, 16).await;
         test_async_drop(
             AsyncStruct { b: AsyncInt(8), a: AsyncInt(7), i: 6 },
-            if cfg!(panic = "unwind") { 168 } else { 136 },
+            136,
         ).await;
         test_async_drop(ManuallyDrop::new(AsyncInt(9)), 16).await;
 
diff --git a/tests/crashes/140429.rs b/tests/ui/async-await/async-drop/live-dead-storage.rs
index 041eaf86c5c..e6f25d351c3 100644
--- a/tests/crashes/140429.rs
+++ b/tests/ui/async-await/async-drop/live-dead-storage.rs
@@ -1,6 +1,9 @@
-//@ known-bug: #140429
+// ex-ice: #140429
 //@ compile-flags: -Zlint-mir --crate-type lib
 //@ edition:2024
+//@ check-pass
 
 #![feature(async_drop)]
+#![allow(incomplete_features)]
+
 async fn a<T>(x: T) {}
diff --git a/tests/crashes/140531.rs b/tests/ui/async-await/async-drop/live-dead-storage2.rs
index f664481d440..18df870785e 100644
--- a/tests/crashes/140531.rs
+++ b/tests/ui/async-await/async-drop/live-dead-storage2.rs
@@ -1,7 +1,11 @@
-//@ known-bug: #140531
-//@compile-flags: -Zlint-mir --crate-type lib
+// ex-ice: #140531
+//@ compile-flags: -Zlint-mir --crate-type lib
 //@ edition:2024
+//@ check-pass
+
 #![feature(async_drop)]
+#![allow(incomplete_features)]
+
 async fn call_once(f: impl AsyncFnOnce()) {
     let fut = Box::pin(f());
 }
diff --git a/tests/ui/async-await/async-drop/live-dead-storage3.rs b/tests/ui/async-await/async-drop/live-dead-storage3.rs
new file mode 100644
index 00000000000..d9fba5759f7
--- /dev/null
+++ b/tests/ui/async-await/async-drop/live-dead-storage3.rs
@@ -0,0 +1,56 @@
+// ex-ice: #141761
+//@ compile-flags: -Zlint-mir --crate-type lib
+//@ edition:2024
+//@ check-pass
+
+#![feature(async_drop)]
+#![allow(incomplete_features)]
+
+type BoxFuture<T> = std::pin::Pin<Box<dyn Future<Output = T>>>;
+fn main() {}
+async fn f() {
+    run("").await
+}
+struct InMemoryStorage;
+struct User<'dep> {
+    dep: &'dep str,
+}
+impl<'a> StorageRequest<InMemoryStorage> for SaveUser<'a> {
+    fn execute(&self) -> BoxFuture<Result<(), String>> {
+        todo!()
+    }
+}
+trait Storage {
+    type Error;
+}
+impl Storage for InMemoryStorage {
+    type Error = String;
+}
+trait StorageRequestReturnType {
+    type Output;
+}
+trait StorageRequest<S: Storage>: StorageRequestReturnType {
+    fn execute(&self) -> BoxFuture<Result<<Self>::Output, S::Error>>;
+}
+struct SaveUser<'a> {
+    name: &'a str,
+}
+impl<'a> StorageRequestReturnType for SaveUser<'a> {
+    type Output = ();
+}
+impl<'dep> User<'dep> {
+    async fn save<S>(self)
+    where
+        S: Storage,
+        for<'a> SaveUser<'a>: StorageRequest<S>,
+    {
+        SaveUser { name: "" }.execute().await;
+    }
+}
+async fn run<S>(dep: &str)
+where
+    S: Storage,
+    for<'a> SaveUser<'a>: StorageRequest<S>,
+{
+    User { dep }.save().await
+}
diff --git a/tests/ui/async-await/async-drop/live-dead-storage4.rs b/tests/ui/async-await/async-drop/live-dead-storage4.rs
new file mode 100644
index 00000000000..d927cb96674
--- /dev/null
+++ b/tests/ui/async-await/async-drop/live-dead-storage4.rs
@@ -0,0 +1,56 @@
+// ex-ice: #141409
+//@ compile-flags: -Zmir-enable-passes=+Inline -Zvalidate-mir -Zlint-mir --crate-type lib
+//@ edition:2024
+//@ check-pass
+
+#![feature(async_drop)]
+#![allow(incomplete_features)]
+#![allow(non_snake_case)]
+
+use std::mem::ManuallyDrop;
+use std::{
+    future::{async_drop_in_place, Future},
+    pin::{pin, Pin},
+    sync::{mpsc, Arc},
+    task::{Context, Poll, Wake, Waker},
+};
+fn main() {
+    block_on(bar(0))
+}
+async fn baz(ident_base: usize) {}
+async fn bar(ident_base: usize) {
+    baz(1).await
+}
+fn block_on<F>(fut_unpin: F) -> F::Output
+where
+    F: Future,
+{
+    let fut_pin = pin!(ManuallyDrop::new(fut_unpin));
+    let mut fut = unsafe { Pin::map_unchecked_mut(fut_pin, |x| &mut **x) };
+    let (waker, rx) = simple_waker();
+    let mut context = Context::from_waker(&waker);
+    let rv = loop {
+        match fut.as_mut().poll(&mut context) {
+            Poll::Ready(out) => break out,
+            PollPending => (),
+        }
+    };
+    let drop_fut_unpin = unsafe { async_drop_in_place(fut.get_unchecked_mut()) };
+    let drop_fut = pin!(drop_fut_unpin);
+    loop {
+        match drop_fut.poll(&mut context) {
+            Poll => break,
+        }
+    }
+    rv
+}
+fn simple_waker() -> (Waker, mpsc::Receiver<()>) {
+    struct SimpleWaker {
+        tx: mpsc::Sender<()>,
+    }
+    impl Wake for SimpleWaker {
+        fn wake(self: Arc<Self>) {}
+    }
+    let (tx, rx) = mpsc::channel();
+    (Waker::from(Arc::new(SimpleWaker { tx })), rx)
+}