Rollup merge of #127136 - compiler-errors:coroutine-closure-env-shim, r=oli-obk

Fix `FnMut::call_mut`/`Fn::call` shim for async closures that capture references

I adjusted async closures to be able to implement `Fn` and `FnMut` *even if* they capture references, as long as those references did not need to borrow data from the closure captures themselves. See #125259.

However, when I did this, I didn't actually relax an assertion in the `build_construct_coroutine_by_move_shim` shim code, which builds the `Fn`/`FnMut`/`FnOnce` implementations for async closures. Therefore, if we actually tried to *call* `FnMut`/`Fn` on async closures, it would ICE.

This PR adjusts this assertion to ensure that we only capture immutable references in closures if they implement `Fn`/`FnMut`. It also adds a bunch of tests and makes more of the async-closure tests into `build-pass` since we often care about these tests actually generating the right closure shims and stuff. I think it might be excessive to *always* use build-pass here, but 🤷 it's not that big of a deal.

Fixes #127019
Fixes #127012

r? oli-obk
This commit is contained in:
Matthias Krüger 2024-07-02 17:47:46 +02:00 committed by GitHub
commit 3cf567e3c0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
23 changed files with 225 additions and 58 deletions

View File

@ -47,7 +47,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
let expected_ty = self.monomorphize(self.mir.local_decls[local].ty);
if expected_ty != op.layout.ty {
warn!(
"Unexpected initial operand type: expected {expected_ty:?}, found {:?}.\
"Unexpected initial operand type:\nexpected {expected_ty:?},\nfound {:?}.\n\
See <https://github.com/rust-lang/rust/issues/114858>.",
op.layout.ty
);

View File

@ -1,18 +1,17 @@
use rustc_hir as hir;
use rustc_hir::def_id::DefId;
use rustc_hir::lang_items::LangItem;
use rustc_index::{Idx, IndexVec};
use rustc_middle::mir::*;
use rustc_middle::query::Providers;
use rustc_middle::ty::GenericArgs;
use rustc_middle::ty::{self, CoroutineArgs, CoroutineArgsExt, EarlyBinder, Ty, TyCtxt};
use rustc_middle::{bug, span_bug};
use rustc_target::abi::{FieldIdx, VariantIdx, FIRST_VARIANT};
use rustc_index::{Idx, IndexVec};
use rustc_span::{source_map::Spanned, Span, DUMMY_SP};
use rustc_target::abi::{FieldIdx, VariantIdx, FIRST_VARIANT};
use rustc_target::spec::abi::Abi;
use std::assert_matches::assert_matches;
use std::fmt;
use std::iter;
@ -1020,21 +1019,19 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
receiver_by_ref: bool,
) -> Body<'tcx> {
let mut self_ty = tcx.type_of(coroutine_closure_def_id).instantiate_identity();
let mut self_local: Place<'tcx> = Local::from_usize(1).into();
let ty::CoroutineClosure(_, args) = *self_ty.kind() else {
bug!();
};
// 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).
// We use `&Self` here because we only need to emit an ABI-compatible shim body,
// rather than match the signature exactly (which might take `&mut 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
// implementation, not the AsyncFn/AsyncFnMut implementation, which is implemented only
// if the coroutine-closure has no captures.
// We adjust the `self_local` to be a deref since we want to copy fields out of
// a reference to the closure.
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_ref(tcx, tcx.lifetimes.re_erased, self_ty);
self_local = tcx.mk_place_deref(self_local);
self_ty = Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, self_ty);
}
let poly_sig = args.as_coroutine_closure().coroutine_closure_sig().map_bound(|sig| {
@ -1067,11 +1064,27 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
fields.push(Operand::Move(Local::from_usize(idx + 1).into()));
}
for (idx, ty) in args.as_coroutine_closure().upvar_tys().iter().enumerate() {
fields.push(Operand::Move(tcx.mk_place_field(
Local::from_usize(1).into(),
FieldIdx::from_usize(idx),
ty,
)));
if receiver_by_ref {
// The only situation where it's possible is when we capture immuatable references,
// since those don't need to be reborrowed with the closure's env lifetime. Since
// references are always `Copy`, just emit a copy.
assert_matches!(
ty.kind(),
ty::Ref(_, _, hir::Mutability::Not),
"field should be captured by immutable ref if we have an `Fn` instance"
);
fields.push(Operand::Copy(tcx.mk_place_field(
self_local,
FieldIdx::from_usize(idx),
ty,
)));
} else {
fields.push(Operand::Move(tcx.mk_place_field(
self_local,
FieldIdx::from_usize(idx),
ty,
)));
}
}
let source_info = SourceInfo::outermost(span);

View File

@ -85,9 +85,13 @@ pub(super) fn mangle<'tcx>(
}
// FIXME(async_closures): This shouldn't be needed when we fix
// `Instance::ty`/`Instance::def_id`.
ty::InstanceKind::ConstructCoroutineInClosureShim { .. }
| ty::InstanceKind::CoroutineKindShim { .. } => {
printer.write_str("{{fn-once-shim}}").unwrap();
ty::InstanceKind::ConstructCoroutineInClosureShim { receiver_by_ref, .. } => {
printer
.write_str(if receiver_by_ref { "{{by-move-shim}}" } else { "{{by-ref-shim}}" })
.unwrap();
}
ty::InstanceKind::CoroutineKindShim { .. } => {
printer.write_str("{{by-move-body-shim}}").unwrap();
}
_ => {}
}

