Skip to content

Commit 3cf567e

Browse files
matthiaskrgrmatthiaskrgr
authored
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
2 parents f8f67b2 + 90143b0 commit 3cf567e

23 files changed

+225
-58
lines changed

‎compiler/rustc_codegen_ssa/src/mir/locals.rs

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

‎compiler/rustc_mir_transform/src/shim.rs

+31-18
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,17 @@
11
use rustc_hir as hir;
22
use rustc_hir::def_id::DefId;
33
use rustc_hir::lang_items::LangItem;
4+
use rustc_index::{Idx, IndexVec};
45
use rustc_middle::mir::*;
56
use rustc_middle::query::Providers;
67
use rustc_middle::ty::GenericArgs;
78
use rustc_middle::ty::{self, CoroutineArgs, CoroutineArgsExt, EarlyBinder, Ty, TyCtxt};
89
use rustc_middle::{bug, span_bug};
9-
use rustc_target::abi::{FieldIdx, VariantIdx, FIRST_VARIANT};
10-
11-
use rustc_index::{Idx, IndexVec};
12-
1310
use rustc_span::{source_map::Spanned, Span, DUMMY_SP};
11+
use rustc_target::abi::{FieldIdx, VariantIdx, FIRST_VARIANT};
1412
use rustc_target::spec::abi::Abi;
1513

14+
use std::assert_matches::assert_matches;
1615
use std::fmt;
1716
use std::iter;
1817

@@ -1020,21 +1019,19 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
10201019
receiver_by_ref: bool,
10211020
) -> Body<'tcx> {
10221021
let mut self_ty = tcx.type_of(coroutine_closure_def_id).instantiate_identity();
1022+
let mut self_local: Place<'tcx> = Local::from_usize(1).into();
10231023
let ty::CoroutineClosure(_, args) = *self_ty.kind() else {
10241024
bug!();
10251025
};
10261026

1027-
// We use `&mut Self` here because we only need to emit an ABI-compatible shim body,
1028-
// rather than match the signature exactly (which might take `&self` instead).
1027+
// We use `&Self` here because we only need to emit an ABI-compatible shim body,
1028+
// rather than match the signature exactly (which might take `&mut self` instead).
10291029
//
1030-
// The self type here is a coroutine-closure, not a coroutine, and we never read from
1031-
// it because it never has any captures, because this is only true in the Fn/FnMut
1032-
// implementation, not the AsyncFn/AsyncFnMut implementation, which is implemented only
1033-
// if the coroutine-closure has no captures.
1030+
// We adjust the `self_local` to be a deref since we want to copy fields out of
1031+
// a reference to the closure.
10341032
if receiver_by_ref {
1035-
// Triple-check that there's no captures here.
1036-
assert_eq!(args.as_coroutine_closure().tupled_upvars_ty(), tcx.types.unit);
1037-
self_ty = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, self_ty);
1033+
self_local = tcx.mk_place_deref(self_local);
1034+
self_ty = Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, self_ty);
10381035
}
10391036

10401037
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>(
10671064
fields.push(Operand::Move(Local::from_usize(idx + 1).into()));
10681065
}
10691066
for (idx, ty) in args.as_coroutine_closure().upvar_tys().iter().enumerate() {
1070-
fields.push(Operand::Move(tcx.mk_place_field(
1071-
Local::from_usize(1).into(),
1072-
FieldIdx::from_usize(idx),
1073-
ty,
1074-
)));
1067+
if receiver_by_ref {
1068+
// The only situation where it's possible is when we capture immuatable references,
1069+
// since those don't need to be reborrowed with the closure's env lifetime. Since
1070+
// references are always `Copy`, just emit a copy.
1071+
assert_matches!(
1072+
ty.kind(),
1073+
ty::Ref(_, _, hir::Mutability::Not),
1074+
"field should be captured by immutable ref if we have an `Fn` instance"
1075+
);
1076+
fields.push(Operand::Copy(tcx.mk_place_field(
1077+
self_local,
1078+
FieldIdx::from_usize(idx),
1079+
ty,
1080+
)));
1081+
} else {
1082+
fields.push(Operand::Move(tcx.mk_place_field(
1083+
self_local,
1084+
FieldIdx::from_usize(idx),
1085+
ty,
1086+
)));
1087+
}
10751088
}
10761089

