Rollup merge of #123049 - compiler-errors:coroutine-closure-rcvr, r=oli-obk

In `ConstructCoroutineInClosureShim`, pass receiver by mut ref, not mut pointer

The receivers were compatible at codegen time, but did not necessarily have the same layouts due to niches, which was caught by miri.

Fixes rust-lang/miri#3400

r? oli-obk

8 files changed, 81 insertions, 21 deletions

diff --git a/compiler/rustc_mir_transform/src/shim.rs b/compiler/rustc_mir_transform/src/shim.rs
index 94a95428ab0..b60ee7649b2 100644
--- a/compiler/rustc_mir_transform/src/shim.rs
+++ b/compiler/rustc_mir_transform/src/shim.rs
@@ -1015,8 +1015,8 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
         bug!();
     };
 
-    // We use `*mut Self` here because we only need to emit an ABI-compatible shim body,
-    // rather than match the signature exactly.
+    // We use `&mut Self` here because we only need to emit an ABI-compatible shim body,
+    // rather than match the signature exactly (which might take `&self` instead).
     //
     // The self type here is a coroutine-closure, not a coroutine, and we never read from
     // it because it never has any captures, because this is only true in the Fn/FnMut
@@ -1025,7 +1025,7 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
     if receiver_by_ref {
         // Triple-check that there's no captures here.
         assert_eq!(args.as_coroutine_closure().tupled_upvars_ty(), tcx.types.unit);
-        self_ty = Ty::new_mut_ptr(tcx, self_ty);
+        self_ty = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, self_ty);
     }
 
     let poly_sig = args.as_coroutine_closure().coroutine_closure_sig().map_bound(|sig| {
diff --git a/compiler/rustc_ty_utils/src/abi.rs b/compiler/rustc_ty_utils/src/abi.rs
index af1dfb6f7e9..65c3cf1a607 100644
--- a/compiler/rustc_ty_utils/src/abi.rs
+++ b/compiler/rustc_ty_utils/src/abi.rs
@@ -125,8 +125,12 @@ fn fn_sig_for_fn_abi<'tcx>(
                     coroutine_kind = ty::ClosureKind::FnOnce;
 
                     // Implementations of `FnMut` and `Fn` for coroutine-closures
-                    // still take their receiver by ref.
-                    if receiver_by_ref { Ty::new_mut_ptr(tcx, coroutine_ty) } else { coroutine_ty }
+                    // still take their receiver by (mut) ref.
+                    if receiver_by_ref {
+                        Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, coroutine_ty)
+                    } else {
+                        coroutine_ty
+                    }
                 } else {
                     tcx.closure_env_ty(coroutine_ty, coroutine_kind, env_region)
                 };
diff --git a/src/tools/miri/tests/pass/async-closure-drop.rs b/src/tools/miri/tests/pass/async-closure-drop.rs
new file mode 100644
index 00000000000..9b2fc2948bf
--- /dev/null
+++ b/src/tools/miri/tests/pass/async-closure-drop.rs
@@ -0,0 +1,40 @@
+#![feature(async_closure, noop_waker, async_fn_traits)]
+
+use std::future::Future;
+use std::pin::pin;
+use std::task::*;
+
+pub fn block_on<T>(fut: impl Future<Output = T>) -> T {
+    let mut fut = pin!(fut);
+    let ctx = &mut Context::from_waker(Waker::noop());
+
+    loop {
+        match fut.as_mut().poll(ctx) {
+            Poll::Pending => {}
+            Poll::Ready(t) => break t,
+        }
+    }
+}
+
+async fn call_once(f: impl async FnOnce(DropMe)) {
+    f(DropMe("world")).await;
+}
+
+#[derive(Debug)]
+struct DropMe(&'static str);
+
+impl Drop for DropMe {
+    fn drop(&mut self) {
+        println!("{}", self.0);
+    }
+}
+
+pub fn main() {
+    block_on(async {
+        let b = DropMe("hello");
+        let async_closure = async move |a: DropMe| {
+            println!("{a:?} {b:?}");
+        };
+        call_once(async_closure).await;
+    });
+}
diff --git a/src/tools/miri/tests/pass/async-closure-drop.stdout b/src/tools/miri/tests/pass/async-closure-drop.stdout
new file mode 100644
index 00000000000..34cfdedc44a
--- /dev/null
+++ b/src/tools/miri/tests/pass/async-closure-drop.stdout
@@ -0,0 +1,3 @@
+DropMe("world") DropMe("hello")
+world
+hello
diff --git a/src/tools/miri/tests/pass/async-closure.rs b/src/tools/miri/tests/pass/async-closure.rs
index 9b2fc2948bf..2f7ec2b9e6f 100644
--- a/src/tools/miri/tests/pass/async-closure.rs
+++ b/src/tools/miri/tests/pass/async-closure.rs
@@ -1,6 +1,7 @@
 #![feature(async_closure, noop_waker, async_fn_traits)]
 
 use std::future::Future;
+use std::ops::{AsyncFnMut, AsyncFnOnce};
 use std::pin::pin;
 use std::task::*;
 
@@ -16,25 +17,36 @@ pub fn block_on<T>(fut: impl Future<Output = T>) -> T {
     }
 }
 
