Lower async closures to TyKind::CoroutineClosure

Instead of `TyKind::Closure`.

Note: the same `InternedCoroutineId` is used both for the *async closure* as well as for the *async block it returns*. When used in `TyKind::CoroutineClosure`, it represents the closure. When used in `TyKind::Coroutine`, it represents the async block. The generic args are different, though.

Also noteworthy is that we distinguish between the different kinds of coroutines (general coroutines, async coroutines and eventually gen coroutines too) via the expression producing them (stored in the `InternedCoroutineId`). It might be worth it to introduce a `CoroutineKind` field to `InternedCoroutineId`, although this is not done in this PR.

Author: ChayimFriedman2

src/tools/rust-analyzer/crates/hir-ty/src/display.rs

@@ -38,7 +38,8 @@ use rustc_apfloat::{
 use rustc_ast_ir::FloatTy;
 use rustc_hash::FxHashSet;
 use rustc_type_ir::{
-    AliasTyKind, BoundVarIndexKind, CoroutineArgsParts, RegionKind, Upcast,
+    AliasTyKind, BoundVarIndexKind, CoroutineArgsParts, CoroutineClosureArgsParts, RegionKind,
+    Upcast,
     inherent::{AdtDef, GenericArgs as _, IntoKind, SliceLike, Term as _, Ty as _, Tys as _},
 };
 use smallvec::SmallVec;
@@ -1444,14 +1445,83 @@ impl<'db> HirDisplay<'db> for Ty<'db> {
                     }
                     if f.closure_style == ClosureStyle::RANotation || !sig.output().is_unit() {
                         write!(f, " -> ")?;
-                        // FIXME: We display `AsyncFn` as `-> impl Future`, but this is hard to fix because
-                        // we don't have a trait environment here, required to normalize `<Ret as Future>::Output`.
                         sig.output().hir_fmt(f)?;
                     }
                 } else {
                     write!(f, "{{closure}}")?;
                 }
             }