View File

@ -49,8 +49,15 @@ pub(super) fn mangle<'tcx>(
ty::InstanceKind::ReifyShim(_, Some(ReifyReason::FnPtr)) => Some("reify_fnptr"),
ty::InstanceKind::ReifyShim(_, Some(ReifyReason::Vtable)) => Some("reify_vtable"),
ty::InstanceKind::ConstructCoroutineInClosureShim { .. }
| ty::InstanceKind::CoroutineKindShim { .. } => Some("fn_once"),
// FIXME(async_closures): This shouldn't be needed when we fix
// `Instance::ty`/`Instance::def_id`.
ty::InstanceKind::ConstructCoroutineInClosureShim { receiver_by_ref: true, .. } => {
Some("by_move")
}
ty::InstanceKind::ConstructCoroutineInClosureShim { receiver_by_ref: false, .. } => {
Some("by_ref")
}
ty::InstanceKind::CoroutineKindShim { .. } => Some("by_move_body"),
_ => None,
};

View File

@ -127,9 +127,9 @@ 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 (mut) ref.
// still take their receiver by ref.
if receiver_by_ref {
Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, coroutine_ty)
Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, coroutine_ty)
} else {
coroutine_ty
}

View File

@ -1,7 +1,8 @@
#![feature(async_closure, noop_waker, async_fn_traits)]
#![allow(unused)]
use std::future::Future;
use std::ops::{AsyncFnMut, AsyncFnOnce};
use std::ops::{AsyncFn, AsyncFnMut, AsyncFnOnce};
use std::pin::pin;
use std::task::*;
@ -17,6 +18,10 @@ pub fn block_on<T>(fut: impl Future<Output = T>) -> T {
}
}
async fn call(f: &mut impl AsyncFn(i32)) {
f(0).await;
}
async fn call_mut(f: &mut impl AsyncFnMut(i32)) {
f(0).await;
}
@ -26,10 +31,10 @@ async fn call_once(f: impl AsyncFnOnce(i32)) {
}
async fn call_normal<F: Future<Output = ()>>(f: &impl Fn(i32) -> F) {
f(0).await;
f(1).await;
}
async fn call_normal_once<F: Future<Output = ()>>(f: impl FnOnce(i32) -> F) {
async fn call_normal_mut<F: Future<Output = ()>>(f: &mut impl FnMut(i32) -> F) {
f(1).await;
}
@ -39,14 +44,16 @@ pub fn main() {
let mut async_closure = async move |a: i32| {
println!("{a} {b}");
};
call(&mut async_closure).await;
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}");
let b = 2i32;
let mut async_closure = async |a: i32| {
println!("{a} {b}");
};
call_normal(&async_closure).await;
call_normal_once(async_closure).await;
call_normal_mut(&mut async_closure).await;
call_once(async_closure).await;
});
}

View File

@ -1,4 +1,6 @@
0 2
0 2
1 2
1 2
1 2
1 2
0
1

View File

