Skip to content

Commit 62cf533

Browse files
compiler-errorscompiler-errors
committed
Only split by-ref/by-move futures for async closures
1 parent 84b0c33 commit 62cf533

34 files changed

+124
-424
lines changed

‎compiler/rustc_borrowck/src/type_check/input_output.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
8888
self.tcx(),
8989
ty::CoroutineArgsParts {
9090
parent_args: args.parent_args(),
91-
kind_ty: Ty::from_closure_kind(self.tcx(), args.kind()),
91+
kind_ty: Ty::from_closure_kind(self.tcx(), args.kind().max_for_coroutine()),
9292
return_ty: user_provided_sig.output(),
9393
tupled_upvars_ty,
9494
// For async closures, none of these can be annotated, so just fill

‎compiler/rustc_hir_typeck/src/callee.rs

+9-5
Original file line numberDiff line numberDiff line change
@@ -184,16 +184,20 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
184184
kind: TypeVariableOriginKind::TypeInference,
185185
span: callee_expr.span,
186186
});
187+
// We may actually receive a coroutine back whose kind is different
188+
// from the closure that this dispatched from. This is because when
189+
// we have no captures, we automatically implement `FnOnce`. This
190+
// impl forces the closure kind to `FnOnce` i.e. `u8`.
191+
let kind_ty = self.next_ty_var(TypeVariableOrigin {
192+
kind: TypeVariableOriginKind::TypeInference,
193+
span: callee_expr.span,
194+
});
187195
let call_sig = self.tcx.mk_fn_sig(
188196
[coroutine_closure_sig.tupled_inputs_ty],
189197
coroutine_closure_sig.to_coroutine(
190198
self.tcx,
191199
closure_args.parent_args(),
192-
// Inherit the kind ty of the closure, since we're calling this
193-
// coroutine with the most relaxed `AsyncFn*` trait that we can.
194-
// We don't necessarily need to do this here, but it saves us
195-
// computing one more infer var that will get constrained later.
196-
closure_args.kind_ty(),
200+
kind_ty,
197201
self.tcx.coroutine_for_closure(def_id),
198202
tupled_upvars_ty,
199203
),

‎compiler/rustc_hir_typeck/src/closure.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,10 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
268268
},
269269
);
270270

