Auto merge of #120361 - compiler-errors:async-closures, r=oli-obk

Rework support for async closures; allow them to return futures that borrow from the closure's captures

This PR implements a new lowering for async closures via `TyKind::CoroutineClosure` which handles the curious relationship between the closure and the coroutine that it returns.

I wrote up a bunch in [this hackmd](https://hackmd.io/`@compiler-errors/S1HvqQxca)` which will be copied to the dev guide after this PR lands, and hopefully left sufficient comments in the source code explaining why this change is as large as it is.

This also necessitates that they begin implementing the `AsyncFn`-family of traits, rather than the `Fn`-family of traits -- if you need `Fn` implementations, you should probably use the non-sugar `|| async {}` syntax instead.

Notably this PR does not yet implement `async Fn()` syntax sugar for bounds, but I expect to add those soon (**edit:** #120392). For now, users must use `AsyncFn()` traits directly, which necessitates adding the `async_fn_traits` feature gate as well. I will add this as a follow-up very soon.

r? oli-obk

This is based on top of #120322, but that PR is minimal.

1 files changed, 77 insertions, 0 deletions

diff --git a/compiler/rustc_ty_utils/src/abi.rs b/compiler/rustc_ty_utils/src/abi.rs
index af26616831f..0542ef4ecf9 100644
--- a/compiler/rustc_ty_utils/src/abi.rs
+++ b/compiler/rustc_ty_utils/src/abi.rs
@@ -101,6 +101,49 @@ fn fn_sig_for_fn_abi<'tcx>(
                 bound_vars,
             )
         }
+        ty::CoroutineClosure(def_id, args) => {
+            let sig = args.as_coroutine_closure().coroutine_closure_sig();
+            let bound_vars = tcx.mk_bound_variable_kinds_from_iter(
+                sig.bound_vars().iter().chain(iter::once(ty::BoundVariableKind::Region(ty::BrEnv))),
+            );
+            let br = ty::BoundRegion {
+                var: ty::BoundVar::from_usize(bound_vars.len() - 1),
+                kind: ty::BoundRegionKind::BrEnv,
+            };
+            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;
+            }
+
+            let env_ty =
+                tcx.closure_env_ty(Ty::new_coroutine_closure(tcx, def_id, args), kind, env_region);
+
+            let sig = sig.skip_binder();
+            ty::Binder::bind_with_vars(
+                tcx.mk_fn_sig(
+                    iter::once(env_ty).chain([sig.tupled_inputs_ty]),
+                    sig.to_coroutine_given_kind_and_upvars(
+                        tcx,
+                        args.as_coroutine_closure().parent_args(),
+                        tcx.coroutine_for_closure(def_id),
+                        kind,
+                        env_region,
+                        args.as_coroutine_closure().tupled_upvars_ty(),
+                        args.as_coroutine_closure().coroutine_captures_by_ref_ty(),
+                    ),
+                    sig.c_variadic,
+                    sig.unsafety,
+                    sig.abi,
+                ),
+                bound_vars,
+            )
+        }
         ty::Coroutine(did, args) => {
             let coroutine_kind = tcx.coroutine_kind(did).unwrap();
             let sig = args.as_coroutine().sig();
@@ -112,6 +155,40 @@ fn fn_sig_for_fn_abi<'tcx>(
                 var: ty::BoundVar::from_usize(bound_vars.len() - 1),
                 kind: ty::BoundRegionKind::BrEnv,
             };
+
+            let mut ty = ty;
+            // When this `Closure` comes from a `CoroutineKindShim`,
+            // 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.
+            if let InstanceDef::CoroutineKindShim { target_kind, .. } = instance.def {
+                // Grab the parent coroutine-closure. It has the same args for the purposes
+                // of substitution, so this will be okay to do.
+                let ty::CoroutineClosure(_, coroutine_closure_args) = *tcx
+                    .instantiate_and_normalize_erasing_regions(
+                        args,
+                        param_env,
+                        tcx.type_of(tcx.parent(did)),
+                    )
+                    .kind()
+                else {
+                    bug!("CoroutineKindShim comes from calling a coroutine-closure");
+                };
+                let coroutine_closure_args = coroutine_closure_args.as_coroutine_closure();
+                ty = tcx.instantiate_bound_regions_with_erased(
+                    coroutine_closure_args.coroutine_closure_sig().map_bound(|sig| {
+                        sig.to_coroutine_given_kind_and_upvars(
+                            tcx,
+                            coroutine_closure_args.parent_args(),
+                            did,
+                            target_kind,
+                            tcx.lifetimes.re_erased,
+                            coroutine_closure_args.tupled_upvars_ty(),
+                            coroutine_closure_args.coroutine_captures_by_ref_ty(),
+                        )
+                    }),
+                );
+            }
             let env_ty = Ty::new_mut_ref(tcx, ty::Region::new_bound(tcx, ty::INNERMOST, br), ty);
 
             let pin_did = tcx.require_lang_item(LangItem::Pin, None);