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

5 files changed, 133 insertions, 1 deletions

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/ui/async-await/async-drop/live-dead-storage.rs b/tests/ui/async-await/async-drop/live-dead-storage.rs
new file mode 100644
index 00000000000..e6f25d351c3
--- /dev/null
+++ b/tests/ui/async-await/async-drop/live-dead-storage.rs
@@ -0,0 +1,9 @@
+// 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/ui/async-await/async-drop/live-dead-storage2.rs b/tests/ui/async-await/async-drop/live-dead-storage2.rs
new file mode 100644
index 00000000000..18df870785e
--- /dev/null
+++ b/tests/ui/async-await/async-drop/live-dead-storage2.rs
@@ -0,0 +1,11 @@
+// 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)
+}