271+
let coroutine_kind_ty = self.next_ty_var(TypeVariableOrigin {
272+
kind: TypeVariableOriginKind::ClosureSynthetic,
273+
span: expr_span,
274+
});
271275
let coroutine_upvars_ty = self.next_ty_var(TypeVariableOrigin {
272276
kind: TypeVariableOriginKind::ClosureSynthetic,
273277
span: expr_span,
@@ -285,7 +289,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
285289
sig.to_coroutine(
286290
tcx,
287291
parent_args,
288-
closure_kind_ty,
292+
coroutine_kind_ty,
289293
tcx.coroutine_for_closure(expr_def_id),
290294
coroutine_upvars_ty,
291295
)

‎compiler/rustc_hir_typeck/src/upvar.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -410,7 +410,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
410410
self.demand_eqtype(
411411
span,
412412
coroutine_args.as_coroutine().kind_ty(),
413-
Ty::from_closure_kind(self.tcx, closure_kind),
413+
Ty::from_closure_kind(self.tcx, closure_kind.max_for_coroutine()),
414414
);
415415
}
416416

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

-12
Original file line numberDiff line numberDiff line change
@@ -278,13 +278,6 @@ pub struct CoroutineInfo<'tcx> {
278278
/// using `run_passes`.
279279
pub by_move_body: Option<Body<'tcx>>,
280280

281-
/// The body of the coroutine, modified to take its upvars by mutable ref rather than by
282-
/// immutable ref.
283-
///
284-
/// FIXME(async_closures): This is literally the same body as the parent body. Find a better
285-
/// way to represent the by-mut signature (or cap the closure-kind of the coroutine).
286-
pub by_mut_body: Option<Body<'tcx>>,
287-
288281
/// The layout of a coroutine. This field is populated after the state transform pass.
289282
pub coroutine_layout: Option<CoroutineLayout<'tcx>>,
290283

@@ -305,7 +298,6 @@ impl<'tcx> CoroutineInfo<'tcx> {
305298
yield_ty: Some(yield_ty),
306299
resume_ty: Some(resume_ty),
307300
by_move_body: None,
308-
by_mut_body: None,
309301
coroutine_drop: None,
310302
coroutine_layout: None,
311303
}
@@ -620,10 +612,6 @@ impl<'tcx> Body<'tcx> {
620612
self.coroutine.as_ref()?.by_move_body.as_ref()
621613
}
622614

623-
pub fn coroutine_by_mut_body(&self) -> Option<&Body<'tcx>> {
624-
self.coroutine.as_ref()?.by_mut_body.as_ref()
625-
}
626-
627615
#[inline]
628616
pub fn coroutine_kind(&self) -> Option<CoroutineKind> {
629617
self.coroutine.as_ref().map(|coroutine| coroutine.coroutine_kind)

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

+5-2
Original file line numberDiff line numberDiff line change
@@ -345,8 +345,11 @@ macro_rules! make_mir_visitor {
345345
ty::InstanceDef::Virtual(_def_id, _) |
346346
ty::InstanceDef::ThreadLocalShim(_def_id) |
347347
ty::InstanceDef::ClosureOnceShim { call_once: _def_id, track_caller: _ } |
348-
ty::InstanceDef::ConstructCoroutineInClosureShim { coroutine_closure_def_id: _def_id, target_kind: _ } |
349-
ty::InstanceDef::CoroutineKindShim { coroutine_def_id: _def_id, target_kind: _ } |
348+
ty::InstanceDef::ConstructCoroutineInClosureShim {
349+
coroutine_closure_def_id: _def_id,
350+
receiver_by_ref: _,
351+
} |
352+
ty::InstanceDef::CoroutineKindShim { coroutine_def_id: _def_id } |
350353
ty::InstanceDef::DropGlue(_def_id, None) => {}
351354

352355
ty::InstanceDef::FnPtrShim(_def_id, ty) |

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

+6-13
Original file line numberDiff line numberDiff line change
@@ -90,24 +90,20 @@ pub enum InstanceDef<'tcx> {
9090
/// and dispatch to the `FnMut::call_mut` instance for the closure.
9191
ClosureOnceShim { call_once: DefId, track_caller: bool },
9292

93-
/// `<[FnMut/Fn coroutine-closure] as FnOnce>::call_once` or
94-
/// `<[Fn coroutine-closure] as FnMut>::call_mut`.
93+
/// `<[FnMut/Fn coroutine-closure] as FnOnce>::call_once`
9594
///
9695
/// The body generated here differs significantly from the `ClosureOnceShim`,
9796
/// since we need to generate a distinct coroutine type that will move the
9897
/// closure's upvars *out* of the closure.
99-
ConstructCoroutineInClosureShim {
100-
coroutine_closure_def_id: DefId,
101-
target_kind: ty::ClosureKind,
102-
},
98+
ConstructCoroutineInClosureShim { coroutine_closure_def_id: DefId, receiver_by_ref: bool },
10399

104100
/// `<[coroutine] as Future>::poll`, but for coroutines produced when `AsyncFnOnce`
105101
/// is called on a coroutine-closure whose closure kind greater than `FnOnce`, or
106102
/// similarly for `AsyncFnMut`.
107103
///
108104
/// This will select the body that is produced by the `ByMoveBody` transform, and thus
109105
/// take and use all of its upvars by-move rather than by-ref.
110-
CoroutineKindShim { coroutine_def_id: DefId, target_kind: ty::ClosureKind },
106+
CoroutineKindShim { coroutine_def_id: DefId },
111107

112108
/// Compiler-generated accessor for thread locals which returns a reference to the thread local
113109
/// the `DefId` defines. This is used to export thread locals from dylibs on platforms lacking
@@ -192,9 +188,9 @@ impl<'tcx> InstanceDef<'tcx> {
192188
| InstanceDef::ClosureOnceShim { call_once: def_id, track_caller: _ }
193189
| ty::InstanceDef::ConstructCoroutineInClosureShim {
194190
coroutine_closure_def_id: def_id,
195-
target_kind: _,
191+
receiver_by_ref: _,
196192
}
197-
| ty::InstanceDef::CoroutineKindShim { coroutine_def_id: def_id, target_kind: _ }
193+
| ty::InstanceDef::CoroutineKindShim { coroutine_def_id: def_id }
198194
| InstanceDef::DropGlue(def_id, _)
199195
| InstanceDef::CloneShim(def_id, _)
200196
| InstanceDef::FnPtrAddrShim(def_id, _) => def_id,
@@ -651,10 +647,7 @@ impl<'tcx> Instance<'tcx> {
651647
Some(Instance { def: ty::InstanceDef::Item(coroutine_def_id), args })
652648
} else {
653649
Some(Instance {
654-
def: ty::InstanceDef::CoroutineKindShim {
655-
coroutine_def_id,
656-
target_kind: args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(),
657-
},
650+
def: ty::InstanceDef::CoroutineKindShim { coroutine_def_id },
658651
args,
659652
})
660653
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -485,7 +485,7 @@ impl<'tcx> CoroutineClosureSignature<'tcx> {
485485
self.to_coroutine(
486486
tcx,
487487
parent_args,
488-
Ty::from_closure_kind(tcx, goal_kind),
488+
Ty::from_closure_kind(tcx, goal_kind.max_for_coroutine()),
489489
coroutine_def_id,
490490
tupled_upvars_ty,
491491
)

‎compiler/rustc_mir_transform/src/coroutine/by_move_body.rs

-35
Original file line numberDiff line numberDiff line change
@@ -67,45 +67,10 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
6767
by_move_body.source = mir::MirSource {
6868
instance: InstanceDef::CoroutineKindShim {
6969
coroutine_def_id: coroutine_def_id.to_def_id(),
70-
target_kind: ty::ClosureKind::FnOnce,
7170
},
7271
promoted: None,
7372
};
7473
body.coroutine.as_mut().unwrap().by_move_body = Some(by_move_body);
75-
76-
// If this is coming from an `AsyncFn` coroutine-closure, we must also create a by-mut body.
77-
// This is actually just a copy of the by-ref body, but with a different self type.
78-
// FIXME(async_closures): We could probably unify this with the by-ref body somehow.
79-
if coroutine_kind == ty::ClosureKind::Fn {
80-
let by_mut_coroutine_ty = Ty::new_coroutine(
81-
tcx,
82-
coroutine_def_id.to_def_id(),
83-
ty::CoroutineArgs::new(
84-
tcx,
85-
ty::CoroutineArgsParts {
86-
parent_args: args.as_coroutine().parent_args(),
87-
kind_ty: Ty::from_closure_kind(tcx, ty::ClosureKind::FnMut),
88-
resume_ty: args.as_coroutine().resume_ty(),
89-
yield_ty: args.as_coroutine().yield_ty(),
90-
return_ty: args.as_coroutine().return_ty(),
91-
witness: args.as_coroutine().witness(),
92-
tupled_upvars_ty: args.as_coroutine().tupled_upvars_ty(),
93-
},
94-
)
95-
.args,
96-
);
97-
let mut by_mut_body = body.clone();
98-
by_mut_body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty = by_mut_coroutine_ty;
99-
dump_mir(tcx, false, "coroutine_by_mut", &0, &by_mut_body, |_, _| Ok(()));
100-
by_mut_body.source = mir::MirSource {
101-
instance: InstanceDef::CoroutineKindShim {
102-
coroutine_def_id: coroutine_def_id.to_def_id(),
103-
target_kind: ty::ClosureKind::FnMut,
104-
},
105-
promoted: None,
106-
};
107-
body.coroutine.as_mut().unwrap().by_mut_body = Some(by_mut_body);
108-
}
10974
}
11075
}
11176

‎compiler/rustc_mir_transform/src/pass_manager.rs

-3
Original file line numberDiff line numberDiff line change
@@ -186,9 +186,6 @@ fn run_passes_inner<'tcx>(
186186
if let Some(by_move_body) = coroutine.by_move_body.as_mut() {
187187
run_passes_inner(tcx, by_move_body, passes, phase_change, validate_each);
188188
}
189-
if let Some(by_mut_body) = coroutine.by_mut_body.as_mut() {
190-
run_passes_inner(tcx, by_mut_body, passes, phase_change, validate_each);
191-
}
192189
}
193190
}
194191

‎compiler/rustc_mir_transform/src/shim.rs

+18-85
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ use rustc_hir::def_id::DefId;
33
use rustc_hir::lang_items::LangItem;
44
use rustc_middle::mir::*;
55
use rustc_middle::query::Providers;
6+
use rustc_middle::ty::GenericArgs;
67
use rustc_middle::ty::{self, CoroutineArgs, EarlyBinder, Ty, TyCtxt};
7-
use rustc_middle::ty::{GenericArgs, CAPTURE_STRUCT_LOCAL};
88
use rustc_target::abi::{FieldIdx, VariantIdx, FIRST_VARIANT};
99

1010
use rustc_index::{Idx, IndexVec};
@@ -68,37 +68,12 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
6868

6969
ty::InstanceDef::ConstructCoroutineInClosureShim {
7070
coroutine_closure_def_id,
71-
target_kind,
72-
} => match target_kind {
73-
ty::ClosureKind::Fn => unreachable!("shouldn't be building shim for Fn"),
74-
ty::ClosureKind::FnMut => {
75-
// No need to optimize the body, it has already been optimized
76-
// since we steal it from the `AsyncFn::call` body and just fix
77-
// the return type.
78-
return build_construct_coroutine_by_mut_shim(tcx, coroutine_closure_def_id);
79-
}
80-
ty::ClosureKind::FnOnce => {
81-
build_construct_coroutine_by_move_shim(tcx, coroutine_closure_def_id)
82-
}
83-
},
71+
receiver_by_ref,
72+
} => build_construct_coroutine_by_move_shim(tcx, coroutine_closure_def_id, receiver_by_ref),
8473