10771090
let source_info = SourceInfo::outermost(span);

‎compiler/rustc_symbol_mangling/src/legacy.rs

+7-3
Original file line numberDiff line numberDiff line change
@@ -85,9 +85,13 @@ pub(super) fn mangle<'tcx>(
8585
}
8686
// FIXME(async_closures): This shouldn't be needed when we fix
8787
// `Instance::ty`/`Instance::def_id`.
88-
ty::InstanceKind::ConstructCoroutineInClosureShim { .. }
89-
| ty::InstanceKind::CoroutineKindShim { .. } => {
90-
printer.write_str("{{fn-once-shim}}").unwrap();
88+
ty::InstanceKind::ConstructCoroutineInClosureShim { receiver_by_ref, .. } => {
89+
printer
90+
.write_str(if receiver_by_ref { "{{by-move-shim}}" } else { "{{by-ref-shim}}" })
91+
.unwrap();
92+
}
93+
ty::InstanceKind::CoroutineKindShim { .. } => {
94+
printer.write_str("{{by-move-body-shim}}").unwrap();
9195
}
9296
_ => {}
9397
}

‎compiler/rustc_symbol_mangling/src/v0.rs

+9-2
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,15 @@ pub(super) fn mangle<'tcx>(
4949
ty::InstanceKind::ReifyShim(_, Some(ReifyReason::FnPtr)) => Some("reify_fnptr"),
5050
ty::InstanceKind::ReifyShim(_, Some(ReifyReason::Vtable)) => Some("reify_vtable"),
5151

52-
ty::InstanceKind::ConstructCoroutineInClosureShim { .. }
53-
| ty::InstanceKind::CoroutineKindShim { .. } => Some("fn_once"),
52+
// FIXME(async_closures): This shouldn't be needed when we fix
53+
// `Instance::ty`/`Instance::def_id`.
54+
ty::InstanceKind::ConstructCoroutineInClosureShim { receiver_by_ref: true, .. } => {
55+
Some("by_move")
56+
}
57+
ty::InstanceKind::ConstructCoroutineInClosureShim { receiver_by_ref: false, .. } => {
58+
Some("by_ref")
59+
}
60+
ty::InstanceKind::CoroutineKindShim { .. } => Some("by_move_body"),
5461

5562
_ => None,
5663
};

‎compiler/rustc_ty_utils/src/abi.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -127,9 +127,9 @@ fn fn_sig_for_fn_abi<'tcx>(
127127
coroutine_kind = ty::ClosureKind::FnOnce;
128128

129129
// Implementations of `FnMut` and `Fn` for coroutine-closures
130-
// still take their receiver by (mut) ref.
130+
// still take their receiver by ref.
131131
if receiver_by_ref {
132-
Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, coroutine_ty)
132+
Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, coroutine_ty)
133133
} else {
134134
coroutine_ty
135135
}

‎src/tools/miri/tests/pass/async-closure.rs

+14-7
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
#![feature(async_closure, noop_waker, async_fn_traits)]
2+
#![allow(unused)]
23

34
use std::future::Future;
4-
use std::ops::{AsyncFnMut, AsyncFnOnce};
5+
use std::ops::{AsyncFn, AsyncFnMut, AsyncFnOnce};
56
use std::pin::pin;
67
use std::task::*;
78

@@ -17,6 +18,10 @@ pub fn block_on<T>(fut: impl Future<Output = T>) -> T {
1718
}
1819
}
1920

21+
async fn call(f: &mut impl AsyncFn(i32)) {
22+
f(0).await;
23+
}
24+
2025
async fn call_mut(f: &mut impl AsyncFnMut(i32)) {
2126
f(0).await;
2227
}
@@ -26,10 +31,10 @@ async fn call_once(f: impl AsyncFnOnce(i32)) {
2631
}
2732

2833
async fn call_normal<F: Future<Output = ()>>(f: &impl Fn(i32) -> F) {
29-
f(0).await;
34+
f(1).await;
3035
}
3136

