More comments, final tweaks

This commit is contained in:
Michael Goulet 2024-01-25 19:17:21 +00:00
parent 37184e86ea
commit c98d6994a3
20 changed files with 169 additions and 42 deletions

View File

@ -1074,6 +1074,9 @@ fn lower_expr_coroutine_closure(
body,
fn_decl_span: self.lower_span(fn_decl_span),
fn_arg_span: Some(self.lower_span(fn_arg_span)),
// Lower this as a `CoroutineClosure`. That will ensure that HIR typeck
// knows that a `FnDecl` output type like `-> &str` actually means
// "coroutine that returns &str", rather than directly returning a `&str`.
kind: hir::ClosureKind::CoroutineClosure(hir::CoroutineDesugaring::Async),
constness: hir::Constness::NotConst,
});

View File

@ -61,6 +61,8 @@ pub(super) fn check_signature_annotation(&mut self, body: &Body<'tcx>) {
)),
"this needs to be modified if we're lowering non-async closures"
);
// Make sure to use the args from `DefiningTy` so the right NLL region vids are prepopulated
// into the type.
let args = args.as_coroutine_closure();
let tupled_upvars_ty = ty::CoroutineClosureSignature::tupled_upvars_by_closure_kind(
self.tcx(),
@ -87,11 +89,13 @@ pub(super) fn check_signature_annotation(&mut self, body: &Body<'tcx>) {
ty::CoroutineArgsParts {
parent_args: args.parent_args(),
kind_ty: Ty::from_closure_kind(self.tcx(), args.kind()),
return_ty: user_provided_sig.output(),
tupled_upvars_ty,
// For async closures, none of these can be annotated, so just fill
// them with fresh ty vars.
resume_ty: next_ty_var(),
yield_ty: next_ty_var(),
witness: next_ty_var(),
return_ty: user_provided_sig.output(),
tupled_upvars_ty: tupled_upvars_ty,
},
)
.args,

View File

@ -720,6 +720,14 @@ fn compute_inputs_and_output(
ty::Binder::dummy(inputs_and_output)
}
// Construct the signature of the CoroutineClosure for the purposes of borrowck.
// This is pretty straightforward -- we:
// 1. first grab the `coroutine_closure_sig`,
// 2. compute the self type (`&`/`&mut`/no borrow),
// 3. flatten the tupled_input_tys,
// 4. construct the correct generator type to return with
// `CoroutineClosureSignature::to_coroutine_given_kind_and_upvars`.
// Then we wrap it all up into a list of inputs and output.
DefiningTy::CoroutineClosure(def_id, args) => {
assert_eq!(self.mir_def.to_def_id(), def_id);
let closure_sig = args.as_coroutine_closure().coroutine_closure_sig();

View File

@ -1533,15 +1533,8 @@ fn coroutine_kind(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<hir::CoroutineK
}
fn coroutine_for_closure(tcx: TyCtxt<'_>, def_id: LocalDefId) -> DefId {
let Node::Expr(&hir::Expr {
kind:
hir::ExprKind::Closure(&rustc_hir::Closure {
kind: hir::ClosureKind::CoroutineClosure(_),
body,
..
}),
..
}) = tcx.hir_node_by_def_id(def_id)
let &rustc_hir::Closure { kind: hir::ClosureKind::CoroutineClosure(_), body, .. } =
tcx.hir_node_by_def_id(def_id).expect_closure()
else {
bug!()
};

View File

@ -166,6 +166,11 @@ fn try_overloaded_call_step(
return Some(CallStep::DeferredClosure(def_id, closure_sig));
}
// When calling a `CoroutineClosure` that is local to the body, we will
// not know what its `closure_kind` is yet. Instead, just fill in the
// signature with an infer var for the `tupled_upvars_ty` of the coroutine,
// and record a deferred call resolution which will constrain that var
// as part of `AsyncFn*` trait confirmation.
ty::CoroutineClosure(def_id, args) if self.closure_kind(adjusted_ty).is_none() => {
let def_id = def_id.expect_local();
let closure_args = args.as_coroutine_closure();
@ -183,6 +188,10 @@ fn try_overloaded_call_step(
coroutine_closure_sig.to_coroutine(
self.tcx,
closure_args.parent_args(),
// Inherit the kind ty of the closure, since we're calling this
// coroutine with the most relaxed `AsyncFn*` trait that we can.
// We don't necessarily need to do this here, but it saves us
// computing one more infer var that will get constrained later.
closure_args.kind_ty(),
self.tcx.coroutine_for_closure(def_id),
tupled_upvars_ty,

View File

@ -175,6 +175,10 @@ pub fn check_expr_closure(
interior,
));
// Coroutines that come from coroutine closures have not yet determined
// their kind ty, so make a fresh infer var which will be constrained
// later during upvar analysis. Regular coroutines always have the kind
// ty of `().`
let kind_ty = match kind {
hir::CoroutineKind::Desugared(_, hir::CoroutineSource::Closure) => self
.next_ty_var(TypeVariableOrigin {
@ -203,6 +207,8 @@ pub fn check_expr_closure(
)
}
hir::ClosureKind::CoroutineClosure(kind) => {
// async closures always return the type ascribed after the `->` (if present),
// and yield `()`.
let (bound_return_ty, bound_yield_ty) = match kind {
hir::CoroutineDesugaring::Async => {
(bound_sig.skip_binder().output(), tcx.types.unit)
@ -211,6 +217,7 @@ pub fn check_expr_closure(
todo!("`gen` and `async gen` closures not supported yet")
}
};
// Compute all of the variables that will be used to populate the coroutine.
let resume_ty = self.next_ty_var(TypeVariableOrigin {
kind: TypeVariableOriginKind::ClosureSynthetic,
span: expr_span,
@ -258,9 +265,14 @@ pub fn check_expr_closure(
kind: TypeVariableOriginKind::ClosureSynthetic,
span: expr_span,
});
liberated_sig = tcx.mk_fn_sig(
liberated_sig.inputs().iter().copied(),
tcx.liberate_late_bound_regions(
// We need to turn the liberated signature that we got from HIR, which
// looks something like `|Args...| -> T`, into a signature that is suitable
// for type checking the inner body of the closure, which always returns a
// coroutine. To do so, we use the `CoroutineClosureSignature` to compute
// the coroutine type, filling in the tupled_upvars_ty and kind_ty with infer
// vars which will get constrained during upvar analysis.
let coroutine_output_ty = tcx.liberate_late_bound_regions(
expr_def_id.to_def_id(),
closure_args.coroutine_closure_sig().map_bound(|sig| {
sig.to_coroutine(
@ -271,7 +283,10 @@ pub fn check_expr_closure(
coroutine_upvars_ty,
)
}),
),
);
liberated_sig = tcx.mk_fn_sig(
liberated_sig.inputs().iter().copied(),
coroutine_output_ty,
liberated_sig.c_variadic,
liberated_sig.unsafety,
liberated_sig.abi,

View File

@ -335,7 +335,7 @@ pub(in super::super) fn check_argument_types(
continue;
}
// For this check, we do *not* want to treat async coroutine-closures (async blocks)
// For this check, we do *not* want to treat async coroutine closures (async blocks)
// as proper closures. Doing so would regress type inference when feeding
// the return value of an argument-position async block to an argument-position
// closure wrapped in a block.

View File

@ -336,6 +336,9 @@ fn analyze_closure(
}
}
// For coroutine-closures, we additionally must compute the
// `coroutine_captures_by_ref_ty` type, which is used to generate the by-ref
// version of the coroutine-closure's output coroutine.
if let UpvarArgs::CoroutineClosure(args) = args {
let closure_env_region: ty::Region<'_> = ty::Region::new_bound(
self.tcx,
@ -353,7 +356,8 @@ fn analyze_closure(
self.tcx.coroutine_for_closure(closure_def_id).expect_local(),
)
// Skip the captures that are just moving the closure's args
// into the coroutine. These are always by move.
// into the coroutine. These are always by move, and we append
// those later in the `CoroutineClosureSignature` helper functions.
.skip(
args.as_coroutine_closure()
.coroutine_closure_sig()
@ -365,6 +369,9 @@ fn analyze_closure(
.map(|captured_place| {
let upvar_ty = captured_place.place.ty();
let capture = captured_place.info.capture_kind;
// Not all upvars are captured by ref, so use
// `apply_capture_kind_on_capture_ty` to ensure that we
// compute the right captured type.
apply_capture_kind_on_capture_ty(
self.tcx,
upvar_ty,
@ -394,6 +401,7 @@ fn analyze_closure(
coroutine_captures_by_ref_ty,
);
// Additionally, we can now constrain the coroutine's kind type.
let ty::Coroutine(_, coroutine_args) =
*self.typeck_results.borrow().expr_ty(body.value).kind()
else {

View File

@ -145,11 +145,11 @@ fn add_kind(&mut self, kind: &ty::TyKind<'_>) {
self.flags -= TypeFlags::STILL_FURTHER_SPECIALIZABLE;
}
self.add_ty(args.signature_parts_ty());
self.add_ty(args.coroutine_witness_ty());
self.add_ty(args.coroutine_captures_by_ref_ty());
self.add_ty(args.kind_ty());
self.add_ty(args.signature_parts_ty());
self.add_ty(args.tupled_upvars_ty());
self.add_ty(args.coroutine_captures_by_ref_ty());
self.add_ty(args.coroutine_witness_ty());
}
&ty::Bound(debruijn, _) => {

View File

@ -646,9 +646,16 @@ pub fn try_resolve_item_for_coroutine(
bug!()
};
// If the closure's kind ty disagrees with the identity closure's kind ty,
// then this must be a coroutine generated by one of the `ConstructCoroutineInClosureShim`s.
if args.as_coroutine().kind_ty() == id_args.as_coroutine().kind_ty() {
Some(Instance { def: ty::InstanceDef::Item(coroutine_def_id), args })
} else {
assert_eq!(
args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(),
ty::ClosureKind::FnOnce,
"FIXME(async_closures): Generate a by-mut body here."
);
Some(Instance {
def: ty::InstanceDef::CoroutineByMoveShim { coroutine_def_id },
args,

View File

@ -275,6 +275,10 @@ pub struct CoroutineClosureArgs<'tcx> {
pub args: GenericArgsRef<'tcx>,
}
/// See docs for explanation of how each argument is used.
///
/// See [`CoroutineClosureSignature`] for how these arguments are put together
/// to make a callable [`FnSig`] suitable for typeck and borrowck.
pub struct CoroutineClosureArgsParts<'tcx> {
/// This is the args of the typeck root.
pub parent_args: &'tcx [GenericArg<'tcx>],
@ -409,17 +413,32 @@ pub struct CoroutineClosureSignature<'tcx> {
pub resume_ty: Ty<'tcx>,
pub yield_ty: Ty<'tcx>,
pub return_ty: Ty<'tcx>,
// Like the `fn_sig_as_fn_ptr_ty` of a regular closure, these types
// never actually differ. But we save them rather than recreating them
// from scratch just for good measure.
/// Always false
pub c_variadic: bool,
/// Always [`hir::Unsafety::Normal`]
pub unsafety: hir::Unsafety,
/// Always [`abi::Abi::RustCall`]
pub abi: abi::Abi,
}
impl<'tcx> CoroutineClosureSignature<'tcx> {
/// Construct a coroutine from the closure signature. Since a coroutine signature
/// is agnostic to the type of generator that is returned (by-ref/by-move),
/// the caller must specify what "flavor" of generator that they'd like to
/// create. Additionally, they must manually compute the upvars of the closure.
///
/// This helper is not really meant to be used directly except for early on
/// during typeck, when we want to put inference vars into the kind and upvars tys.
/// When the kind and upvars are known, use the other helper functions.
pub fn to_coroutine(
self,
tcx: TyCtxt<'tcx>,
parent_args: &'tcx [GenericArg<'tcx>],
kind_ty: Ty<'tcx>,
coroutine_kind_ty: Ty<'tcx>,
coroutine_def_id: DefId,
tupled_upvars_ty: Ty<'tcx>,
) -> Ty<'tcx> {
@ -427,7 +446,7 @@ pub fn to_coroutine(
tcx,
ty::CoroutineArgsParts {
parent_args,
kind_ty,
kind_ty: coroutine_kind_ty,
resume_ty: self.resume_ty,
yield_ty: self.yield_ty,
return_ty: self.return_ty,
@ -439,19 +458,24 @@ pub fn to_coroutine(
Ty::new_coroutine(tcx, coroutine_def_id, coroutine_args.args)
}
/// Given known upvars and a [`ClosureKind`](ty::ClosureKind), compute the coroutine
/// returned by that corresponding async fn trait.
///
/// This function expects the upvars to have been computed already, and doesn't check
/// that the `ClosureKind` is actually supported by the coroutine-closure.
pub fn to_coroutine_given_kind_and_upvars(
self,
tcx: TyCtxt<'tcx>,
parent_args: &'tcx [GenericArg<'tcx>],
coroutine_def_id: DefId,
closure_kind: ty::ClosureKind,
goal_kind: ty::ClosureKind,
env_region: ty::Region<'tcx>,
closure_tupled_upvars_ty: Ty<'tcx>,
coroutine_captures_by_ref_ty: Ty<'tcx>,
) -> Ty<'tcx> {
let tupled_upvars_ty = Self::tupled_upvars_by_closure_kind(
tcx,
closure_kind,
goal_kind,
self.tupled_inputs_ty,
closure_tupled_upvars_ty,
coroutine_captures_by_ref_ty,
@ -461,13 +485,21 @@ pub fn to_coroutine_given_kind_and_upvars(
self.to_coroutine(
tcx,
parent_args,
Ty::from_closure_kind(tcx, closure_kind),
Ty::from_closure_kind(tcx, goal_kind),
coroutine_def_id,
tupled_upvars_ty,
)
}
/// Given a closure kind, compute the tupled upvars that the given coroutine would return.
/// Compute the tupled upvars that a coroutine-closure's output coroutine
/// would return for the given `ClosureKind`.
///
/// When `ClosureKind` is `FnMut`/`Fn`, then this will use the "captures by ref"
/// to return a set of upvars which are borrowed with the given `env_region`.
///
/// This ensures that the `AsyncFn::call` will return a coroutine whose upvars'
/// lifetimes are related to the lifetime of the borrow on the closure made for
/// the call. This allows borrowck to enforce the self-borrows correctly.
pub fn tupled_upvars_by_closure_kind(
tcx: TyCtxt<'tcx>,
kind: ty::ClosureKind,

View File

@ -1149,6 +1149,9 @@ pub fn iter_fields<'tcx>(
ty::Closure(_, args) => {
iter_fields(args.as_closure().tupled_upvars_ty(), tcx, param_env, f);
}
ty::Coroutine(_, args) => {
iter_fields(args.as_coroutine().tupled_upvars_ty(), tcx, param_env, f);
}
ty::CoroutineClosure(_, args) => {
iter_fields(args.as_coroutine_closure().tupled_upvars_ty(), tcx, param_env, f);
}

View File

@ -1,3 +1,8 @@
//! A MIR pass which duplicates a coroutine's body and removes any derefs which
//! would be present for upvars that are taken by-ref. The result of which will
//! be a coroutine body that takes all of its upvars by-move, and which we stash
//! into the `CoroutineInfo` for all coroutines returned by coroutine-closures.
use rustc_data_structures::fx::FxIndexSet;
use rustc_hir as hir;
use rustc_middle::mir::visit::MutVisitor;
@ -87,11 +92,16 @@ fn visit_place(
&& self.by_ref_fields.contains(&idx)
{
let (begin, end) = place.projection[1..].split_first().unwrap();
// FIXME(async_closures): I'm actually a bit surprised to see that we always
// initially deref the by-ref upvars. If this is not actually true, then we
// will at least get an ICE that explains why this isn't true :^)
assert_eq!(*begin, mir::ProjectionElem::Deref);
// Peel one ref off of the ty.
let peeled_ty = ty.builtin_deref(true).unwrap().ty;
*place = mir::Place {
local: place.local,
projection: self.tcx.mk_place_elems_from_iter(
[mir::ProjectionElem::Field(idx, ty.builtin_deref(true).unwrap().ty)]
[mir::ProjectionElem::Field(idx, peeled_ty)]
.into_iter()
.chain(end.iter().copied()),
),
@ -101,6 +111,7 @@ fn visit_place(
}
fn visit_local_decl(&mut self, local: mir::Local, local_decl: &mut mir::LocalDecl<'tcx>) {
// Replace the type of the self arg.
if local == ty::CAPTURE_STRUCT_LOCAL {
local_decl.ty = self.by_move_coroutine_ty;
}

View File

@ -309,14 +309,19 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_callable<'tcx>(
}
// Returns a binder of the tupled inputs types, output type, and coroutine type
// from a builtin async closure type.
// from a builtin coroutine-closure type. If we don't yet know the closure kind of
// the coroutine-closure, emit an additional trait predicate for `AsyncFnKindHelper`
// which enforces the closure is actually callable with the given trait. When we
// know the kind already, we can short-circuit this check.
pub(in crate::solve) fn extract_tupled_inputs_and_output_from_async_callable<'tcx>(
tcx: TyCtxt<'tcx>,
self_ty: Ty<'tcx>,
goal_kind: ty::ClosureKind,
env_region: ty::Region<'tcx>,
) -> Result<(ty::Binder<'tcx, (Ty<'tcx>, Ty<'tcx>, Ty<'tcx>)>, Vec<ty::Predicate<'tcx>>), NoSolution>
{
) -> Result<
(ty::Binder<'tcx, (Ty<'tcx>, Ty<'tcx>, Ty<'tcx>)>, Option<ty::Predicate<'tcx>>),
NoSolution,
> {
match *self_ty.kind() {
ty::CoroutineClosure(def_id, args) => {
let args = args.as_coroutine_closure();
@ -339,7 +344,7 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_async_callable<'tc
);
(sig.tupled_inputs_ty, sig.return_ty, coroutine_ty)
}),
vec![],
None,
))
} else {
let async_fn_kind_trait_def_id =
@ -350,6 +355,13 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_async_callable<'tc
.next()
.unwrap()
.def_id;
// When we don't know the closure kind (and therefore also the closure's upvars,
// which are computed at the same time), we must delay the computation of the
// generator's upvars. We do this using the `AsyncFnKindHelper`, which as a trait
// goal functions similarly to the old `ClosureKind` predicate, and ensures that
// the goal kind <= the closure kind. As a projection `AsyncFnKindHelper::Upvars`
// will project to the right upvars for the generator, appending the inputs and
// coroutine upvars respecting the closure kind.
Ok((
args.coroutine_closure_sig().map_bound(|sig| {
let tupled_upvars_ty = Ty::new_projection(
@ -373,14 +385,14 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_async_callable<'tc
);
(sig.tupled_inputs_ty, sig.return_ty, coroutine_ty)
}),
vec![
Some(
ty::TraitRef::new(
tcx,
async_fn_kind_trait_def_id,
[kind_ty, Ty::from_closure_kind(tcx, goal_kind)],
)
.to_predicate(tcx),
],
),
))
}
}

View File

@ -2491,6 +2491,9 @@ fn confirm_async_closure_candidate<'cx, 'tcx>(
// since all this does is make the solver do more work.
//
// The code duplication due to the different length args is kind of weird, too.
//
// See the logic in `structural_traits` in the new solver to understand a bit
// more clearly how this *should* look.
let poly_cache_entry = args.coroutine_closure_sig().map_bound(|sig| {
let (projection_ty, term) = match tcx.item_name(obligation.predicate.def_id) {
sym::CallOnceFuture => {

View File

@ -121,9 +121,11 @@ pub(super) fn assemble_candidates<'o>(
self.assemble_async_fn_kind_helper_candidates(obligation, &mut candidates);
}
// FIXME: Put these into `else if` blocks above, since they're built-in.
self.assemble_closure_candidates(obligation, &mut candidates);
self.assemble_async_closure_candidates(obligation, &mut candidates);
self.assemble_fn_pointer_candidates(obligation, &mut candidates);
self.assemble_candidates_from_impls(obligation, &mut candidates);
self.assemble_candidates_from_object_ty(obligation, &mut candidates);
}
@ -382,6 +384,8 @@ fn assemble_async_fn_kind_helper_candidates(
return;
}
// Check that the self kind extends the goal kind. If it does,
// then there's nothing else to check.
if let Some(closure_kind) = self_ty.to_opt_closure_kind()
&& let Some(goal_kind) = target_kind_ty.to_opt_closure_kind()
{

View File

@ -88,6 +88,8 @@ pub(super) fn confirm_candidate(
ImplSource::Builtin(BuiltinImplSource::Misc, vtable_closure)
}
// No nested obligations or confirmation process. The checks that we do in
// candidate assembly are sufficient.
AsyncFnKindHelperCandidate => ImplSource::Builtin(BuiltinImplSource::Misc, vec![]),
CoroutineCandidate => {

View File

@ -728,7 +728,7 @@ fn compute(&mut self, arg: GenericArg<'tcx>) {
}
ty::CoroutineClosure(did, args) => {
// See the above comments.
// See the above comments. The same apply to coroutine-closures.
walker.skip_current_subtree();
self.compute(args.as_coroutine_closure().tupled_upvars_ty().into());
let obligations = self.nominal_obligations(did, args);

View File

@ -112,6 +112,10 @@ fn fn_sig_for_fn_abi<'tcx>(
};
let env_region = ty::Region::new_bound(tcx, ty::INNERMOST, br);
// When this `CoroutineClosure` comes from a `ConstructCoroutineInClosureShim`,
// make sure we respect the `target_kind` in that shim.
// FIXME(async_closures): This shouldn't be needed, and we should be populating
// a separate def-id for these bodies.
let mut kind = args.as_coroutine_closure().kind();
if let InstanceDef::ConstructCoroutineInClosureShim { target_kind, .. } = instance.def {
kind = target_kind;
@ -141,6 +145,8 @@ fn fn_sig_for_fn_abi<'tcx>(
)
}
ty::Coroutine(did, args) => {
// FIXME(async_closures): This isn't right for `CoroutineByMoveShim`.
let coroutine_kind = tcx.coroutine_kind(did).unwrap();
let sig = args.as_coroutine().sig();

View File

@ -287,6 +287,13 @@ fn resolve_associated_item<'tcx>(
{
match *rcvr_args.type_at(0).kind() {
ty::CoroutineClosure(coroutine_closure_def_id, args) => {
// If we're computing `AsyncFnOnce`/`AsyncFnMut` for a by-ref closure,
// or `AsyncFnOnce` for a by-mut closure, then construct a new body that
// has the right return types.
//
// Specifically, `AsyncFnMut` for a by-ref coroutine-closure just needs
// to have its input and output types fixed (`&mut self` and returning
// `i16` coroutine kind).
if target_kind > args.as_coroutine_closure().kind() {
Some(Instance {
def: ty::InstanceDef::ConstructCoroutineInClosureShim {