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.

5 files changed, 127 insertions, 1 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);
diff --git a/compiler/rustc_ty_utils/src/instance.rs b/compiler/rustc_ty_utils/src/instance.rs
index 2d76cf994e4..9faad10dd14 100644
--- a/compiler/rustc_ty_utils/src/instance.rs
+++ b/compiler/rustc_ty_utils/src/instance.rs
@@ -38,6 +38,7 @@ fn resolve_instance<'tcx>(
                 debug!(" => nontrivial drop glue");
                 match *ty.kind() {
                     ty::Closure(..)
+                    | ty::CoroutineClosure(..)
                     | ty::Coroutine(..)
                     | ty::Tuple(..)
                     | ty::Adt(..)
@@ -215,6 +216,7 @@ fn resolve_associated_item<'tcx>(
                         ty::Coroutine(..)
                         | ty::CoroutineWitness(..)
                         | ty::Closure(..)
+                        | ty::CoroutineClosure(..)
                         | ty::Tuple(..) => {}
                         _ => return Ok(None),
                     };
@@ -281,6 +283,34 @@ fn resolve_associated_item<'tcx>(
                         tcx.item_name(trait_item_id)
                     ),
                 }
+            } else if let Some(target_kind) = tcx.async_fn_trait_kind_from_def_id(trait_ref.def_id)
+            {
+                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 {
+                                    coroutine_closure_def_id,
+                                    target_kind,
+                                },
+                                args,
+                            })
+                        } else {
+                            Some(Instance::new(coroutine_closure_def_id, args))
+                        }
+                    }
+                    _ => bug!(
+                        "no built-in definition for `{trait_ref}::{}` for non-lending-closure type",
+                        tcx.item_name(trait_item_id)
+                    ),
+                }
             } else {
                 Instance::try_resolve_item_for_coroutine(tcx, trait_item_id, trait_id, rcvr_args)
             }
diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs
index 2fc4bfd4aa3..f20ded355b1 100644
--- a/compiler/rustc_ty_utils/src/layout.rs
+++ b/compiler/rustc_ty_utils/src/layout.rs
@@ -328,6 +328,17 @@ fn layout_of_uncached<'tcx>(
             )?
         }
 
+        ty::CoroutineClosure(_, args) => {
+            let tys = args.as_coroutine_closure().upvar_tys();
+            univariant(
+                &tys.iter()
+                    .map(|ty| Ok(cx.layout_of(ty)?.layout))
+                    .try_collect::<IndexVec<_, _>>()?,
+                &ReprOptions::default(),
+                StructKind::AlwaysSized,
+            )?
+        }
+
         ty::Tuple(tys) => {
             let kind =
                 if tys.len() == 0 { StructKind::AlwaysSized } else { StructKind::MaybeUnsized };
diff --git a/compiler/rustc_ty_utils/src/needs_drop.rs b/compiler/rustc_ty_utils/src/needs_drop.rs
index 08e5476ae43..7b3d2ab22cf 100644
--- a/compiler/rustc_ty_utils/src/needs_drop.rs
+++ b/compiler/rustc_ty_utils/src/needs_drop.rs
@@ -172,6 +172,12 @@ where
                         }
                     }
 
+                    ty::CoroutineClosure(_, args) => {
+                        for upvar in args.as_coroutine_closure().upvar_tys() {
+                            queue_type(self, upvar);
+                        }
+                    }
+
                     // Check for a `Drop` impl and whether this is a union or
                     // `ManuallyDrop`. If it's a struct or enum without a `Drop`
                     // impl then check whether the field types need `Drop`.
diff --git a/compiler/rustc_ty_utils/src/ty.rs b/compiler/rustc_ty_utils/src/ty.rs
index 2158aacab03..60b1bbe8c2a 100644
--- a/compiler/rustc_ty_utils/src/ty.rs
+++ b/compiler/rustc_ty_utils/src/ty.rs
@@ -18,7 +18,9 @@ fn sized_constraint_for_ty<'tcx>(
 
     let result = match ty.kind() {
         Bool | Char | Int(..) | Uint(..) | Float(..) | RawPtr(..) | Ref(..) | FnDef(..)
-        | FnPtr(_) | Array(..) | Closure(..) | Coroutine(..) | Never => vec![],
+        | FnPtr(_) | Array(..) | Closure(..) | CoroutineClosure(..) | Coroutine(..) | Never => {
+            vec![]
+        }
 
         Str | Dynamic(..) | Slice(_) | Foreign(..) | Error(_) | CoroutineWitness(..) => {
             // these are never sized - return the target type