Skip to content

Commit a6570ec

Browse files
compiler-errorscompiler-errors
committed
Fix ABI for FnMut/Fn impls for async closures
1 parent ba9a24f commit a6570ec

12 files changed

+81
-18
lines changed

‎compiler/rustc_middle/src/mir/visit.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,7 @@ macro_rules! make_mir_visitor {
347347
ty::InstanceDef::ClosureOnceShim { call_once: _def_id, track_caller: _ } |
348348
ty::InstanceDef::ConstructCoroutineInClosureShim {
349349
coroutine_closure_def_id: _def_id,
350+
receiver_by_ref: _,
350351
} |
351352
ty::InstanceDef::CoroutineKindShim { coroutine_def_id: _def_id } |
352353
ty::InstanceDef::DropGlue(_def_id, None) => {}

‎compiler/rustc_middle/src/ty/instance.rs

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,15 @@ pub enum InstanceDef<'tcx> {
9595
/// The body generated here differs significantly from the `ClosureOnceShim`,
9696
/// since we need to generate a distinct coroutine type that will move the
9797
/// closure's upvars *out* of the closure.
98-
ConstructCoroutineInClosureShim { coroutine_closure_def_id: DefId },
98+
ConstructCoroutineInClosureShim {
99+
coroutine_closure_def_id: DefId,
100+
// Whether the generated MIR body takes the coroutine by-ref. This is
101+
// because the signature of `<{async fn} as FnMut>::call_mut` is:
102+
// `fn(&mut self, args: A) -> <Self as FnOnce>::Output`, that is to say
103+
// that it returns the `FnOnce`-flavored coroutine but takes the closure
104+
// by ref (and similarly for `Fn::call`).
105+
receiver_by_ref: bool,
106+
},
99107

100108
/// `<[coroutine] as Future>::poll`, but for coroutines produced when `AsyncFnOnce`
101109
/// is called on a coroutine-closure whose closure kind greater than `FnOnce`, or
@@ -188,6 +196,7 @@ impl<'tcx> InstanceDef<'tcx> {
188196
| InstanceDef::ClosureOnceShim { call_once: def_id, track_caller: _ }
189197
| ty::InstanceDef::ConstructCoroutineInClosureShim {
190198
coroutine_closure_def_id: def_id,
199+
receiver_by_ref: _,
191200
}
192201
| ty::InstanceDef::CoroutineKindShim { coroutine_def_id: def_id }
193202
| InstanceDef::DropGlue(def_id, _)

‎compiler/rustc_mir_transform/src/shim.rs

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -66,9 +66,10 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
6666
build_call_shim(tcx, instance, Some(Adjustment::RefMut), CallKind::Direct(call_mut))
6767
}
6868

69-
ty::InstanceDef::ConstructCoroutineInClosureShim { coroutine_closure_def_id } => {
70-
build_construct_coroutine_by_move_shim(tcx, coroutine_closure_def_id)
71-
}
69+
ty::InstanceDef::ConstructCoroutineInClosureShim {
70+
coroutine_closure_def_id,
71+
receiver_by_ref,
72+
} => build_construct_coroutine_by_move_shim(tcx, coroutine_closure_def_id, receiver_by_ref),
7273

7374
ty::InstanceDef::CoroutineKindShim { coroutine_def_id } => {
7475
return tcx.optimized_mir(coroutine_def_id).coroutine_by_move_body().unwrap().clone();
@@ -1011,12 +1012,17 @@ fn build_fn_ptr_addr_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, self_ty: Ty<'t
10111012
fn build_construct_coroutine_by_move_shim<'tcx>(
10121013
tcx: TyCtxt<'tcx>,
10131014
coroutine_closure_def_id: DefId,
1015+
receiver_by_ref: bool,
10141016
) -> Body<'tcx> {
1015-
let self_ty = tcx.type_of(coroutine_closure_def_id).instantiate_identity();
1017+
let mut self_ty = tcx.type_of(coroutine_closure_def_id).instantiate_identity();
10161018
let ty::CoroutineClosure(_, args) = *self_ty.kind() else {
10171019
bug!();
10181020
};
10191021

1022+
if receiver_by_ref {
1023+
self_ty = Ty::new_mut_ptr(tcx, self_ty);
1024+
}
1025+
10201026
let poly_sig = args.as_coroutine_closure().coroutine_closure_sig().map_bound(|sig| {
10211027
tcx.mk_fn_sig(
10221028
[self_ty].into_iter().chain(sig.tupled_inputs_ty.tuple_fields()),
@@ -1072,11 +1078,19 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
10721078

10731079
let source = MirSource::from_instance(ty::InstanceDef::ConstructCoroutineInClosureShim {
10741080
coroutine_closure_def_id,
1081+
receiver_by_ref,
10751082
});
10761083

10771084
let body =
10781085
new_body(source, IndexVec::from_elem_n(start_block, 1), locals, sig.inputs().len(), span);
1079-
dump_mir(tcx, false, "coroutine_closure_by_move", &0, &body, |_, _| Ok(()));
1086+
dump_mir(
1087+
tcx,
1088+
false,
1089+
if receiver_by_ref { "coroutine_closure_by_ref" } else { "coroutine_closure_by_move" },
1090+
&0,
1091+
&body,
1092+
|_, _| Ok(()),
1093+
);
10801094

10811095
body
10821096
}

‎compiler/rustc_ty_utils/src/abi.rs

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -118,11 +118,18 @@ fn fn_sig_for_fn_abi<'tcx>(
118118
// a separate def-id for these bodies.
119119
let mut coroutine_kind = args.as_coroutine_closure().kind();
120120

121-
if let InstanceDef::ConstructCoroutineInClosureShim { .. } = instance.def {
122-
coroutine_kind = ty::ClosureKind::FnOnce;
123-
}
121+
let env_ty =
122+
if let InstanceDef::ConstructCoroutineInClosureShim { receiver_by_ref, .. } =
123+
instance.def
124+
{
125+
coroutine_kind = ty::ClosureKind::FnOnce;
124126

125-
let env_ty = tcx.closure_env_ty(coroutine_ty, coroutine_kind, env_region);
127+
// Implementations of `FnMut` and `Fn` for coroutine-closures
128+
// still take their receiver by ref.
129+
if receiver_by_ref { Ty::new_mut_ptr(tcx, coroutine_ty) } else { coroutine_ty }
130+
} else {
131+
tcx.closure_env_ty(coroutine_ty, coroutine_kind, env_region)
132+
};
126133

127134
let sig = sig.skip_binder();
128135
ty::Binder::bind_with_vars(

‎compiler/rustc_ty_utils/src/instance.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,7 @@ fn resolve_associated_item<'tcx>(
283283
Some(Instance {
284284
def: ty::InstanceDef::ConstructCoroutineInClosureShim {
285285
coroutine_closure_def_id,
286+
receiver_by_ref: target_kind != ty::ClosureKind::FnOnce,
286287
},
287288
args,
288289
})
@@ -310,6 +311,7 @@ fn resolve_associated_item<'tcx>(
310311
Some(Instance {
311312
def: ty::InstanceDef::ConstructCoroutineInClosureShim {
312313
coroutine_closure_def_id,
314+
receiver_by_ref: false,
313315
},
314316
args,
315317
})

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

Lines changed: 1 addition & 1 deletion
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:37:53: 40:10}, _2: ResumeTy) -> ()
3+
fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:42:53: 45: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

Lines changed: 1 addition & 1 deletion
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:37:53: 40:10}, _2: ResumeTy) -> ()
3+
fn main::{closure#0}::{closure#0}::{closure#0}(_1: {async closure body@$DIR/async_closure_shims.rs:42:53: 45: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

Lines changed: 3 additions & 3 deletions
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:37:33: 37:52}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:37:53: 40:10} {
4-
let mut _0: {async closure body@$DIR/async_closure_shims.rs:37:53: 40:10};
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};
55