85-
ty::InstanceDef::CoroutineKindShim { coroutine_def_id, target_kind } => match target_kind {
86-
ty::ClosureKind::Fn => unreachable!(),
87-
ty::ClosureKind::FnMut => {
88-
return tcx
89-
.optimized_mir(coroutine_def_id)
90-
.coroutine_by_mut_body()
91-
.unwrap()
92-
.clone();
93-
}
94-
ty::ClosureKind::FnOnce => {
95-
return tcx
96-
.optimized_mir(coroutine_def_id)
97-
.coroutine_by_move_body()
98-
.unwrap()
99-
.clone();
100-
}
101-
},
74+
ty::InstanceDef::CoroutineKindShim { coroutine_def_id } => {
75+
return tcx.optimized_mir(coroutine_def_id).coroutine_by_move_body().unwrap().clone();
76+
}
10277

10378
ty::InstanceDef::DropGlue(def_id, ty) => {
10479
// FIXME(#91576): Drop shims for coroutines aren't subject to the MIR passes at the end
@@ -119,21 +94,11 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
11994
let body = if id_args.as_coroutine().kind_ty() == args.as_coroutine().kind_ty() {
12095
coroutine_body.coroutine_drop().unwrap()
12196
} else {
122-
match args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap() {
123-
ty::ClosureKind::Fn => {
124-
unreachable!()
125-
}
126-
ty::ClosureKind::FnMut => coroutine_body
127-
.coroutine_by_mut_body()
128-
.unwrap()
129-
.coroutine_drop()
130-
.unwrap(),
131-
ty::ClosureKind::FnOnce => coroutine_body
132-
.coroutine_by_move_body()
133-
.unwrap()
134-
.coroutine_drop()
135-
.unwrap(),
136-
}
97+
assert_eq!(
98+
args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(),
99+
ty::ClosureKind::FnOnce
100+
);
101+
coroutine_body.coroutine_by_move_body().unwrap().coroutine_drop().unwrap()
137102
};
138103

139104
let mut body = EarlyBinder::bind(body.clone()).instantiate(tcx, args);
@@ -1047,12 +1012,17 @@ fn build_fn_ptr_addr_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, self_ty: Ty<'t
10471012
fn build_construct_coroutine_by_move_shim<'tcx>(
10481013
tcx: TyCtxt<'tcx>,
10491014
coroutine_closure_def_id: DefId,
1015+
receiver_by_ref: bool,
10501016
) -> Body<'tcx> {
1051-
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();
10521018
let ty::CoroutineClosure(_, args) = *self_ty.kind() else {
10531019
bug!();
10541020
};
10551021

