Merge pull request rust-lang#20895 from ChayimFriedman2/async-block-no-opaque

fix: Lower async block/closures correctly

Author: ShoyuVanilla

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;
@@ -48,7 +49,7 @@ use triomphe::Arc;
 
 use crate::{
     CallableDefId, FnAbi, ImplTraitId, MemoryMap, TraitEnvironment, consteval,
-    db::{HirDatabase, InternedClosure},
+    db::{HirDatabase, InternedClosure, InternedCoroutine},
     generics::generics,
     layout::Layout,
     mir::pad16,
@@ -1389,33 +1390,6 @@ impl<'db> HirDisplay<'db> for Ty<'db> {
                             SizedByDefault::Sized { anchor: krate },
                         )?;
                     }
-                    ImplTraitId::AsyncBlockTypeImplTrait(body, ..) => {
-                        let future_trait =
-                            LangItem::Future.resolve_trait(db, body.module(db).krate());
-                        let output = future_trait.and_then(|t| {
-                            t.trait_items(db)
-                                .associated_type_by_name(&Name::new_symbol_root(sym::Output))
-                        });
-                        write!(f, "impl ")?;
-                        if let Some(t) = future_trait {
-                            f.start_location_link(t.into());
-                        }
-                        write!(f, "Future")?;
-                        if future_trait.is_some() {
-                            f.end_location_link();
-                        }
-                        write!(f, "<")?;
-                        if let Some(t) = output {
-                            f.start_location_link(t.into());
-                        }
-                        write!(f, "Output")?;
-                        if output.is_some() {
-                            f.end_location_link();
-                        }
-                        write!(f, " = ")?;
-                        alias_ty.args.type_at(0).hir_fmt(f)?;
-                        write!(f, ">")?;
-                    }
                 }
             }
             TyKind::Closure(id, substs) => {
@@ -1471,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
@@ -1567,28 +1610,69 @@ impl<'db> HirDisplay<'db> for Ty<'db> {
                 }
             }
             TyKind::Infer(..) => write!(f, "_")?,
-            TyKind::Coroutine(_, subst) => {
-                if f.display_kind.is_source_code() {
-                    return Err(HirDisplayError::DisplaySourceCodeError(
-                        DisplaySourceCodeError::Coroutine,
-                    ));
-                }
+            TyKind::Coroutine(coroutine_id, subst) => {
+                let InternedCoroutine(owner, expr_id) = coroutine_id.0.loc(db);
                 let CoroutineArgsParts { resume_ty, yield_ty, return_ty, .. } =
                     subst.split_coroutine_args();
-                write!(f, "|")?;
-                resume_ty.hir_fmt(f)?;
-                write!(f, "|")?;
+                let body = db.body(owner);
+                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| {
+                            t.trait_items(db)
+                                .associated_type_by_name(&Name::new_symbol_root(sym::Output))
+                        });
+                        write!(f, "impl ")?;
+                        if let Some(t) = future_trait {
+                            f.start_location_link(t.into());
+                        }
+                        write!(f, "Future")?;
+                        if future_trait.is_some() {
+                            f.end_location_link();
+                        }
+                        write!(f, "<")?;
+                        if let Some(t) = output {
+                            f.start_location_link(t.into());
+                        }
+                        write!(f, "Output")?;
+                        if output.is_some() {
+                            f.end_location_link();
+                        }
+                        write!(f, " = ")?;
+                        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,
+                            ));
+                        }
+                        write!(f, "|")?;
+                        resume_ty.hir_fmt(f)?;
+                        write!(f, "|")?;
 
-                write!(f, " yields ")?;
-                yield_ty.hir_fmt(f)?;
+                        write!(f, " yields ")?;
+                        yield_ty.hir_fmt(f)?;
 
-                write!(f, " -> ")?;
-                return_ty.hir_fmt(f)?;
+                        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.rs

@@ -1243,7 +1243,6 @@ impl<'body, 'db> InferenceContext<'body, 'db> {
                     }
                     (self.db.type_alias_impl_traits(def), idx)
                 }
-                _ => unreachable!(),
             };
             let Some(impl_traits) = impl_traits else {
                 return ty;

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