32-
async fn call_normal_once<F: Future<Output = ()>>(f: impl FnOnce(i32) -> F) {
37+
async fn call_normal_mut<F: Future<Output = ()>>(f: &mut impl FnMut(i32) -> F) {
3338
f(1).await;
3439
}
3540

@@ -39,14 +44,16 @@ pub fn main() {
3944
let mut async_closure = async move |a: i32| {
4045
println!("{a} {b}");
4146
};
47+
call(&mut async_closure).await;
4248
call_mut(&mut async_closure).await;
4349
call_once(async_closure).await;
4450

45-
// No-capture closures implement `Fn`.
46-
let async_closure = async move |a: i32| {
47-
println!("{a}");
51+
let b = 2i32;
52+
let mut async_closure = async |a: i32| {
53+
println!("{a} {b}");
4854
};
4955
call_normal(&async_closure).await;
50-
call_normal_once(async_closure).await;
56+
call_normal_mut(&mut async_closure).await;
57+
call_once(async_closure).await;
5158
});
5259
}

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

+4-2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
0 2
2+
0 2
3+
1 2
4+
1 2
5+
1 2
26
1 2
3-
0
4-
1

‎tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}-{closure#0}.coroutine_by_move.0.panic-abort.mir

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
// MIR for `main::{closure#0}::{closure#0}::{closure#0}` 0 coroutine_by_move
22

3-
fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:42:53: 45:10}, _2: ResumeTy) -> ()
3+
fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:53:53: 56:10}, _2: ResumeTy) -> ()
44
yields ()
55
{
66
debug _task_context => _2;

‎tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}-{closure#0}.coroutine_by_move.0.panic-unwind.mir

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
// MIR for `main::{closure#0}::{closure#0}::{closure#0}` 0 coroutine_by_move
22

3-
fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:42:53: 45:10}, _2: ResumeTy) -> ()
3+
fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:53:53: 56:10}, _2: ResumeTy) -> ()
44
yields ()
55
{
66
debug _task_context => _2;

‎tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}.coroutine_closure_by_move.0.panic-abort.mir

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
// MIR for `main::{closure#0}::{closure#0}` 0 coroutine_closure_by_move
22

