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

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,
+    // We use `&mut Self` here because we only need to emit an ABI-compatible shim body,
-    // rather than match the signature exactly.
+    // 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| {

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.
+                    // still take their receiver by (mut) ref.
-                    if receiver_by_ref { Ty::new_mut_ptr(tcx, coroutine_ty) } else { coroutine_ty }
+                    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)
                 };

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

src/tools/miri/tests/pass/async-closure-drop.stdout

@@ -0,0 +1,3 @@
+DropMe("world") DropMe("hello")
+world
+hello

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)) {
+async fn call_mut(f: &mut impl AsyncFnMut(i32)) {
-    f(DropMe("world")).await;
+    f(0).await;
 }
 
-#[derive(Debug)]
+async fn call_once(f: impl AsyncFnOnce(i32)) {
-struct DropMe(&'static str);
+    f(1).await;
+}
 
-impl Drop for DropMe {
+async fn call_normal<F: Future<Output = ()>>(f: &impl Fn(i32) -> F) {
-    fn drop(&mut self) {
+    f(0).await;
-        println!("{}", self.0);
+}
-    }
+
+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 b = 2i32;
-        let async_closure = async move |a: DropMe| {
+        let mut async_closure = async move |a: i32| {
-            println!("{a:?} {b:?}");
+            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;
     });
 }

src/tools/miri/tests/pass/async-closure.stdout

@@ -1,3 +1,4 @@
-DropMe("world") DropMe("hello")
+0 2
-world
+1 2
-hello
+0
+1

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: {

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: {