Auto merge of #14461 - HKalbasi:dev, r=Veykril

Use async block in async fn type inference

fix #14456

At some point we should probably go further and completely desugar async fn in hir lowering.

7 files changed, 68 insertions, 57 deletions

diff --git a/crates/hir-def/src/data.rs b/crates/hir-def/src/data.rs
index dcc21db7409..431140a665d 100644
--- a/crates/hir-def/src/data.rs
+++ b/crates/hir-def/src/data.rs
@@ -30,7 +30,6 @@ pub struct FunctionData {
     pub name: Name,
     pub params: Vec<(Option<Name>, Interned<TypeRef>)>,
     pub ret_type: Interned<TypeRef>,
-    pub async_ret_type: Option<Interned<TypeRef>>,
     pub attrs: Attrs,
     pub visibility: RawVisibility,
     pub abi: Option<Interned<str>>,
@@ -104,7 +103,6 @@ impl FunctionData {
                 })
                 .collect(),
             ret_type: func.ret_type.clone(),
-            async_ret_type: func.async_ret_type.clone(),
             attrs: item_tree.attrs(db, krate, ModItem::from(loc.id.value).into()),
             visibility,
             abi: func.abi.clone(),
diff --git a/crates/hir-def/src/item_tree.rs b/crates/hir-def/src/item_tree.rs
index 06942170179..d5b973751dd 100644
--- a/crates/hir-def/src/item_tree.rs
+++ b/crates/hir-def/src/item_tree.rs
@@ -608,7 +608,6 @@ pub struct Function {
     pub abi: Option<Interned<str>>,
     pub params: IdxRange<Param>,
     pub ret_type: Interned<TypeRef>,
-    pub async_ret_type: Option<Interned<TypeRef>>,
     pub ast_id: FileAstId<ast::Fn>,
     pub(crate) flags: FnFlags,
 }
diff --git a/crates/hir-def/src/item_tree/lower.rs b/crates/hir-def/src/item_tree/lower.rs
index 49deff080e5..5bbf8e52e8b 100644
--- a/crates/hir-def/src/item_tree/lower.rs
+++ b/crates/hir-def/src/item_tree/lower.rs
@@ -364,13 +364,12 @@ impl<'a> Ctx<'a> {
             None => TypeRef::unit(),
         };
 
-        let (ret_type, async_ret_type) = if func.async_token().is_some() {
-            let async_ret_type = ret_type.clone();
+        let ret_type = if func.async_token().is_some() {
             let future_impl = desugar_future_path(ret_type);
             let ty_bound = Interned::new(TypeBound::Path(future_impl, TraitBoundModifier::None));
-            (TypeRef::ImplTrait(vec![ty_bound]), Some(async_ret_type))
+            TypeRef::ImplTrait(vec![ty_bound])
         } else {
-            (ret_type, None)
+            ret_type
         };
 
         let abi = func.abi().map(lower_abi);
@@ -404,7 +403,6 @@ impl<'a> Ctx<'a> {
             abi,
             params,
             ret_type: Interned::new(ret_type),
-            async_ret_type: async_ret_type.map(Interned::new),
             ast_id,
             flags,
         };
diff --git a/crates/hir-def/src/item_tree/pretty.rs b/crates/hir-def/src/item_tree/pretty.rs
index 5f29997964b..edd5c3b1151 100644
--- a/crates/hir-def/src/item_tree/pretty.rs
+++ b/crates/hir-def/src/item_tree/pretty.rs
@@ -233,7 +233,6 @@ impl<'a> Printer<'a> {
                     abi,
                     params,
                     ret_type,
-                    async_ret_type: _,
                     ast_id: _,
                     flags,
                 } = &self.tree[it];
diff --git a/crates/hir-ty/src/infer.rs b/crates/hir-ty/src/infer.rs
index ab08593bcf6..08ba80cdfff 100644
--- a/crates/hir-ty/src/infer.rs
+++ b/crates/hir-ty/src/infer.rs
@@ -459,6 +459,7 @@ pub(crate) struct InferenceContext<'a> {
     resume_yield_tys: Option<(Ty, Ty)>,
     diverges: Diverges,
     breakables: Vec<BreakableContext>,
+    is_async_fn: bool,
 }
 
 #[derive(Clone, Debug)]