+            TyKind::CoroutineClosure(id, args) => {
+                let id = id.0;
+                if f.display_kind.is_source_code() {
+                    if !f.display_kind.allows_opaque() {
+                        return Err(HirDisplayError::DisplaySourceCodeError(
+                            DisplaySourceCodeError::OpaqueType,
+                        ));
+                    } else if f.closure_style != ClosureStyle::ImplFn {
+                        never!("Only `impl Fn` is valid for displaying closures in source code");
+                    }
+                }
+                match f.closure_style {
+                    ClosureStyle::Hide => return write!(f, "{TYPE_HINT_TRUNCATION}"),
+                    ClosureStyle::ClosureWithId => {
+                        return write!(
+                            f,
+                            "{{async closure#{:?}}}",
+                            salsa::plumbing::AsId::as_id(&id).index()
+                        );
+                    }
+                    ClosureStyle::ClosureWithSubst => {
+                        write!(
+                            f,
+                            "{{async closure#{:?}}}",
+                            salsa::plumbing::AsId::as_id(&id).index()
+                        )?;
+                        return hir_fmt_generics(f, args.as_slice(), None, None);
+                    }
+                    _ => (),
+                }
+                let CoroutineClosureArgsParts { closure_kind_ty, signature_parts_ty, .. } =
+                    args.split_coroutine_closure_args();
+                let kind = closure_kind_ty.to_opt_closure_kind().unwrap();
+                let kind = match kind {
+                    rustc_type_ir::ClosureKind::Fn => "AsyncFn",
+                    rustc_type_ir::ClosureKind::FnMut => "AsyncFnMut",
+                    rustc_type_ir::ClosureKind::FnOnce => "AsyncFnOnce",
+                };
+                let TyKind::FnPtr(coroutine_sig, _) = signature_parts_ty.kind() else {
+                    unreachable!("invalid coroutine closure signature");
+                };
+                let coroutine_sig = coroutine_sig.skip_binder();
+                let coroutine_inputs = coroutine_sig.inputs();
+                let TyKind::Tuple(coroutine_inputs) = coroutine_inputs.as_slice()[1].kind() else {
+                    unreachable!("invalid coroutine closure signature");
+                };
+                let TyKind::Tuple(coroutine_output) = coroutine_sig.output().kind() else {
+                    unreachable!("invalid coroutine closure signature");
+                };
+                let coroutine_output = coroutine_output.as_slice()[1];
+                match f.closure_style {
+                    ClosureStyle::ImplFn => write!(f, "impl {kind}(")?,
+                    ClosureStyle::RANotation => write!(f, "async |")?,
+                    _ => unreachable!(),
+                }
+                if coroutine_inputs.is_empty() {
+                } else if f.should_truncate() {
+                    write!(f, "{TYPE_HINT_TRUNCATION}")?;
+                } else {
+                    f.write_joined(coroutine_inputs, ", ")?;
+                };
+                match f.closure_style {
+                    ClosureStyle::ImplFn => write!(f, ")")?,
+                    ClosureStyle::RANotation => write!(f, "|")?,
+                    _ => unreachable!(),
+                }
+                if f.closure_style == ClosureStyle::RANotation || !coroutine_output.is_unit() {
+                    write!(f, " -> ")?;
+                    coroutine_output.hir_fmt(f)?;
+                }
+            }
             TyKind::Placeholder(_) => write!(f, "{{placeholder}}")?,
             TyKind::Param(param) => {
                 // FIXME: We should not access `param.id`, it should be removed, and we should know the
@@ -1545,8 +1615,13 @@ impl<'db> HirDisplay<'db> for Ty<'db> {
                 let CoroutineArgsParts { resume_ty, yield_ty, return_ty, .. } =
                     subst.split_coroutine_args();
                 let body = db.body(owner);
-                match &body[expr_id] {
-                    hir_def::hir::Expr::Async { .. } => {
+                let expr = &body[expr_id];
+                match expr {
+                    hir_def::hir::Expr::Closure {
+                        closure_kind: hir_def::hir::ClosureKind::Async,
+                        ..
+                    }
+                    | hir_def::hir::Expr::Async { .. } => {
                         let future_trait =
                             LangItem::Future.resolve_trait(db, owner.module(db).krate());
                         let output = future_trait.and_then(|t| {
@@ -1573,7 +1648,10 @@ impl<'db> HirDisplay<'db> for Ty<'db> {
                         return_ty.hir_fmt(f)?;
                         write!(f, ">")?;
                     }
-                    _ => {
+                    hir_def::hir::Expr::Closure {
+                        closure_kind: hir_def::hir::ClosureKind::Coroutine(..),
+                        ..
+                    } => {
                         if f.display_kind.is_source_code() {
                             return Err(HirDisplayError::DisplaySourceCodeError(
                                 DisplaySourceCodeError::Coroutine,
@@ -1589,12 +1667,12 @@ impl<'db> HirDisplay<'db> for Ty<'db> {
                         write!(f, " -> ")?;
                         return_ty.hir_fmt(f)?;
                     }
+                    _ => panic!("invalid expr for coroutine: {expr:?}"),
                 }
             }
             TyKind::CoroutineWitness(..) => write!(f, "{{coroutine witness}}")?,
             TyKind::Pat(_, _) => write!(f, "{{pat}}")?,
             TyKind::UnsafeBinder(_) => write!(f, "{{unsafe binder}}")?,
-            TyKind::CoroutineClosure(_, _) => write!(f, "{{coroutine closure}}")?,
             TyKind::Alias(_, _) => write!(f, "{{alias}}")?,
         }
         Ok(())

src/tools/rust-analyzer/crates/hir-ty/src/infer/closure.rs

@@ -11,8 +11,9 @@ use hir_def::{
     type_ref::TypeRefId,
 };
 use rustc_type_ir::{
-    ClosureArgs, ClosureArgsParts, CoroutineArgs, CoroutineArgsParts, Interner, TypeSuperVisitable,
-    TypeVisitable, TypeVisitableExt, TypeVisitor,
+    ClosureArgs, ClosureArgsParts, CoroutineArgs, CoroutineArgsParts, CoroutineClosureArgs,
+    CoroutineClosureArgsParts, Interner, TypeSuperVisitable, TypeVisitable, TypeVisitableExt,
+    TypeVisitor,
     inherent::{BoundExistentialPredicates, GenericArgs as _, IntoKind, SliceLike, Ty as _},
 };
 use tracing::debug;
@@ -22,8 +23,9 @@ use crate::{
     db::{InternedClosure, InternedCoroutine},
     infer::{BreakableKind, Diverges, coerce::CoerceMany},
     next_solver::{
-        AliasTy, Binder, ClauseKind, DbInterner, ErrorGuaranteed, FnSig, GenericArgs, PolyFnSig,
-        PolyProjectionPredicate, Predicate, PredicateKind, SolverDefId, Ty, TyKind,
+        AliasTy, Binder, BoundRegionKind, BoundVarKind, BoundVarKinds, ClauseKind, DbInterner,
+        ErrorGuaranteed, FnSig, GenericArgs, PolyFnSig, PolyProjectionPredicate, Predicate,
+        PredicateKind, SolverDefId, Ty, TyKind,
         abi::Safety,
         infer::{
             BoundRegionConversionTime, InferOk, InferResult,
@@ -72,6 +74,8 @@ impl<'db> InferenceContext<'_, 'db> {
         let sig_ty = Ty::new_fn_ptr(interner, bound_sig);
 
         let parent_args = GenericArgs::identity_for_item(interner, self.generic_def.into());
+        // FIXME: Make this an infer var and infer it later.
+        let tupled_upvars_ty = self.types.unit;
         let (id, ty, resume_yield_tys) = match closure_kind {
             ClosureKind::Coroutine(_) => {
                 let yield_ty = self.table.next_ty_var();
@@ -80,11 +84,11 @@ impl<'db> InferenceContext<'_, 'db> {
                 // FIXME: Infer the upvars later.
                 let parts = CoroutineArgsParts {
                     parent_args,
-                    kind_ty: Ty::new_unit(interner),
+                    kind_ty: self.types.unit,
                     resume_ty,
                     yield_ty,
                     return_ty: body_ret_ty,
-                    tupled_upvars_ty: Ty::new_unit(interner),
+                    tupled_upvars_ty,
                 };
 
                 let coroutine_id =
@@ -97,9 +101,7 @@ impl<'db> InferenceContext<'_, 'db> {
 
                 (None, coroutine_ty, Some((resume_ty, yield_ty)))
             }
-            // FIXME(next-solver): `ClosureKind::Async` should really be a separate arm that creates a `CoroutineClosure`.
-            // But for now we treat it as a closure.
-            ClosureKind::Closure | ClosureKind::Async => {
+            ClosureKind::Closure => {
                 let closure_id = self.db.intern_closure(InternedClosure(self.owner, tgt_expr));
                 match expected_kind {
                     Some(kind) => {
@@ -117,15 +119,15 @@ impl<'db> InferenceContext<'_, 'db> {
                     }
                     None => {}
                 };
-                // FIXME: Infer the kind and the upvars later when needed.
+                // FIXME: Infer the kind later if needed.
                 let parts = ClosureArgsParts {
                     parent_args,
                     closure_kind_ty: Ty::from_closure_kind(
                         interner,
                         expected_kind.unwrap_or(rustc_type_ir::ClosureKind::Fn),
                     ),
                     closure_sig_as_fn_ptr_ty: sig_ty,
-                    tupled_upvars_ty: Ty::new_unit(interner),
+                    tupled_upvars_ty,
                 };
                 let closure_ty = Ty::new_closure(
                     interner,
@@ -136,6 +138,61 @@ impl<'db> InferenceContext<'_, 'db> {
                 self.add_current_closure_dependency(closure_id);
                 (Some(closure_id), closure_ty, None)
             }
+            ClosureKind::Async => {
+                // async closures always return the type ascribed after the `->` (if present),
+                // and yield `()`.
+                let bound_return_ty = bound_sig.skip_binder().output();
+                let bound_yield_ty = self.types.unit;
+                // rustc uses a special lang item type for the resume ty. I don't believe this can cause us problems.
+                let resume_ty = self.types.unit;
+
+                // FIXME: Infer the kind later if needed.
+                let closure_kind_ty = Ty::from_closure_kind(
+                    interner,
+                    expected_kind.unwrap_or(rustc_type_ir::ClosureKind::Fn),
+                );
+
+                // FIXME: Infer captures later.
+                // `for<'env> fn() -> ()`, for no captures.
+                let coroutine_captures_by_ref_ty = Ty::new_fn_ptr(
+                    interner,
+                    Binder::bind_with_vars(
+                        interner.mk_fn_sig([], self.types.unit, false, Safety::Safe, FnAbi::Rust),
+                        BoundVarKinds::new_from_iter(
+                            interner,
+                            [BoundVarKind::Region(BoundRegionKind::ClosureEnv)],
+                        ),
+                    ),
+                );
+                let closure_args = CoroutineClosureArgs::new(
+                    interner,
+                    CoroutineClosureArgsParts {
+                        parent_args,
+                        closure_kind_ty,
+                        signature_parts_ty: Ty::new_fn_ptr(
+                            interner,
+                            bound_sig.map_bound(|sig| {
+                                interner.mk_fn_sig(
+                                    [
+                                        resume_ty,
+                                        Ty::new_tup_from_iter(interner, sig.inputs().iter()),
+                                    ],
+                                    Ty::new_tup(interner, &[bound_yield_ty, bound_return_ty]),
+                                    sig.c_variadic,
+                                    sig.safety,
+                                    sig.abi,
+                                )
+                            }),
+                        ),
+                        tupled_upvars_ty,
+                        coroutine_captures_by_ref_ty,
+                    },
+                );
+
+                let coroutine_id =
+                    self.db.intern_coroutine(InternedCoroutine(self.owner, tgt_expr)).into();
+                (None, Ty::new_coroutine_closure(interner, coroutine_id, closure_args.args), None)
+            }
         };
 
         // Now go through the argument patterns

src/tools/rust-analyzer/crates/hir-ty/src/next_solver/interner.rs

@@ -17,8 +17,8 @@ use rustc_abi::{ReprFlags, ReprOptions};
 use rustc_hash::FxHashSet;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_type_ir::{
-    AliasTermKind, AliasTyKind, BoundVar, CollectAndApply, DebruijnIndex, EarlyBinder,
-    FlagComputation, Flags, GenericArgKind, ImplPolarity, InferTy, Interner, TraitRef,
+    AliasTermKind, AliasTyKind, BoundVar, CollectAndApply, CoroutineWitnessTypes, DebruijnIndex,
+    EarlyBinder, FlagComputation, Flags, GenericArgKind, ImplPolarity, InferTy, Interner, TraitRef,
     TypeVisitableExt, UniverseIndex, Upcast, Variance,
     elaborate::elaborate,
     error::TypeError,
@@ -29,7 +29,7 @@ use rustc_type_ir::{
 
 use crate::{
     FnAbi,
-    db::HirDatabase,
+    db::{HirDatabase, InternedCoroutine},
     method_resolution::{ALL_FLOAT_FPS, ALL_INT_FPS, TyFingerprint},
     next_solver::{
         AdtIdWrapper, BoundConst, CallableIdWrapper, CanonicalVarKind, ClosureIdWrapper,
@@ -1205,12 +1205,28 @@ impl<'db> Interner for DbInterner<'db> {
         self.db().callable_item_signature(def_id.0)
     }
 
-    fn coroutine_movability(self, _def_id: Self::CoroutineId) -> rustc_ast_ir::Movability {
-        unimplemented!()
+    fn coroutine_movability(self, def_id: Self::CoroutineId) -> rustc_ast_ir::Movability {
+        // FIXME: Make this a query? I don't believe this can be accessed from bodies other than
+        // the current infer query, except with revealed opaques - is it rare enough to not matter?
+        let InternedCoroutine(owner, expr_id) = def_id.0.loc(self.db);
+        let body = self.db.body(owner);
+        let expr = &body[expr_id];
+        match *expr {
+            hir_def::hir::Expr::Closure { closure_kind, .. } => match closure_kind {
+                hir_def::hir::ClosureKind::Coroutine(movability) => match movability {
+                    hir_def::hir::Movability::Static => rustc_ast_ir::Movability::Static,
+                    hir_def::hir::Movability::Movable => rustc_ast_ir::Movability::Movable,
+                },
+                hir_def::hir::ClosureKind::Async => rustc_ast_ir::Movability::Static,
+                _ => panic!("unexpected expression for a coroutine: {expr:?}"),
+            },
+            hir_def::hir::Expr::Async { .. } => rustc_ast_ir::Movability::Static,
+            _ => panic!("unexpected expression for a coroutine: {expr:?}"),
+        }
     }
 
-    fn coroutine_for_closure(self, _def_id: Self::CoroutineId) -> Self::CoroutineId {
-        unimplemented!()
+    fn coroutine_for_closure(self, def_id: Self::CoroutineClosureId) -> Self::CoroutineId {
+        def_id
     }
 
     fn generics_require_sized_self(self, def_id: Self::DefId) -> bool {
@@ -1725,23 +1741,39 @@ impl<'db> Interner for DbInterner<'db> {
         panic!("Bug encountered in next-trait-solver: {}", msg.to_string())
     }
 
-    fn is_general_coroutine(self, _coroutine_def_id: Self::CoroutineId) -> bool {
-        // FIXME(next-solver)
-        true
+    fn is_general_coroutine(self, def_id: Self::CoroutineId) -> bool {
+        // FIXME: Make this a query? I don't believe this can be accessed from bodies other than
+        // the current infer query, except with revealed opaques - is it rare enough to not matter?
+        let InternedCoroutine(owner, expr_id) = def_id.0.loc(self.db);
+        let body = self.db.body(owner);
+        matches!(
+            body[expr_id],
+            hir_def::hir::Expr::Closure {
+                closure_kind: hir_def::hir::ClosureKind::Coroutine(_),
+                ..
+            }
+        )
     }
 
-    fn coroutine_is_async(self, _coroutine_def_id: Self::CoroutineId) -> bool {
-        // FIXME(next-solver)
-        true
+    fn coroutine_is_async(self, def_id: Self::CoroutineId) -> bool {
+        // FIXME: Make this a query? I don't believe this can be accessed from bodies other than
+        // the current infer query, except with revealed opaques - is it rare enough to not matter?
+        let InternedCoroutine(owner, expr_id) = def_id.0.loc(self.db);
+        let body = self.db.body(owner);
+        matches!(
+            body[expr_id],
+            hir_def::hir::Expr::Closure { closure_kind: hir_def::hir::ClosureKind::Async, .. }
+                | hir_def::hir::Expr::Async { .. }
+        )
     }
 
     fn coroutine_is_gen(self, _coroutine_def_id: Self::CoroutineId) -> bool {
-        // FIXME(next-solver)
+        // We don't handle gen coroutines yet.
         false
     }
 
     fn coroutine_is_async_gen(self, _coroutine_def_id: Self::CoroutineId) -> bool {
-        // FIXME(next-solver)
+        // We don't handle gen coroutines yet.
         false
     }
 
@@ -1897,10 +1929,12 @@ impl<'db> Interner for DbInterner<'db> {
     fn coroutine_hidden_types(
         self,
         _def_id: Self::CoroutineId,
-    ) -> EarlyBinder<Self, rustc_type_ir::Binder<Self, rustc_type_ir::CoroutineWitnessTypes<Self>>>
-    {
-        // FIXME(next-solver)
-        unimplemented!()
+    ) -> EarlyBinder<Self, Binder<'db, CoroutineWitnessTypes<Self>>> {
+        // FIXME: Actually implement this.
+        EarlyBinder::bind(Binder::dummy(CoroutineWitnessTypes {
+            types: Tys::default(),
+            assumptions: RegionAssumptions::default(),
+        }))
     }
 
     fn is_default_trait(self, def_id: Self::TraitId) -> bool {