Merge #5971

5971: Implement async blocks r=flodiebold a=oxalica

Fix #4018

@flodiebold already gave a generic guide in the issue. Here's some concern about implementation detail:
- Chalk doesn't support generator type yet.
- Adding generator type as a brand new type (ctor) can be complex and need to *re-introduced* builtin impls. (Like how we implement closures before native closure support of chalk, which is already removed in #5401 )
- The output type of async block should be known after type inference of the whole body.
  - We cannot directly get the type from source like return-positon-impl-trait. But we still need to provide trait bounds when chalk asking for `opaque_ty_data`.
  - During the inference, the output type of async block can be temporary unknown and participate the later inference.
    `let a = async { None }; let _: i32 = a.await.unwrap();`

So in this PR, the type of async blocks is inferred as an opaque type parameterized by the `Future::Output` type it should be, like what we do with closure type.
And it really works now.

Well, I still have some questions:
- The bounds `AsyncBlockImplType<T>: Future<Output = T>` is currently generated in `opaque_ty_data`. I'm not sure if we should put this code here.
- Type of async block is now rendered as `impl Future<Output = OutputType>`. Do we need to special display to hint that it's a async block? Note that closure type has its special format, instead of `impl Fn(..) -> ..` or function type.



Co-authored-by: oxalica <[email protected]>

Author: bors[bot]

crates/hir/src/code_model.rs

@@ -1283,6 +1283,8 @@ impl Type {
     /// Checks that particular type `ty` implements `std::future::Future`.
     /// This function is used in `.await` syntax completion.
     pub fn impls_future(&self, db: &dyn HirDatabase) -> bool {
+        // No special case for the type of async block, since Chalk can figure it out.
+
         let krate = self.krate;
 
         let std_future_trait =
@@ -1600,6 +1602,11 @@ impl Type {
                                 cb(type_.derived(ty.clone()));
                             }
                         }
+                        TypeCtor::OpaqueType(..) => {
+                            if let Some(bounds) = ty.impl_trait_bounds(db) {
+                                walk_bounds(db, &type_.derived(ty.clone()), &bounds, cb);
+                            }
+                        }
                         _ => (),
                     }
 

crates/hir_def/src/body/lower.rs

@@ -239,7 +239,10 @@ impl ExprCollector<'_> {
                     None => self.missing_expr(),
                 },
                 // FIXME: we need to record these effects somewhere...
