bjoernager/rust
bjoernager/rust
1
Fork 0
Code Activity

Stop using a special inner body for the coroutine by-move body for async closures

Browse source
This commit is contained in:
Michael Goulet 2024-08-01 13:05:17 -04:00
parent 515395af0e
commit 4609841c07
40 changed files with 295 additions and 315 deletions
Download patch file Download diff file Expand all files Collapse all files

2
compiler/rustc_mir_transform/src/coroutine.rs
View file

@ -53,7 +53,7 @@
mod by_move_body;
use std::{iter, ops};
pub use by_move_body::ByMoveBody;
pub use by_move_body::coroutine_by_move_body_def_id;
use rustc_data_structures::fx::FxHashSet;
use rustc_errors::pluralize;
use rustc_hir as hir;

288
compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
View file

@ -69,161 +69,167 @@
//! in case precise captures (edition 2021 closure capture rules) caused the inner coroutine
//! to split one field capture into two.
use rustc_data_structures::steal::Steal;
use rustc_data_structures::unord::UnordMap;
use rustc_hir as hir;
use rustc_hir::def::DefKind;
use rustc_hir::def_id::{DefId, LocalDefId};
use rustc_middle::bug;
use rustc_middle::hir::place::{Projection, ProjectionKind};
use rustc_middle::mir::visit::MutVisitor;
use rustc_middle::mir::{self, dump_mir, MirPass};
use rustc_middle::mir::{self, dump_mir};
use rustc_middle::ty::{self, InstanceKind, Ty, TyCtxt, TypeVisitableExt};
use rustc_span::symbol::kw;
use rustc_target::abi::{FieldIdx, VariantIdx};
use crate::pass_manager::validate_body;
pub fn coroutine_by_move_body_def_id<'tcx>(
tcx: TyCtxt<'tcx>,
coroutine_def_id: LocalDefId,
) -> DefId {
let body = tcx.mir_built(coroutine_def_id).borrow();
pub struct ByMoveBody;
let Some(hir::CoroutineKind::Desugared(_, hir::CoroutineSource::Closure)) =
tcx.coroutine_kind(coroutine_def_id)
else {
bug!("should only be invoked on coroutine-closures");
};
impl<'tcx> MirPass<'tcx> for ByMoveBody {
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
// We only need to generate by-move coroutine bodies for coroutines that come
// from coroutine-closures.
let Some(coroutine_def_id) = body.source.def_id().as_local() else {
return;
};
let Some(hir::CoroutineKind::Desugared(_, hir::CoroutineSource::Closure)) =
tcx.coroutine_kind(coroutine_def_id)
else {
return;
};
// Also, let's skip processing any bodies with errors, since there's no guarantee
// the MIR body will be constructed well.
let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
// Also, let's skip processing any bodies with errors, since there's no guarantee
// the MIR body will be constructed well.
let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
if coroutine_ty.references_error() {
return;
}
let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
let args = args.as_coroutine();
// We don't need to generate a by-move coroutine if the coroutine body was
// produced by the `CoroutineKindShim`, since it's already by-move.
if matches!(body.source.instance, ty::InstanceKind::CoroutineKindShim { .. }) {
return;
}
let coroutine_kind = args.kind_ty().to_opt_closure_kind().unwrap();
let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
let args = args.as_coroutine();
let coroutine_kind = args.kind_ty().to_opt_closure_kind().unwrap();
let parent_def_id = tcx.local_parent(coroutine_def_id);
let ty::CoroutineClosure(_, parent_args) =
*tcx.type_of(parent_def_id).instantiate_identity().kind()
else {
bug!();
};
let parent_closure_args = parent_args.as_coroutine_closure();
let num_args = parent_closure_args
.coroutine_closure_sig()
.skip_binder()
.tupled_inputs_ty
.tuple_fields()
.len();
let field_remapping: UnordMap<_, _> = ty::analyze_coroutine_closure_captures(
tcx.closure_captures(parent_def_id).iter().copied(),
tcx.closure_captures(coroutine_def_id).iter().skip(num_args).copied(),
|(parent_field_idx, parent_capture), (child_field_idx, child_capture)| {
// Store this set of additional projections (fields and derefs).
// We need to re-apply them later.
let mut child_precise_captures = child_capture.place.projections
[parent_capture.place.projections.len()..]
.to_vec();
// If the parent capture is by-ref, then we need to apply an additional
// deref before applying any further projections to this place.
if parent_capture.is_by_ref() {
child_precise_captures.insert(
0,
Projection { ty: parent_capture.place.ty(), kind: ProjectionKind::Deref },
);
}
// If the child capture is by-ref, then we need to apply a "ref"
// projection (i.e. `&`) at the end. But wait! We don't have that
// as a projection kind. So instead, we can apply its dual and
// *peel* a deref off of the place when it shows up in the MIR body.
// Luckily, by construction this is always possible.
let peel_deref = if child_capture.is_by_ref() {
assert!(
parent_capture.is_by_ref() || coroutine_kind != ty::ClosureKind::FnOnce,
"`FnOnce` coroutine-closures return coroutines that capture from \
their body; it will always result in a borrowck error!"
);
true
} else {
false
};
// Regarding the behavior above, you may think that it's redundant to both
// insert a deref and then peel a deref if the parent and child are both
// captured by-ref. This would be correct, except for the case where we have
// precise capturing projections, since the inserted deref is to the *beginning*
// and the peeled deref is at the *end*. I cannot seem to actually find a
// case where this happens, though, but let's keep this code flexible.
// Finally, store the type of the parent's captured place. We need
// this when building the field projection in the MIR body later on.
let mut parent_capture_ty = parent_capture.place.ty();
parent_capture_ty = match parent_capture.info.capture_kind {
ty::UpvarCapture::ByValue => parent_capture_ty,
ty::UpvarCapture::ByRef(kind) => Ty::new_ref(
tcx,
tcx.lifetimes.re_erased,
parent_capture_ty,
kind.to_mutbl_lossy(),
),
};
(
FieldIdx::from_usize(child_field_idx + num_args),
(
FieldIdx::from_usize(parent_field_idx + num_args),
parent_capture_ty,
peel_deref,
child_precise_captures,
),
)
},
)
.collect();
if coroutine_kind == ty::ClosureKind::FnOnce {
assert_eq!(field_remapping.len(), tcx.closure_captures(parent_def_id).len());
return;
}
let by_move_coroutine_ty = tcx
.instantiate_bound_regions_with_erased(parent_closure_args.coroutine_closure_sig())
.to_coroutine_given_kind_and_upvars(
tcx,
parent_closure_args.parent_args(),
coroutine_def_id.to_def_id(),
ty::ClosureKind::FnOnce,
tcx.lifetimes.re_erased,
parent_closure_args.tupled_upvars_ty(),
parent_closure_args.coroutine_captures_by_ref_ty(),
);
let mut by_move_body = body.clone();
MakeByMoveBody { tcx, field_remapping, by_move_coroutine_ty }.visit_body(&mut by_move_body);
dump_mir(tcx, false, "coroutine_by_move", &0, &by_move_body, |_, _| Ok(()));
// Let's just always validate this body.
validate_body(tcx, &mut by_move_body, "Initial coroutine_by_move body".to_string());
// FIXME: use query feeding to generate the body right here and then only store the `DefId` of the new body.
by_move_body.source = mir::MirSource::from_instance(InstanceKind::CoroutineKindShim {
coroutine_def_id: coroutine_def_id.to_def_id(),
});
body.coroutine.as_mut().unwrap().by_move_body = Some(by_move_body);
let parent_def_id = tcx.local_parent(coroutine_def_id);
let ty::CoroutineClosure(_, parent_args) =
*tcx.type_of(parent_def_id).instantiate_identity().kind()
else {
bug!();
};
if parent_args.references_error() {
return coroutine_def_id.to_def_id();
}
let parent_closure_args = parent_args.as_coroutine_closure();
let num_args = parent_closure_args
.coroutine_closure_sig()
.skip_binder()
.tupled_inputs_ty
.tuple_fields()
.len();
let field_remapping: UnordMap<_, _> = ty::analyze_coroutine_closure_captures(
tcx.closure_captures(parent_def_id).iter().copied(),
tcx.closure_captures(coroutine_def_id).iter().skip(num_args).copied(),
|(parent_field_idx, parent_capture), (child_field_idx, child_capture)| {
// Store this set of additional projections (fields and derefs).
// We need to re-apply them later.
let mut child_precise_captures =
child_capture.place.projections[parent_capture.place.projections.len()..].to_vec();
// If the parent capture is by-ref, then we need to apply an additional
// deref before applying any further projections to this place.
if parent_capture.is_by_ref() {
child_precise_captures.insert(
0,
Projection { ty: parent_capture.place.ty(), kind: ProjectionKind::Deref },
);
}
// If the child capture is by-ref, then we need to apply a "ref"
// projection (i.e. `&`) at the end. But wait! We don't have that
// as a projection kind. So instead, we can apply its dual and
// *peel* a deref off of the place when it shows up in the MIR body.
// Luckily, by construction this is always possible.
let peel_deref = if child_capture.is_by_ref() {
assert!(
parent_capture.is_by_ref() || coroutine_kind != ty::ClosureKind::FnOnce,
"`FnOnce` coroutine-closures return coroutines that capture from \
their body; it will always result in a borrowck error!"
);
true
} else {
false
};
// Regarding the behavior above, you may think that it's redundant to both
// insert a deref and then peel a deref if the parent and child are both
// captured by-ref. This would be correct, except for the case where we have
// precise capturing projections, since the inserted deref is to the *beginning*
// and the peeled deref is at the *end*. I cannot seem to actually find a
// case where this happens, though, but let's keep this code flexible.
// Finally, store the type of the parent's captured place. We need
// this when building the field projection in the MIR body later on.
let mut parent_capture_ty = parent_capture.place.ty();
parent_capture_ty = match parent_capture.info.capture_kind {
ty::UpvarCapture::ByValue => parent_capture_ty,
ty::UpvarCapture::ByRef(kind) => Ty::new_ref(
tcx,
tcx.lifetimes.re_erased,
parent_capture_ty,
kind.to_mutbl_lossy(),
),
};
(
FieldIdx::from_usize(child_field_idx + num_args),
(
FieldIdx::from_usize(parent_field_idx + num_args),
parent_capture_ty,
peel_deref,
child_precise_captures,
),
)
},
)
.collect();
if coroutine_kind == ty::ClosureKind::FnOnce {
assert_eq!(field_remapping.len(), tcx.closure_captures(parent_def_id).len());
// The by-move body is just the body :)
return coroutine_def_id.to_def_id();
}
let by_move_coroutine_ty = tcx
.instantiate_bound_regions_with_erased(parent_closure_args.coroutine_closure_sig())
.to_coroutine_given_kind_and_upvars(
tcx,
parent_closure_args.parent_args(),
coroutine_def_id.to_def_id(),
ty::ClosureKind::FnOnce,
tcx.lifetimes.re_erased,
parent_closure_args.tupled_upvars_ty(),
parent_closure_args.coroutine_captures_by_ref_ty(),
);
let mut by_move_body = body.clone();
MakeByMoveBody { tcx, field_remapping, by_move_coroutine_ty }.visit_body(&mut by_move_body);
dump_mir(tcx, false, "coroutine_by_move", &0, &by_move_body, |_, _| Ok(()));
let body_def = tcx.create_def(coroutine_def_id, kw::Empty, DefKind::SyntheticCoroutineBody);
by_move_body.source =
mir::MirSource::from_instance(InstanceKind::Item(body_def.def_id().to_def_id()));
// Inherited from the by-ref coroutine.
body_def.codegen_fn_attrs(tcx.codegen_fn_attrs(coroutine_def_id).clone());
body_def.constness(tcx.constness(coroutine_def_id).clone());
body_def.coroutine_kind(tcx.coroutine_kind(coroutine_def_id).clone());
body_def.def_ident_span(tcx.def_ident_span(coroutine_def_id));
body_def.def_span(tcx.def_span(coroutine_def_id));
body_def.explicit_predicates_of(tcx.explicit_predicates_of(coroutine_def_id).clone());
body_def.generics_of(tcx.generics_of(coroutine_def_id).clone());
body_def.param_env(tcx.param_env(coroutine_def_id).clone());
body_def.predicates_of(tcx.predicates_of(coroutine_def_id).clone());
// The type of the coroutine is the `by_move_coroutine_ty`.
body_def.type_of(ty::EarlyBinder::bind(by_move_coroutine_ty));
body_def.mir_built(tcx.arena.alloc(Steal::new(by_move_body)));
body_def.def_id().to_def_id()
}
struct MakeByMoveBody<'tcx> {

1
compiler/rustc_mir_transform/src/inline.rs
View file

@ -341,7 +341,6 @@ impl<'tcx> Inliner<'tcx> {
| InstanceKind::FnPtrShim(..)
| InstanceKind::ClosureOnceShim { .. }
| InstanceKind::ConstructCoroutineInClosureShim { .. }
| InstanceKind::CoroutineKindShim { .. }
| InstanceKind::DropGlue(..)
| InstanceKind::CloneShim(..)
| InstanceKind::ThreadLocalShim(..)

1
compiler/rustc_mir_transform/src/inline/cycle.rs
View file

@ -88,7 +88,6 @@ pub(crate) fn mir_callgraph_reachable<'tcx>(
| InstanceKind::FnPtrShim(..)
| InstanceKind::ClosureOnceShim { .. }
| InstanceKind::ConstructCoroutineInClosureShim { .. }
| InstanceKind::CoroutineKindShim { .. }
| InstanceKind::ThreadLocalShim { .. }
| InstanceKind::CloneShim(..) => {}

42
compiler/rustc_mir_transform/src/lib.rs
View file

@ -135,6 +135,7 @@ pub fn provide(providers: &mut Providers) {
mir_inliner_callees: inline::cycle::mir_inliner_callees,
promoted_mir,
deduced_param_attrs: deduce_param_attrs::deduced_param_attrs,
coroutine_by_move_body_def_id: coroutine::coroutine_by_move_body_def_id,
..providers.queries
};
}
@ -293,10 +294,6 @@ fn mir_built(tcx: TyCtxt<'_>, def: LocalDefId) -> &Steal<Body<'_>> {
&Lint(check_packed_ref::CheckPackedRef),
&Lint(check_const_item_mutation::CheckConstItemMutation),
&Lint(function_item_references::FunctionItemReferences),
// If this is an async closure's output coroutine, generate
// by-move and by-mut bodies if needed. We do this first so
// they can be optimized in lockstep with their parent bodies.
&coroutine::ByMoveBody,
// What we need to do constant evaluation.
&simplify::SimplifyCfg::Initial,
&rustc_peek::SanityCheck, // Just a lint
@ -329,8 +326,15 @@ fn mir_promoted(
| DefKind::AnonConst => tcx.mir_const_qualif(def),
_ => ConstQualifs::default(),
};
// has_ffi_unwind_calls query uses the raw mir, so make sure it is run.
// the `has_ffi_unwind_calls` query uses the raw mir, so make sure it is run.
tcx.ensure_with_value().has_ffi_unwind_calls(def);
// the `by_move_body` query uses the raw mir, so make sure it is run.
if tcx.needs_coroutine_by_move_body_def_id(def) {
tcx.ensure_with_value().coroutine_by_move_body_def_id(def);
}
let mut body = tcx.mir_built(def).steal();
if let Some(error_reported) = const_qualifs.tainted_by_errors {
body.tainted_by_errors = Some(error_reported);
@ -339,14 +343,6 @@ fn mir_promoted(
// Collect `required_consts` *before* promotion, so if there are any consts being promoted
// we still add them to the list in the outer MIR body.
RequiredConstsVisitor::compute_required_consts(&mut body);
// If this has an associated by-move async closure body, that doesn't get run through these
// passes itself, it gets "tagged along" by the pass manager. `RequiredConstsVisitor` is not
// a regular pass so we have to also apply it manually to the other body.
if let Some(coroutine) = body.coroutine.as_mut() {
if let Some(by_move_body) = coroutine.by_move_body.as_mut() {
RequiredConstsVisitor::compute_required_consts(by_move_body);
}
}
// What we need to run borrowck etc.
let promote_pass = promote_consts::PromoteTemps::default();
@ -398,7 +394,10 @@ fn mir_drops_elaborated_and_const_checked(tcx: TyCtxt<'_>, def: LocalDefId) -> &
if tcx.is_coroutine(def.to_def_id()) {
tcx.ensure_with_value().mir_coroutine_witnesses(def);
}
let mir_borrowck = tcx.mir_borrowck(def);
// We only need to borrowck non-synthetic MIR.
let tainted_by_errors =
if !tcx.is_synthetic_mir(def) { tcx.mir_borrowck(def).tainted_by_errors } else { None };
let is_fn_like = tcx.def_kind(def).is_fn_like();
if is_fn_like {
@ -410,7 +409,8 @@ fn mir_drops_elaborated_and_const_checked(tcx: TyCtxt<'_>, def: LocalDefId) -> &
let (body, _) = tcx.mir_promoted(def);
let mut body = body.steal();
if let Some(error_reported) = mir_borrowck.tainted_by_errors {
if let Some(error_reported) = tainted_by_errors {
body.tainted_by_errors = Some(error_reported);
}
@ -660,14 +660,6 @@ fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> {
// visited does not depend on the optimization level.
// We do not use `run_passes` for this as that might skip the pass if `injection_phase` is set.
mentioned_items::MentionedItems.run_pass(tcx, &mut body);
// If this has an associated by-move async closure body, that doesn't get run through these
// passes itself, it gets "tagged along" by the pass manager. Since we're not using the pass
// manager we have to do this by hand.
if let Some(coroutine) = body.coroutine.as_mut() {
if let Some(by_move_body) = coroutine.by_move_body.as_mut() {
mentioned_items::MentionedItems.run_pass(tcx, by_move_body);
}
}
// If `mir_drops_elaborated_and_const_checked` found that the current body has unsatisfiable
// predicates, it will shrink the MIR to a single `unreachable` terminator.
@ -690,7 +682,9 @@ fn promoted_mir(tcx: TyCtxt<'_>, def: LocalDefId) -> &IndexVec<Promoted, Body<'_
return tcx.arena.alloc(IndexVec::new());
}
tcx.ensure_with_value().mir_borrowck(def);
if !tcx.is_synthetic_mir(def) {
tcx.ensure_with_value().mir_borrowck(def);
}
let mut promoted = tcx.mir_promoted(def).1.steal();
for body in &mut promoted {

6
compiler/rustc_mir_transform/src/pass_manager.rs
View file

@ -182,12 +182,6 @@ fn run_passes_inner<'tcx>(
body.pass_count = 1;
}
if let Some(coroutine) = body.coroutine.as_mut() {
if let Some(by_move_body) = coroutine.by_move_body.as_mut() {
run_passes_inner(tcx, by_move_body, passes, phase_change, validate_each);
}
}
}
pub fn validate_body<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, when: String) {

12
compiler/rustc_mir_transform/src/shim.rs
View file

@ -78,15 +78,11 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceKind<'tcx>) -> Body<
receiver_by_ref,
} => build_construct_coroutine_by_move_shim(tcx, coroutine_closure_def_id, receiver_by_ref),
ty::InstanceKind::CoroutineKindShim { coroutine_def_id } => {
return tcx.optimized_mir(coroutine_def_id).coroutine_by_move_body().unwrap().clone();
}
ty::InstanceKind::DropGlue(def_id, ty) => {
// FIXME(#91576): Drop shims for coroutines aren't subject to the MIR passes at the end
// of this function. Is this intentional?
if let Some(ty::Coroutine(coroutine_def_id, args)) = ty.map(Ty::kind) {
let coroutine_body = tcx.optimized_mir(*coroutine_def_id);
if let Some(&ty::Coroutine(coroutine_def_id, args)) = ty.map(Ty::kind) {
let coroutine_body = tcx.optimized_mir(coroutine_def_id);
let ty::Coroutine(_, id_args) = *tcx.type_of(coroutine_def_id).skip_binder().kind()
else {
@ -105,7 +101,9 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceKind<'tcx>) -> Body<
args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(),
ty::ClosureKind::FnOnce
);
coroutine_body.coroutine_by_move_body().unwrap().coroutine_drop().unwrap()
tcx.optimized_mir(tcx.coroutine_by_move_body_def_id(coroutine_def_id))
.coroutine_drop()
.unwrap()
};
let mut body = EarlyBinder::bind(body.clone()).instantiate(tcx, args);

19
compiler/rustc_mir_transform/src/validate.rs
View file

@ -102,25 +102,6 @@ impl<'tcx> MirPass<'tcx> for Validator {
}
}
}
// Enforce that coroutine-closure layouts are identical.
if let Some(layout) = body.coroutine_layout_raw()
&& let Some(by_move_body) = body.coroutine_by_move_body()
&& let Some(by_move_layout) = by_move_body.coroutine_layout_raw()
{
// FIXME(async_closures): We could do other validation here?
if layout.variant_fields.len() != by_move_layout.variant_fields.len() {
cfg_checker.fail(
Location::START,
format!(
"Coroutine layout has different number of variant fields from \
by-move coroutine layout:\n\
layout: {layout:#?}\n\
by_move_layout: {by_move_layout:#?}",
),
);
}
}
}
}