Skip to content

Commit ba9a24f

Browse files
Only split by-ref/by-move futures for async closures
1 parent d377991 commit ba9a24f

33 files changed

+109
-415
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_coroutine_closure_kind(self.tcx(), args.kind()),
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_coroutine_closure_kind(self.tcx, closure_kind),
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

+4-2
Original file line numberDiff line numberDiff line change
@@ -345,8 +345,10 @@ 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+
} |
351+
ty::InstanceDef::CoroutineKindShim { coroutine_def_id: _def_id } |
350352
ty::InstanceDef::DropGlue(_def_id, None) => {}
351353

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

compiler/rustc_middle/src/ty/instance.rs

+5-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 },
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,8 @@ 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: _,
196191
}
197-
| ty::InstanceDef::CoroutineKindShim { coroutine_def_id: def_id, target_kind: _ }
192+
| ty::InstanceDef::CoroutineKindShim { coroutine_def_id: def_id }
198193
| InstanceDef::DropGlue(def_id, _)
199194
| InstanceDef::CloneShim(def_id, _)
200195
| InstanceDef::FnPtrAddrShim(def_id, _) => def_id,
@@ -651,10 +646,7 @@ impl<'tcx> Instance<'tcx> {
651646
Some(Instance { def: ty::InstanceDef::Item(coroutine_def_id), args })
652647
} else {
653648
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-
},
649+
def: ty::InstanceDef::CoroutineKindShim { coroutine_def_id },
658650
args,
659651
})
660652
}

compiler/rustc_middle/src/ty/sty.rs

+16-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_coroutine_closure_kind(tcx, goal_kind),
489489
coroutine_def_id,
490490
tupled_upvars_ty,
491491
)
@@ -2430,6 +2430,21 @@ impl<'tcx> Ty<'tcx> {
24302430
}
24312431
}
24322432

2433+
/// Like [`Ty::to_opt_closure_kind`], but it caps the "maximum" closure kind
2434+
/// to `FnMut`. This is because although we have three capability states,
2435+
/// `AsyncFn`/`AsyncFnMut`/`AsyncFnOnce`, we only need to distinguish two coroutine
2436+
/// bodies: by-ref and by-value.
2437+
///
2438+
/// This method should be used when constructing a `Coroutine` out of a
2439+
/// `CoroutineClosure`, when the `Coroutine`'s `kind` field is being populated
2440+
/// directly from the `CoroutineClosure`'s `kind`.
2441+
pub fn from_coroutine_closure_kind(tcx: TyCtxt<'tcx>, kind: ty::ClosureKind) -> Ty<'tcx> {
2442+
match kind {
2443+
ty::ClosureKind::Fn | ty::ClosureKind::FnMut => tcx.types.i16,
2444+
ty::ClosureKind::FnOnce => tcx.types.i32,
2445+
}
2446+
}
2447+
24332448
/// Fast path helper for testing if a type is `Sized`.
24342449
///
24352450
/// Returning true means the type is known to be sized. Returning

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

+12-86
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};
@@ -66,39 +66,13 @@ 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 {
70-
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-
},
69+
ty::InstanceDef::ConstructCoroutineInClosureShim { coroutine_closure_def_id } => {
70+
build_construct_coroutine_by_move_shim(tcx, coroutine_closure_def_id)
71+
}
8472

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-
},
73+
ty::InstanceDef::CoroutineKindShim { coroutine_def_id } => {
74+
return tcx.optimized_mir(coroutine_def_id).coroutine_by_move_body().unwrap().clone();
75+
}
10276

10377
ty::InstanceDef::DropGlue(def_id, ty) => {
10478
// FIXME(#91576): Drop shims for coroutines aren't subject to the MIR passes at the end
@@ -119,21 +93,11 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
11993
let body = if id_args.as_coroutine().kind_ty() == args.as_coroutine().kind_ty() {
12094
coroutine_body.coroutine_drop().unwrap()
12195
} 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-
}
96+
assert_eq!(
97+
args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(),
98+
ty::ClosureKind::FnOnce
99+
);
100+
coroutine_body.coroutine_by_move_body().unwrap().coroutine_drop().unwrap()
137101
};
138102

139103
let mut body = EarlyBinder::bind(body.clone()).instantiate(tcx, args);
@@ -1108,7 +1072,6 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
11081072

11091073
let source = MirSource::from_instance(ty::InstanceDef::ConstructCoroutineInClosureShim {
11101074
coroutine_closure_def_id,
1111-
target_kind: ty::ClosureKind::FnOnce,
11121075
});
11131076

11141077
let body =
@@ -1117,40 +1080,3 @@ fn build_construct_coroutine_by_move_shim<'tcx>(
11171080

11181081
body
11191082
}
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)