-                ast::Effect::Async(_) => self.collect_block_opt(e.block_expr()),
+                ast::Effect::Async(_) => {
+                    let body = self.collect_block_opt(e.block_expr());
+                    self.alloc_expr(Expr::Async { body }, syntax_ptr)
+                }
             },
             ast::Expr::BlockExpr(e) => self.collect_block(e),
             ast::Expr::LoopExpr(e) => {

crates/hir_def/src/expr.rs

@@ -111,6 +111,9 @@ pub enum Expr {
     TryBlock {
         body: ExprId,
     },
+    Async {
+        body: ExprId,
+    },
     Cast {
         expr: ExprId,
         type_ref: TypeRef,
@@ -250,7 +253,7 @@ impl Expr {
                     f(*expr);
                 }
             }
-            Expr::TryBlock { body } | Expr::Unsafe { body } => f(*body),
+            Expr::TryBlock { body } | Expr::Unsafe { body } | Expr::Async { body } => f(*body),
             Expr::Loop { body, .. } => f(*body),
             Expr::While { condition, body, .. } => {
                 f(*condition);

crates/hir_ty/src/display.rs

@@ -381,19 +381,24 @@ impl HirDisplay for ApplicationTy {
                 }
             }
             TypeCtor::OpaqueType(opaque_ty_id) => {
-                let bounds = match opaque_ty_id {
+                match opaque_ty_id {
                     OpaqueTyId::ReturnTypeImplTrait(func, idx) => {
                         let datas =
                             f.db.return_type_impl_traits(func).expect("impl trait id without data");
                         let data = (*datas)
                             .as_ref()
                             .map(|rpit| rpit.impl_traits[idx as usize].bounds.clone());
-                        data.subst(&self.parameters)
+                        let bounds = data.subst(&self.parameters);
+                        write!(f, "impl ")?;
+                        write_bounds_like_dyn_trait(&bounds.value, f)?;
+                        // FIXME: it would maybe be good to distinguish this from the alias type (when debug printing), and to show the substitution
                     }
-                };
-                write!(f, "impl ")?;
-                write_bounds_like_dyn_trait(&bounds.value, f)?;
-                // FIXME: it would maybe be good to distinguish this from the alias type (when debug printing), and to show the substitution
+                    OpaqueTyId::AsyncBlockTypeImplTrait(..) => {
+                        write!(f, "impl Future<Output = ")?;
+                        self.parameters[0].hir_fmt(f)?;
+                        write!(f, ">")?;
+                    }
+                }
             }
             TypeCtor::Closure { .. } => {
                 let sig = self.parameters[0].callable_sig(f.db);
@@ -474,18 +479,21 @@ impl HirDisplay for Ty {
                 write_bounds_like_dyn_trait(predicates, f)?;
             }
             Ty::Opaque(opaque_ty) => {
-                let bounds = match opaque_ty.opaque_ty_id {
+                match opaque_ty.opaque_ty_id {
                     OpaqueTyId::ReturnTypeImplTrait(func, idx) => {
                         let datas =
                             f.db.return_type_impl_traits(func).expect("impl trait id without data");
                         let data = (*datas)
                             .as_ref()
                             .map(|rpit| rpit.impl_traits[idx as usize].bounds.clone());
-                        data.subst(&opaque_ty.parameters)
+                        let bounds = data.subst(&opaque_ty.parameters);
+                        write!(f, "impl ")?;
+                        write_bounds_like_dyn_trait(&bounds.value, f)?;
+                    }
+                    OpaqueTyId::AsyncBlockTypeImplTrait(..) => {
+                        write!(f, "{{async block}}")?;
                     }
                 };
-                write!(f, "impl ")?;
-                write_bounds_like_dyn_trait(&bounds.value, f)?;
             }
             Ty::Unknown => write!(f, "{{unknown}}")?,
             Ty::Infer(..) => write!(f, "_")?,

crates/hir_ty/src/infer/expr.rs

@@ -17,8 +17,8 @@ use crate::{
     autoderef, method_resolution, op,
     traits::{FnTrait, InEnvironment},
     utils::{generics, variant_data, Generics},
-    ApplicationTy, Binders, CallableDefId, InferTy, IntTy, Mutability, Obligation, Rawness, Substs,
-    TraitRef, Ty, TypeCtor,
+    ApplicationTy, Binders, CallableDefId, InferTy, IntTy, Mutability, Obligation, OpaqueTyId,
+    Rawness, Substs, TraitRef, Ty, TypeCtor,
 };
 
 use super::{
@@ -146,6 +146,13 @@ impl<'a> InferenceContext<'a> {
                 // FIXME should be std::result::Result<{inner}, _>
                 Ty::Unknown
             }
+            Expr::Async { body } => {
+                // Use the first type parameter as the output type of future.
+                // existenail type AsyncBlockImplTrait<InnerType>: Future<Output = InnerType>
+                let inner_ty = self.infer_expr(*body, &Expectation::none());
+                let opaque_ty_id = OpaqueTyId::AsyncBlockTypeImplTrait(self.owner, *body);
+                Ty::apply_one(TypeCtor::OpaqueType(opaque_ty_id), inner_ty)
+            }
             Expr::Loop { body, label } => {
                 self.breakables.push(BreakableContext {
                     may_break: false,

crates/hir_ty/src/lib.rs

@@ -129,8 +129,9 @@ pub enum TypeCtor {
 
     /// This represents a placeholder for an opaque type in situations where we
     /// don't know the hidden type (i.e. currently almost always). This is
-    /// analogous to the `AssociatedType` type constructor. As with that one,
-    /// these are only produced by Chalk.
+    /// analogous to the `AssociatedType` type constructor.
+    /// It is also used as the type of async block, with one type parameter
+    /// representing the Future::Output type.
     OpaqueType(OpaqueTyId),
 
     /// The type of a specific closure.
@@ -173,6 +174,8 @@ impl TypeCtor {
                         let generic_params = generics(db.upcast(), func.into());
                         generic_params.len()
                     }
+                    // 1 param representing Future::Output type.
+                    OpaqueTyId::AsyncBlockTypeImplTrait(..) => 1,
                 }
             }
             TypeCtor::FnPtr { num_args, is_varargs: _ } => num_args as usize + 1,
@@ -205,6 +208,7 @@ impl TypeCtor {
                 OpaqueTyId::ReturnTypeImplTrait(func, _) => {
                     Some(func.lookup(db.upcast()).module(db.upcast()).krate)
                 }
+                OpaqueTyId::AsyncBlockTypeImplTrait(def, _) => Some(def.module(db.upcast()).krate),
             },
         }
     }
@@ -843,6 +847,29 @@ impl Ty {
 
     pub fn impl_trait_bounds(&self, db: &dyn HirDatabase) -> Option<Vec<GenericPredicate>> {
         match self {
+            Ty::Apply(ApplicationTy { ctor: TypeCtor::OpaqueType(opaque_ty_id), .. }) => {
+                match opaque_ty_id {
+                    OpaqueTyId::AsyncBlockTypeImplTrait(def, _expr) => {
+                        let krate = def.module(db.upcast()).krate;
+                        if let Some(future_trait) = db
+                            .lang_item(krate, "future_trait".into())
+                            .and_then(|item| item.as_trait())
+                        {
+                            // This is only used by type walking.
+                            // Parameters will be walked outside, and projection predicate is not used.
+                            // So just provide the Future trait.
+                            let impl_bound = GenericPredicate::Implemented(TraitRef {
+                                trait_: future_trait,
+                                substs: Substs::empty(),
+                            });
+                            Some(vec![impl_bound])
+                        } else {
+                            None
+                        }
+                    }
+                    OpaqueTyId::ReturnTypeImplTrait(..) => None,
+                }
+            }
             Ty::Opaque(opaque_ty) => {
                 let predicates = match opaque_ty.opaque_ty_id {
                     OpaqueTyId::ReturnTypeImplTrait(func, idx) => {
@@ -853,6 +880,8 @@ impl Ty {
                             data.subst(&opaque_ty.parameters)
                         })
                     }
+                    // It always has an parameter for Future::Output type.
+                    OpaqueTyId::AsyncBlockTypeImplTrait(..) => unreachable!(),
                 };
 
                 predicates.map(|it| it.value)
@@ -1065,6 +1094,7 @@ impl<T: TypeWalk> TypeWalk for Vec<T> {
 #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)]
 pub enum OpaqueTyId {
     ReturnTypeImplTrait(hir_def::FunctionId, u16),
+    AsyncBlockTypeImplTrait(hir_def::DefWithBodyId, ExprId),
 }
 
 #[derive(Clone, PartialEq, Eq, Debug, Hash)]

crates/hir_ty/src/tests/simple.rs

@@ -1889,31 +1889,40 @@ fn fn_pointer_return() {
 fn effects_smoke_test() {
     check_infer(
         r#"
-        fn main() {
+        async fn main() {
             let x = unsafe { 92 };
             let y = async { async { () }.await };
             let z = try { () };
             let t = 'a: { 92 };
         }
+
+        #[prelude_import] use future::*;
+
+        mod future {
+            #[lang = "future_trait"]
+            pub trait Future { type Output; }
+        }
         "#,
         expect![[r#"
-            10..130 '{     ...2 }; }': ()
-            20..21 'x': i32
-            24..37 'unsafe { 92 }': i32
-            31..37 '{ 92 }': i32
-            33..35 '92': i32
-            47..48 'y': {unknown}
-            57..79 '{ asyn...wait }': {unknown}
-            59..77 'async ....await': {unknown}
-            65..71 '{ () }': ()
-            67..69 '()': ()
-            89..90 'z': {unknown}
-            93..103 'try { () }': {unknown}
-            97..103 '{ () }': ()
-            99..101 '()': ()
-            113..114 't': i32
-            121..127 '{ 92 }': i32
-            123..125 '92': i32
+            16..136 '{     ...2 }; }': ()
+            26..27 'x': i32
+            30..43 'unsafe { 92 }': i32
+            37..43 '{ 92 }': i32
+            39..41 '92': i32
+            53..54 'y': impl Future<Output = ()>
+            57..85 'async ...wait }': impl Future<Output = ()>
+            63..85 '{ asyn...wait }': ()
+            65..77 'async { () }': impl Future<Output = ()>
+            65..83 'async ....await': ()
+            71..77 '{ () }': ()
+            73..75 '()': ()
+            95..96 'z': {unknown}
+            99..109 'try { () }': {unknown}
+            103..109 '{ () }': ()
+            105..107 '()': ()
+            119..120 't': i32
+            127..133 '{ 92 }': i32
+            129..131 '92': i32
         "#]],
     )
 }

crates/hir_ty/src/tests/traits.rs

@@ -85,6 +85,46 @@ mod future {
     );
 }
 
+#[test]
+fn infer_async_block() {
+    check_types(
+        r#"
+//- /main.rs crate:main deps:core
+async fn test() {
+    let a = async { 42 };
+    a;
+//  ^ impl Future<Output = i32>
+    let x = a.await;
+    x;
+//  ^ i32
+    let b = async {}.await;
+    b;
+//  ^ ()
+    let c = async {
+        let y = Option::None;
+        y
+    //  ^ Option<u64>
+    };
+    let _: Option<u64> = c.await;
+    c;
+//  ^ impl Future<Output = Option<u64>>
+}
+
+enum Option<T> { None, Some(T) }
+
+//- /core.rs crate:core
+#[prelude_import] use future::*;
+mod future {
+    #[lang = "future_trait"]
+    trait Future {
+        type Output;
+    }
+}
+
+"#,
+    );
+}
+
 #[test]
 fn infer_try() {
     check_types(