-async fn call_once(f: impl async FnOnce(DropMe)) {
-    f(DropMe("world")).await;
+async fn call_mut(f: &mut impl AsyncFnMut(i32)) {
+    f(0).await;
 }
 
-#[derive(Debug)]
-struct DropMe(&'static str);
+async fn call_once(f: impl AsyncFnOnce(i32)) {
+    f(1).await;
+}
 
-impl Drop for DropMe {
-    fn drop(&mut self) {
-        println!("{}", self.0);
-    }
+async fn call_normal<F: Future<Output = ()>>(f: &impl Fn(i32) -> F) {
+    f(0).await;
+}
+
+async fn call_normal_once<F: Future<Output = ()>>(f: impl FnOnce(i32) -> F) {
+    f(1).await;
 }
 
 pub fn main() {
     block_on(async {
-        let b = DropMe("hello");
-        let async_closure = async move |a: DropMe| {
-            println!("{a:?} {b:?}");
+        let b = 2i32;
+        let mut async_closure = async move |a: i32| {
+            println!("{a} {b}");
         };
+        call_mut(&mut async_closure).await;
         call_once(async_closure).await;
+
+        // No-capture closures implement `Fn`.
+        let async_closure = async move |a: i32| {
+            println!("{a}");
+        };
+        call_normal(&async_closure).await;
+        call_normal_once(async_closure).await;
     });
 }
diff --git a/src/tools/miri/tests/pass/async-closure.stdout b/src/tools/miri/tests/pass/async-closure.stdout
index 34cfdedc44a..7baae1aa94f 100644
--- a/src/tools/miri/tests/pass/async-closure.stdout
+++ b/src/tools/miri/tests/pass/async-closure.stdout
@@ -1,3 +1,4 @@
-DropMe("world") DropMe("hello")
-world
-hello
+0 2
+1 2
+0
+1
diff --git a/tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.panic-abort.mir b/tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.panic-abort.mir
index f51540bcfff..cab7bdb7e3c 100644
--- a/tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.panic-abort.mir
+++ b/tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.panic-abort.mir
@@ -1,6 +1,6 @@
 // MIR for `main::{closure#0}::{closure#1}` 0 coroutine_closure_by_ref
 
-fn main::{closure#0}::{closure#1}(_1: *mut {async closure@$DIR/async_closure_shims.rs:49:29: 49:48}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:49:49: 51:10} {
+fn main::{closure#0}::{closure#1}(_1: &mut {async closure@$DIR/async_closure_shims.rs:49:29: 49:48}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:49:49: 51:10} {
     let mut _0: {async closure body@$DIR/async_closure_shims.rs:49:49: 51:10};
 
     bb0: {
diff --git a/tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.panic-unwind.mir b/tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.panic-unwind.mir
index f51540bcfff..cab7bdb7e3c 100644
--- a/tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.panic-unwind.mir
+++ b/tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.panic-unwind.mir
@@ -1,6 +1,6 @@
 // MIR for `main::{closure#0}::{closure#1}` 0 coroutine_closure_by_ref
 
-fn main::{closure#0}::{closure#1}(_1: *mut {async closure@$DIR/async_closure_shims.rs:49:29: 49:48}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:49:49: 51:10} {
+fn main::{closure#0}::{closure#1}(_1: &mut {async closure@$DIR/async_closure_shims.rs:49:29: 49:48}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:49:49: 51:10} {
     let mut _0: {async closure body@$DIR/async_closure_shims.rs:49:49: 51:10};
 
     bb0: {