3-
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} {
4-
let mut _0: {async closure body@$DIR/async_closure_shims.rs:42:53: 45:10};
3+
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} {
4+
let mut _0: {async closure body@$DIR/async_closure_shims.rs:53:53: 56:10};
55

66
bb0: {
7-
_0 = {coroutine@$DIR/async_closure_shims.rs:42:53: 45:10 (#0)} { a: move _2, b: move (_1.0: i32) };
7+
_0 = {coroutine@$DIR/async_closure_shims.rs:53:53: 56:10 (#0)} { a: move _2, b: move (_1.0: i32) };
88
return;
99
}
1010
}

‎tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#0}.coroutine_closure_by_move.0.panic-unwind.mir

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
// MIR for `main::{closure#0}::{closure#0}` 0 coroutine_closure_by_move
22

3-
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} {
4-
let mut _0: {async closure body@$DIR/async_closure_shims.rs:42:53: 45:10};
3+
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} {
4+
let mut _0: {async closure body@$DIR/async_closure_shims.rs:53:53: 56:10};
55

66
bb0: {
7-
_0 = {coroutine@$DIR/async_closure_shims.rs:42:53: 45:10 (#0)} { a: move _2, b: move (_1.0: i32) };
7+
_0 = {coroutine@$DIR/async_closure_shims.rs:53:53: 56:10 (#0)} { a: move _2, b: move (_1.0: i32) };
88
return;
99
}
1010
}

‎tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}-{closure#0}.coroutine_by_move.0.panic-abort.mir

+47
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
// MIR for `main::{closure#0}::{closure#1}::{closure#0}` 0 coroutine_by_move
2+
3+
fn main::{closure#0}::{closure#1}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10}, _2: ResumeTy) -> ()
4+
yields ()
5+
{
6+
debug _task_context => _2;
7+
debug a => (_1.0: i32);
8+
debug b => (*(_1.1: &i32));
9+
let mut _0: ();
10+
let _3: i32;
11+
scope 1 {
12+
debug a => _3;
13+
let _4: &i32;
14+
scope 2 {
15+
debug a => _4;
16+
let _5: &i32;
17+
scope 3 {
18+
debug b => _5;
19+
}
20+
}
21+
}
22+
23+
bb0: {
24+
StorageLive(_3);
25+
_3 = (_1.0: i32);
26+
FakeRead(ForLet(None), _3);
27+
StorageLive(_4);
28+
_4 = &_3;
29+
FakeRead(ForLet(None), _4);
30+
StorageLive(_5);
31+
_5 = &(*(_1.1: &i32));
32+
FakeRead(ForLet(None), _5);
33+
_0 = const ();
34+
StorageDead(_5);
35+
StorageDead(_4);
36+
StorageDead(_3);
37+
drop(_1) -> [return: bb1, unwind: bb2];
38+
}
39+
40+
bb1: {
41+
return;
42+
}
43+
44+
bb2 (cleanup): {
45+
resume;
46+
}
47+
}

‎tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}-{closure#0}.coroutine_by_move.0.panic-unwind.mir

+47
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
// MIR for `main::{closure#0}::{closure#1}::{closure#0}` 0 coroutine_by_move
2+
3+
fn main::{closure#0}::{closure#1}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10}, _2: ResumeTy) -> ()
4+
yields ()
5+
{
6+
debug _task_context => _2;
7+
debug a => (_1.0: i32);
8+
debug b => (*(_1.1: &i32));
9+
let mut _0: ();
10+
let _3: i32;
11+
scope 1 {
12+
debug a => _3;
13+
let _4: &i32;
14+
scope 2 {
15+
debug a => _4;
16+
let _5: &i32;
17+
scope 3 {
18+
debug b => _5;
19+
}
20+
}
21+
}
22+
23+
bb0: {
24+
StorageLive(_3);
25+
_3 = (_1.0: i32);
26+
FakeRead(ForLet(None), _3);
27+
StorageLive(_4);
28+
_4 = &_3;
29+
FakeRead(ForLet(None), _4);
30+
StorageLive(_5);
31+
_5 = &(*(_1.1: &i32));
32+
FakeRead(ForLet(None), _5);
33+
_0 = const ();
34+
StorageDead(_5);
35+
StorageDead(_4);
36+
StorageDead(_3);
37+
drop(_1) -> [return: bb1, unwind: bb2];
38+
}
39+
40+
bb1: {
41+
return;
42+
}
43+
44+
bb2 (cleanup): {
45+
resume;
46+
}
47+
}

‎tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_move.0.panic-abort.mir

+10
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
// MIR for `main::{closure#0}::{closure#1}` 0 coroutine_closure_by_move
2+
3+
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} {
4+
let mut _0: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10};
5+
6+
bb0: {
7+
_0 = {coroutine@$DIR/async_closure_shims.rs:62:48: 65:10 (#0)} { a: move _2, b: move (_1.0: &i32) };
8+
return;
9+
}
10+
}

‎tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_move.0.panic-unwind.mir

+10
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
// MIR for `main::{closure#0}::{closure#1}` 0 coroutine_closure_by_move
2+
3+
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} {
4+
let mut _0: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10};
5+
6+
bb0: {
7+
_0 = {coroutine@$DIR/async_closure_shims.rs:62:48: 65:10 (#0)} { a: move _2, b: move (_1.0: &i32) };
8+
return;
9+
}
10+
}

‎tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.panic-abort.mir

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
// MIR for `main::{closure#0}::{closure#1}` 0 coroutine_closure_by_ref
22

3-
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} {
4-
let mut _0: {async closure body@$DIR/async_closure_shims.rs:49:49: 51:10};
3+
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} {
4+
let mut _0: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10};
55

66
bb0: {
7-
_0 = {coroutine@$DIR/async_closure_shims.rs:49:49: 51:10 (#0)} { a: move _2 };
7+
_0 = {coroutine@$DIR/async_closure_shims.rs:62:48: 65:10 (#0)} { a: move _2, b: ((*_1).0: &i32) };
88
return;
99
}
1010
}

‎tests/mir-opt/async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.panic-unwind.mir

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
// MIR for `main::{closure#0}::{closure#1}` 0 coroutine_closure_by_ref
22

3-
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} {
4-
let mut _0: {async closure body@$DIR/async_closure_shims.rs:49:49: 51:10};
3+
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} {
4+
let mut _0: {async closure body@$DIR/async_closure_shims.rs:62:48: 65:10};
55

66
bb0: {
7-
_0 = {coroutine@$DIR/async_closure_shims.rs:49:49: 51:10 (#0)} { a: move _2 };
7+
_0 = {coroutine@$DIR/async_closure_shims.rs:62:48: 65:10 (#0)} { a: move _2, b: ((*_1).0: &i32) };
88
return;
99
}
1010
}

0 commit comments

Comments
 (0)