Try   HackMD

Async Closures (and "coroutine-closures" in general)

For the purposes of keeping the implementation mostly future-compatible (i.e. with gen || {} and async gen || {}), most of this document calls async closures "coroutine-closures". Coroutine-closures are a generalization of async closures, being special syntax for closure expressions which return a coroutine, notably one that is allowed to capture from the closure's upvars.

For now, the only usable kind of coroutine-closure is the async closure, and supporting async closures is the extent of this PR. We may eventually support gen ||, etc., and all of the problems and curiosities described in this document apply to all coroutine-closures in general.

TyKind::CoroutineClosure

The main thing that this PR introduces is a new TyKind called CoroutineClosure and corresponding variants on other relevant enums in typeck and borrowck (UpvarArgs, DefiningTy, AggregateKind).

Signature

A traditional closure has a fn_sig_as_fn_ptr_ty which it uses to represent the signature of the closure. The problem with this sig type is that it doesn't actually reference the closure input: e.g., for a closure like || -> i32 { 0 }, the ptr type is fn(()) -> i32.

This is the first problem with a coroutine-closure, which returns a coroutine that is allowed to borrow from the closure's upvars, since there's no way to link the input lifetime (of the closure borrow) with the output lifetimes (in the coroutine's upvars).

The second problem is that coroutine-closures actually have several different signatures depending on if they're called with AsyncFn/AsyncFnMut/AsyncFnOnce.

For a general coroutine-closure which returns a coroutine that borrows from the closure's upvars

let s = String::new();
let c = async move || {
  call_service(&s).await;
};

the output coroutine returned by AsyncFn::call(& /* borrow for '1 */ c) captures the &'1 String for the input lifetime '1. Conversely, the coroutine returned by AsyncFnOnce::call_once(c) cannot borrow from the c closure since it is consumed as part of the call, so it must capture String by move, since the coroutine is now responsible for dropping the coroutine-closure's upvars after the coroutine-closure is dropped.

Conceptually, the coroutine-closure may be thought as containing several different signature types depending on whether it is being called by-ref or by-move.

Instead of doing this, we store the common parts of the different coroutines in the CoroutineClosureSignature, and compute the relevant coroutine output type on demand. This CoroutineClosureSignature is stored in a compressed form in the signature_parts_ty. See the docs on that type for more explanation.

Delaying the computation of the returned coroutine's upvars

We introduce a new AsyncFnKindHelper trait to enforce that the ClosureKind of a goal is within the capabilities of a CoroutineClosure, and which allows us to delay the projection of the tupled upvar types until after upvar analysis is complete.

This is because the upvars of the coroutine returned by the coroutine-closure should be the appended tuple of the input tys and the coroutine-closure's upvars. However, since the coroutine-closure's tupled upvars ty is an infer var until after closure analysis, we can't compute this eagerly.

We have two options therefore:

  1. We could mark all AsyncFn* goals as ambiguous until upvar analysis. However, this is really detrimental to inference in the program, since it means that programs like this would not type check:
let c = async || -> String { .. };
let s = c().await;
// ^ If we can't project `<{ c } as AsyncFn>::call()` to a coroutine type, then the `IntoFuture::into_future` call inside of the `.await` stalls out, and the type of `s` is left as an infer var.
s.as_bytes();
// ^ That means we can't call any methods on it!
  1. So instead, we can use an alias type (in this case, a projection: AsyncFnKindHelper::Upvars<'env, ...>) to delay the computation of the tupled upvars and give us something to put in its place.

Modifications to capture mode

Async closures are peculiar since they must move all the closure's arguments into the returned coroutine, but prefer not to move the coroutine-closure's upvars unless needed (otherwise they'd always be forced to only implement AsyncFnOnce) and instead capture them by ref.

Right now, the deusgaring that is shared between async closures and async functions (#119978) always generates a by-move async block.

In order to support this, I've modified the desugaring of these generated async blocks so that they always capture by-ref, and then modified the upvar analysis in hir_typeck to additionally force any argument types from the parent signature to be captured by move. This seems to work quite successfully.

NOTE: Since this is essentially a second copy of the coroutine body stored within the first, we must make sure to apply all the same MIR passes to this one. This functionality is implemented in run_passes, but other ad-hoc calls to visit_body will need to be audited for correctness.

Coroutine kind ty

The coroutines returned by AsyncFnOnce/AsyncFnMut/AsyncFn have the same def id, since they originate from the same HIR, but correspond to different bodies during codegen. To distinguish which body to associate with each implementation, I added a kind_ty to the coroutine args.

For coroutines that do not originate from coroutine-closures, this kind_ty is always (). For coroutines that do, this kind ty will match the ClosureTy of the call trait that produced it.

By move shims

This PR introduces the ByMoveBody MIR pass which is run right after MIR is built. When it finds the body of a coroutine from a coroutine-closure, and that coroutine-closure's closure kind is greater than FnOnce, it clones the body and adjusts all of the upvars to be taken by-move. We call this the by_move_body, and store it into the CoroutineInfo in the original coroutine's MIR body.

Later on, when Instance::resolve tries to resolve Future::poll_next for a coroutine type that is returned by AsyncFnOnce::call_once (and that coroutine-closure's closure kind is greater than FnOnce), we can use this by-move body instead.

We also generate a shim for AsyncFnOnce::call_once for these coroutine-closures, which constructs a coroutine by moving the coroutine-closure's upvars rather than borrows them.

FOLLOW-UP: The fn_sig_for_fn_abi/Instance::ty implementation for these shims is a bit sketchy. This shouldn't cause issues for codegen_llvm, but may cause issues for stricter backends like clif.

Wins

Higher-ranked async closures

We support coroutine-closures with binders in their signature, both implicit and explicit.

let c = async |s: &str| { do_service(s).await; };

While the the future coroutine returned by the closure may reference late-bound lifetimes, the coroutine still is not "lending". See async-await/async-closures/not-lending.rs for an example.

It is however not currently possible for the return type of the coroutine to reference the higher-ranked lifetimes of the closure:

let c = async |s: &str| -> &str { s };

FOLLOW-UP: Figure out why this code doesn't work.

Limitations

The "double move" case

The coroutine returned by coroutine-closures will always opportunistically borrow from the parent coroutine. There's essentially no way to express move || async move { .. }.

Async closures don't currently implement the regular Fn traits

The AsyncFn hierarchy of traits is not currently unified with the Fn hierarchy of traits.

In the future, we could make coroutine-closures implement FnOnce always (since it's always possible to implement FnOnce) and then opportunistically implement FnMut/Fn as long as the don't borrow anything from the closure upvars.

EDIT: this was fixed https://github.com/rust-lang/rust/pull/120712

Closure signature inference isn't implemented

This PR does not implement closure signature inference that comes from passing async closures as arguments. This could be implemented if needed, but it may put the new trait solver in a worse position w.r.t. its inability to do closure signature inference.