@@ -526,6 +527,7 @@ impl<'a> InferenceContext<'a> {
             resolver,
             diverges: Diverges::Maybe,
             breakables: Vec::new(),
+            is_async_fn: false,
         }
     }
 
@@ -636,12 +638,10 @@ impl<'a> InferenceContext<'a> {
 
             self.infer_top_pat(*pat, &ty);
         }
-        let error_ty = &TypeRef::Error;
-        let return_ty = if data.has_async_kw() {
-            data.async_ret_type.as_deref().unwrap_or(error_ty)
-        } else {
-            &*data.ret_type
-        };
+        let return_ty = &*data.ret_type;
+        if data.has_async_kw() {
+            self.is_async_fn = true;
+        }
 
         let ctx = crate::lower::TyLoweringContext::new(self.db, &self.resolver)
             .with_impl_trait_mode(ImplTraitLoweringMode::Opaque);
diff --git a/crates/hir-ty/src/infer/expr.rs b/crates/hir-ty/src/infer/expr.rs
index 38a09953027..23ef32db227 100644
--- a/crates/hir-ty/src/infer/expr.rs
+++ b/crates/hir-ty/src/infer/expr.rs
@@ -162,35 +162,7 @@ impl<'a> InferenceContext<'a> {
                 .1
             }
             Expr::Async { id, statements, tail } => {
-                let ret_ty = self.table.new_type_var();
-                let prev_diverges = mem::replace(&mut self.diverges, Diverges::Maybe);
-                let prev_ret_ty = mem::replace(&mut self.return_ty, ret_ty.clone());
-                let prev_ret_coercion =
-                    mem::replace(&mut self.return_coercion, Some(CoerceMany::new(ret_ty.clone())));
-
-                let (_, inner_ty) =
-                    self.with_breakable_ctx(BreakableKind::Border, None, None, |this| {
-                        this.infer_block(
-                            tgt_expr,
-                            *id,
-                            statements,
-                            *tail,
-                            None,
-                            &Expectation::has_type(ret_ty),
-                        )
-                    });
-
-                self.diverges = prev_diverges;
-                self.return_ty = prev_ret_ty;
-                self.return_coercion = prev_ret_coercion;
-
-                // Use the first type parameter as the output type of future.
-                // existential type AsyncBlockImplTrait<InnerType>: Future<Output = InnerType>
-                let impl_trait_id =
-                    crate::ImplTraitId::AsyncBlockTypeImplTrait(self.owner, tgt_expr);
-                let opaque_ty_id = self.db.intern_impl_trait_id(impl_trait_id).into();
-                TyKind::OpaqueType(opaque_ty_id, Substitution::from1(Interner, inner_ty))
-                    .intern(Interner)
+                self.infer_async_block(tgt_expr, id, statements, tail)
             }
             &Expr::Loop { body, label } => {
                 // FIXME: should be:
@@ -260,18 +232,7 @@ impl<'a> InferenceContext<'a> {
                     None => self.table.new_type_var(),
                 };
                 if let ClosureKind::Async = closure_kind {
-                    // Use the first type parameter as the output type of future.
-                    // existential type AsyncBlockImplTrait<InnerType>: Future<Output = InnerType>
-                    let impl_trait_id =
-                        crate::ImplTraitId::AsyncBlockTypeImplTrait(self.owner, *body);
-                    let opaque_ty_id = self.db.intern_impl_trait_id(impl_trait_id).into();
-                    sig_tys.push(
-                        TyKind::OpaqueType(
-                            opaque_ty_id,
-                            Substitution::from1(Interner, ret_ty.clone()),
-                        )
-                        .intern(Interner),
-                    );
+                    sig_tys.push(self.lower_async_block_type_impl_trait(ret_ty.clone(), *body));
                 } else {
                     sig_tys.push(ret_ty.clone());
                 }
@@ -332,6 +293,7 @@ impl<'a> InferenceContext<'a> {
 
                 // FIXME: lift these out into a struct
                 let prev_diverges = mem::replace(&mut self.diverges, Diverges::Maybe);
+                let prev_is_async_fn = mem::replace(&mut self.is_async_fn, false);
                 let prev_ret_ty = mem::replace(&mut self.return_ty, ret_ty.clone());
                 let prev_ret_coercion =
                     mem::replace(&mut self.return_coercion, Some(CoerceMany::new(ret_ty)));
@@ -345,6 +307,7 @@ impl<'a> InferenceContext<'a> {
                 self.diverges = prev_diverges;
                 self.return_ty = prev_ret_ty;
                 self.return_coercion = prev_ret_coercion;
+                self.is_async_fn = prev_is_async_fn;
                 self.resume_yield_tys = prev_resume_yield_tys;
 
                 ty
@@ -900,6 +863,42 @@ impl<'a> InferenceContext<'a> {
         ty
     }
 
+    fn infer_async_block(
+        &mut self,
+        tgt_expr: ExprId,
+        id: &Option<BlockId>,
+        statements: &[Statement],
+        tail: &Option<ExprId>,
+    ) -> Ty {
+        let ret_ty = self.table.new_type_var();
+        let prev_diverges = mem::replace(&mut self.diverges, Diverges::Maybe);
+        let prev_ret_ty = mem::replace(&mut self.return_ty, ret_ty.clone());
+        let prev_ret_coercion =
+            mem::replace(&mut self.return_coercion, Some(CoerceMany::new(ret_ty.clone())));
+
+        let (_, inner_ty) = self.with_breakable_ctx(BreakableKind::Border, None, None, |this| {
+            this.infer_block(tgt_expr, *id, statements, *tail, None, &Expectation::has_type(ret_ty))
+        });
+
+        self.diverges = prev_diverges;
+        self.return_ty = prev_ret_ty;
+        self.return_coercion = prev_ret_coercion;
+
+        self.lower_async_block_type_impl_trait(inner_ty, tgt_expr)
+    }
+
+    pub(crate) fn lower_async_block_type_impl_trait(
+        &mut self,
+        inner_ty: Ty,
+        tgt_expr: ExprId,
+    ) -> Ty {
+        // Use the first type parameter as the output type of future.
+        // existential type AsyncBlockImplTrait<InnerType>: Future<Output = InnerType>
+        let impl_trait_id = crate::ImplTraitId::AsyncBlockTypeImplTrait(self.owner, tgt_expr);
+        let opaque_ty_id = self.db.intern_impl_trait_id(impl_trait_id).into();
+        TyKind::OpaqueType(opaque_ty_id, Substitution::from1(Interner, inner_ty)).intern(Interner)
+    }
+
     fn infer_expr_array(
         &mut self,
         array: &Array,
@@ -964,7 +963,11 @@ impl<'a> InferenceContext<'a> {
             .as_mut()
             .expect("infer_return called outside function body")
             .expected_ty();
-        let return_expr_ty = self.infer_expr_inner(expr, &Expectation::HasType(ret_ty));
+        let return_expr_ty = if self.is_async_fn {
+            self.infer_async_block(expr, &None, &[], &Some(expr))
+        } else {
+            self.infer_expr_inner(expr, &Expectation::HasType(ret_ty))
+        };
         let mut coerce_many = self.return_coercion.take().unwrap();
         coerce_many.coerce(self, Some(expr), &return_expr_ty);
         self.return_coercion = Some(coerce_many);
diff --git a/crates/hir-ty/src/tests/regression.rs b/crates/hir-ty/src/tests/regression.rs
index 689f0da44f6..8911dd318c4 100644
--- a/crates/hir-ty/src/tests/regression.rs
+++ b/crates/hir-ty/src/tests/regression.rs
@@ -1758,6 +1758,20 @@ const C: usize = 2 + 2;
 }
 
 #[test]
+fn regression_14456() {
+    check_no_mismatches(
+        r#"
+//- minicore: future
+async fn x() {}
+fn f() {
+    let fut = x();
+    let t = [0u8; 2 + 2];
+}
+"#,
+    );
+}
+
+#[test]
 fn regression_14164() {
     check_types(
         r#"