1022+
if receiver_by_ref {
1023+
self_ty = Ty::new_mut_ptr(tcx, self_ty);
1024+
}
1025+
10561026
let poly_sig = args.as_coroutine_closure().coroutine_closure_sig().map_bound(|sig| {
10571027
tcx.mk_fn_sig(
10581028
[self_ty].into_iter().chain(sig.tupled_inputs_ty.tuple_fields()),
@@ -1108,7 +1078,7 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
11081078

11091079
let source = MirSource::from_instance(ty::InstanceDef::ConstructCoroutineInClosureShim {
11101080
coroutine_closure_def_id,
1111-
target_kind: ty::ClosureKind::FnOnce,
1081+
receiver_by_ref,
11121082
});
11131083

11141084
let body =
@@ -1117,40 +1087,3 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
11171087

11181088
body
11191089
}
1120-
1121-
fn build_construct_coroutine_by_mut_shim<'tcx>(
1122-
tcx: TyCtxt<'tcx>,
1123-
coroutine_closure_def_id: DefId,
1124-
) -> Body<'tcx> {
1125-
let mut body = tcx.optimized_mir(coroutine_closure_def_id).clone();
1126-
let coroutine_closure_ty = tcx.type_of(coroutine_closure_def_id).instantiate_identity();
1127-
let ty::CoroutineClosure(_, args) = *coroutine_closure_ty.kind() else {
1128-
bug!();
1129-
};
1130-
let args = args.as_coroutine_closure();
1131-
1132-
body.local_decls[RETURN_PLACE].ty =
1133-
tcx.instantiate_bound_regions_with_erased(args.coroutine_closure_sig().map_bound(|sig| {
1134-
sig.to_coroutine_given_kind_and_upvars(
1135-
tcx,
1136-
args.parent_args(),
1137-
tcx.coroutine_for_closure(coroutine_closure_def_id),
1138-
ty::ClosureKind::FnMut,
1139-
tcx.lifetimes.re_erased,
1140-
args.tupled_upvars_ty(),
1141-
args.coroutine_captures_by_ref_ty(),
1142-
)
1143-
}));
1144-
body.local_decls[CAPTURE_STRUCT_LOCAL].ty =
1145-
Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, coroutine_closure_ty);
1146-
1147-
body.source = MirSource::from_instance(ty::InstanceDef::ConstructCoroutineInClosureShim {
1148-
coroutine_closure_def_id,
1149-
target_kind: ty::ClosureKind::FnMut,
1150-
});
1151-
1152-
body.pass_count = 0;
1153-
dump_mir(tcx, false, "coroutine_closure_by_mut", &0, &body, |_, _| Ok(()));
1154-
1155-
body
1156-
}

‎compiler/rustc_span/src/symbol.rs

+1-2
Original file line numberDiff line numberDiff line change
@@ -167,9 +167,8 @@ symbols! {
167167
Break,
168168
C,
169169
CStr,
170-
CallFuture,
171-
CallMutFuture,
172170
CallOnceFuture,
171+
CallRefFuture,
173172
Capture,
174173
Center,
175174
Cleanup,

0 commit comments

Comments
 (0)