@ -1,6 +1,6 @@
// MIR for `main::{closure#0}::{closure#0}::{closure#0}` 0 coroutine_by_move
fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:42:53: 45:10}, _2: ResumeTy) -> ()
fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:53:53: 56:10}, _2: ResumeTy) -> ()
yields ()
{
debug _task_context => _2;

View File

@ -1,6 +1,6 @@
// MIR for `main::{closure#0}::{closure#0}::{closure#0}` 0 coroutine_by_move
fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:42:53: 45:10}, _2: ResumeTy) -> ()
fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:53:53: 56:10}, _2: ResumeTy) -> ()
yields ()
{
debug _task_context => _2;

View File

@ -1,10 +1,10 @@
// MIR for `main::{closure#0}::{closure#0}` 0 coroutine_closure_by_move
fn main::{closure#0}::{closure#0}(_1: {async closure@$DIR/async_closure_shims.rs:42:33: 42:52}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:42:53: 45:10} {
let mut _0: {async closure body@$DIR/async_closure_shims.rs:42:53: 45:10};
fn main::{closure#0}::{closure#0}(_1: {async closure@$DIR/async_closure_shims.rs:53:33: 53:52}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:53:53: 56:10} {
let mut _0: {async closure body@$DIR/async_closure_shims.rs:53:53: 56:10};
bb0: {
_0 = {coroutine@$DIR/async_closure_shims.rs:42:53: 45:10 (#0)} { a: move _2, b: move (_1.0: i32) };
_0 = {coroutine@$DIR/async_closure_shims.rs:53:53: 56:10 (#0)} { a: move _2, b: move (_1.0: i32) };
return;
}
}

View File

@ -1,10 +1,10 @@
// MIR for `main::{closure#0}::{closure#0}` 0 coroutine_closure_by_move
fn main::{closure#0}::{closure#0}(_1: {async closure@$DIR/async_closure_shims.rs:42:33: 42:52}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:42:53: 45:10} {
let mut _0: {async closure body@$DIR/async_closure_shims.rs:42:53: 45:10};
fn main::{closure#0}::{closure#0}(_1: {async closure@$DIR/async_closure_shims.rs:53:33: 53:52}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:53:53: 56:10} {
let mut _0: {async closure body@$DIR/async_closure_shims.rs:53:53: 56:10};
bb0: {
_0 = {coroutine@$DIR/async_closure_shims.rs:42:53: 45:10 (#0)} { a: move _2, b: move (_1.0: i32) };
_0 = {coroutine@$DIR/async_closure_shims.rs:53:53: 56:10 (#0)} { a: move _2, b: move (_1.0: i32) };
return;
}
}

View File

@ -0,0 +1,47 @@
// MIR for `main::{closure#0}::{closure#1}::{closure#0}` 0 coroutine_by_move
fn main::{closure#0}::{closure#1}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10}, _2: ResumeTy) -> ()
yields ()
{
debug _task_context => _2;
debug a => (_1.0: i32);
debug b => (*(_1.1: &i32));
let mut _0: ();
let _3: i32;
scope 1 {
debug a => _3;
let _4: &i32;
scope 2 {
debug a => _4;
let _5: &i32;
scope 3 {
debug b => _5;
}
}
}
bb0: {
StorageLive(_3);
_3 = (_1.0: i32);
FakeRead(ForLet(None), _3);
StorageLive(_4);
_4 = &_3;
FakeRead(ForLet(None), _4);
StorageLive(_5);
_5 = &(*(_1.1: &i32));
FakeRead(ForLet(None), _5);
_0 = const ();
StorageDead(_5);
StorageDead(_4);
StorageDead(_3);
drop(_1) -> [return: bb1, unwind: bb2];
}
bb1: {
return;
}
bb2 (cleanup): {
resume;
}
}

View File

@ -0,0 +1,47 @@
// MIR for `main::{closure#0}::{closure#1}::{closure#0}` 0 coroutine_by_move
fn main::{closure#0}::{closure#1}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10}, _2: ResumeTy) -> ()
yields ()
{
debug _task_context => _2;
debug a => (_1.0: i32);
debug b => (*(_1.1: &i32));
let mut _0: ();
let _3: i32;
scope 1 {
debug a => _3;
let _4: &i32;
scope 2 {
debug a => _4;
let _5: &i32;
scope 3 {
debug b => _5;
}
}
}
bb0: {
StorageLive(_3);
_3 = (_1.0: i32);
FakeRead(ForLet(None), _3);
StorageLive(_4);
_4 = &_3;
FakeRead(ForLet(None), _4);
StorageLive(_5);
_5 = &(*(_1.1: &i32));
FakeRead(ForLet(None), _5);
_0 = const ();
StorageDead(_5);
StorageDead(_4);
StorageDead(_3);
drop(_1) -> [return: bb1, unwind: bb2];
}
bb1: {
return;
}
bb2 (cleanup): {
resume;
}
}

View File

@ -0,0 +1,10 @@
// MIR for `main::{closure#0}::{closure#1}` 0 coroutine_closure_by_move
fn main::{closure#0}::{closure#1}(_1: {async closure@$DIR/async_closure_shims.rs:62:33: 62:47}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10} {
let mut _0: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10};
bb0: {
_0 = {coroutine@$DIR/async_closure_shims.rs:62:48: 65:10 (#0)} { a: move _2, b: move (_1.0: &i32) };
return;
}
}

View File

@ -0,0 +1,10 @@
// MIR for `main::{closure#0}::{closure#1}` 0 coroutine_closure_by_move
fn main::{closure#0}::{closure#1}(_1: {async closure@$DIR/async_closure_shims.rs:62:33: 62:47}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10} {
let mut _0: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10};
bb0: {
_0 = {coroutine@$DIR/async_closure_shims.rs:62:48: 65:10 (#0)} { a: move _2, b: move (_1.0: &i32) };
return;
}
}

View File

@ -1,10 +1,10 @@
// 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} {
let mut _0: {async closure body@$DIR/async_closure_shims.rs:49:49: 51:10};
fn main::{closure#0}::{closure#1}(_1: &{async closure@$DIR/async_closure_shims.rs:62:33: 62:47}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10} {
let mut _0: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10};
bb0: {
_0 = {coroutine@$DIR/async_closure_shims.rs:49:49: 51:10 (#0)} { a: move _2 };
_0 = {coroutine@$DIR/async_closure_shims.rs:62:48: 65:10 (#0)} { a: move _2, b: ((*_1).0: &i32) };
return;
}
}

View File

@ -1,10 +1,10 @@
// 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} {
let mut _0: {async closure body@$DIR/async_closure_shims.rs:49:49: 51:10};
fn main::{closure#0}::{closure#1}(_1: &{async closure@$DIR/async_closure_shims.rs:62:33: 62:47}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10} {
let mut _0: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10};
bb0: {
_0 = {coroutine@$DIR/async_closure_shims.rs:49:49: 51:10 (#0)} { a: move _2 };
_0 = {coroutine@$DIR/async_closure_shims.rs:62:48: 65:10 (#0)} { a: move _2, b: ((*_1).0: &i32) };
return;
}
}

View File

@ -3,9 +3,10 @@
// EMIT_MIR_FOR_EACH_PANIC_STRATEGY
#![feature(async_closure, noop_waker, async_fn_traits)]
#![allow(unused)]
use std::future::Future;
use std::ops::{AsyncFnMut, AsyncFnOnce};
use std::ops::{AsyncFn, AsyncFnMut, AsyncFnOnce};
use std::pin::pin;
use std::task::*;
@ -21,6 +22,10 @@ pub fn block_on<T>(fut: impl Future<Output = T>) -> T {
}
}
async fn call(f: &mut impl AsyncFn(i32)) {
f(0).await;
}
async fn call_mut(f: &mut impl AsyncFnMut(i32)) {
f(0).await;
}
@ -33,9 +38,15 @@ async fn call_normal<F: Future<Output = ()>>(f: &impl Fn(i32) -> F) {
f(1).await;
}
async fn call_normal_mut<F: Future<Output = ()>>(f: &mut impl FnMut(i32) -> F) {
f(1).await;
}
// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#0}.coroutine_closure_by_move.0.mir
// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#0}-{closure#0}.coroutine_by_move.0.mir
// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.mir
// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_move.0.mir
// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#1}-{closure#0}.coroutine_by_move.0.mir
pub fn main() {
block_on(async {
let b = 2i32;
@ -43,12 +54,17 @@ pub fn main() {
let a = &a;
let b = &b;
};
call(&mut async_closure).await;
call_mut(&mut async_closure).await;
call_once(async_closure).await;
let async_closure = async move |a: i32| {
let b = 2i32;
let mut async_closure = async |a: i32| {
let a = &a;
let b = &b;
};
call_normal(&async_closure).await;
call_normal_mut(&mut async_closure).await;
call_once(async_closure).await;
});
}

View File

@ -1,5 +1,5 @@
//@ edition: 2021
//@ check-pass
//@ build-pass
//@ revisions: current next
//@ ignore-compare-mode-next-solver (explicit revisions)
//@[next] compile-flags: -Znext-solver

View File

@ -1,6 +1,6 @@
//@ aux-build:block-on.rs
//@ edition:2021
//@ check-pass
//@ build-pass
#![feature(async_closure)]

View File

@ -1,6 +1,6 @@
//@ aux-build:block-on.rs
//@ edition:2021
//@ check-pass
//@ build-pass
#![feature(async_closure)]

View File

@ -1,4 +1,4 @@
//@ check-pass
//@ build-pass
//@ edition: 2021
// Demonstrates that an async closure may implement `FnMut` (not just `async FnMut`!)
@ -9,9 +9,13 @@
#![feature(async_closure)]
fn main() {}
fn main() {
hello(&Ty);
}
fn needs_fn_mut<T>(x: impl FnMut() -> T) {}
fn needs_fn_mut<T>(mut x: impl FnMut() -> T) {
x();
}
fn hello(x: &Ty) {
needs_fn_mut(async || { x.hello(); });

View File

@ -1,4 +1,4 @@
//@ check-pass
//@ build-pass
//@ edition: 2021
#![feature(async_closure)]