66
bb0: {
7-
_0 = {coroutine@$DIR/async_closure_shims.rs:37:53: 40:10 (#0)} { a: move _2, b: move (_1.0: i32) };
7+
_0 = {coroutine@$DIR/async_closure_shims.rs:42:53: 45: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

Lines changed: 3 additions & 3 deletions
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:37:33: 37:52}, _2: i32) -> {async closure body@$DIR/async_closure_shims.rs:37:53: 40:10} {
4-
let mut _0: {async closure body@$DIR/async_closure_shims.rs:37:53: 40:10};
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};
55

66
bb0: {
7-
_0 = {coroutine@$DIR/async_closure_shims.rs:37:53: 40:10 (#0)} { a: move _2, b: move (_1.0: i32) };
7+
_0 = {coroutine@$DIR/async_closure_shims.rs:42:53: 45: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}.coroutine_closure_by_ref.0.panic-abort.mir

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
// MIR for `main::{closure#0}::{closure#1}` 0 coroutine_closure_by_ref
2+
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};
5+
6+
bb0: {
7+
_0 = {coroutine@$DIR/async_closure_shims.rs:49:49: 51:10 (#0)} { a: move _2 };
8+
return;
9+
}
10+
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
// MIR for `main::{closure#0}::{closure#1}` 0 coroutine_closure_by_ref
2+
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};
5+
6+
bb0: {
7+
_0 = {coroutine@$DIR/async_closure_shims.rs:49:49: 51:10 (#0)} { a: move _2 };
8+
return;
9+
}
10+
}

‎tests/mir-opt/async_closure_shims.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,13 @@ async fn call_once(f: impl AsyncFnOnce(i32)) {
2929
f(1).await;
3030
}
3131

32+
async fn call_normal<F: Future<Output = ()>>(f: &impl Fn(i32) -> F) {
33+
f(1).await;
34+
}
35+
3236
// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#0}.coroutine_closure_by_move.0.mir
3337
// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#0}-{closure#0}.coroutine_by_move.0.mir
38+
// EMIT_MIR async_closure_shims.main-{closure#0}-{closure#1}.coroutine_closure_by_ref.0.mir
3439
pub fn main() {
3540
block_on(async {
3641
let b = 2i32;
@@ -40,5 +45,10 @@ pub fn main() {
4045
};
4146
call_mut(&mut async_closure).await;
4247
call_once(async_closure).await;
48+
49+
let async_closure = async move |a: i32| {
50+
let a = &a;
51+
};
52+
call_normal(&async_closure).await;
4353
});
4454
}

0 commit comments

Comments
 (0)