From ed0c62597f9fdbe04e7c99988dad6f102acc4473 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Wed, 23 Mar 2022 23:39:35 +0100
Subject: [PATCH] fix: Fix async block type inference using containing function
 return type

---
 crates/hir_ty/src/infer.rs        |  2 +-
 crates/hir_ty/src/infer/expr.rs   | 10 ++++++-
 crates/hir_ty/src/tests/simple.rs | 46 +++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/crates/hir_ty/src/infer.rs b/crates/hir_ty/src/infer.rs
index c91cdec1834..e78d3f267f0 100644
--- a/crates/hir_ty/src/infer.rs
+++ b/crates/hir_ty/src/infer.rs
@@ -364,7 +364,7 @@ pub(crate) struct InferenceContext<'a> {
     table: unify::InferenceTable<'a>,
     trait_env: Arc<TraitEnvironment>,
     result: InferenceResult,
-    /// The return type of the function being inferred, or the closure if we're
+    /// The return type of the function being inferred, the closure or async block if we're
     /// currently within one.
     ///
     /// We might consider using a nested inference context for checking
diff --git a/crates/hir_ty/src/infer/expr.rs b/crates/hir_ty/src/infer/expr.rs
index 43cff92f233..01b11861cd9 100644
--- a/crates/hir_ty/src/infer/expr.rs
+++ b/crates/hir_ty/src/infer/expr.rs
@@ -157,9 +157,17 @@ fn infer_expr_inner(&mut self, tgt_expr: ExprId, expected: &Expectation) -> Ty {
                 self.err_ty()
             }
             Expr::Async { body } => {
+                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 inner_ty = self.infer_expr_coerce(*body, &Expectation::has_type(ret_ty));
+
+                self.diverges = prev_diverges;
+                self.return_ty = prev_ret_ty;
+
                 // Use the first type parameter as the output type of future.
                 // existential type AsyncBlockImplTrait<InnerType>: Future<Output = InnerType>
-                let inner_ty = self.infer_expr(*body, &Expectation::none());
                 let impl_trait_id = crate::ImplTraitId::AsyncBlockTypeImplTrait(self.owner, *body);
                 let opaque_ty_id = self.db.intern_impl_trait_id(impl_trait_id).into();
                 TyKind::OpaqueType(opaque_ty_id, Substitution::from1(Interner, inner_ty))
diff --git a/crates/hir_ty/src/tests/simple.rs b/crates/hir_ty/src/tests/simple.rs
index de27c294f61..31045c193c9 100644
--- a/crates/hir_ty/src/tests/simple.rs
+++ b/crates/hir_ty/src/tests/simple.rs
@@ -1952,6 +1952,52 @@ async fn main() {
         "#]],
     )
 }
+#[test]
+fn async_block_early_return() {
+    check_infer(
+        r#"
+//- minicore: future, result, fn
+fn test<I, E, F: FnMut() -> Fut, Fut: core::future::Future<Output = Result<I, E>>>(f: F) {}
+
+fn main() {
+    async {
+        return Err(());
+        Ok(())
+    };
+    test(|| async {
+        return Err(());
+        Ok(())
+    });
+}
+        "#,
+        expect![[r#"
+            83..84 'f': F
+            89..91 '{}': ()
+            103..231 '{     ... }); }': ()
+            109..161 'async ...     }': Result<(), ()>
+            109..161 'async ...     }': impl Future<Output = Result<(), ()>>
+            125..139 'return Err(())': !
+            132..135 'Err': Err<(), ()>(()) -> Result<(), ()>
+            132..139 'Err(())': Result<(), ()>
+            136..138 '()': ()
+            149..151 'Ok': Ok<(), ()>(()) -> Result<(), ()>
+            149..155 'Ok(())': Result<(), ()>
+            152..154 '()': ()
+            167..171 'test': fn test<(), (), || -> impl Future<Output = Result<(), ()>>, impl Future<Output = Result<(), ()>>>(|| -> impl Future<Output = Result<(), ()>>)
+            167..228 'test(|...    })': ()
+            172..227 '|| asy...     }': || -> impl Future<Output = Result<(), ()>>
+            175..227 'async ...     }': Result<(), ()>
+            175..227 'async ...     }': impl Future<Output = Result<(), ()>>
+            191..205 'return Err(())': !
+            198..201 'Err': Err<(), ()>(()) -> Result<(), ()>
+            198..205 'Err(())': Result<(), ()>
+            202..204 '()': ()
+            215..217 'Ok': Ok<(), ()>(()) -> Result<(), ()>
+            215..221 'Ok(())': Result<(), ()>
+            218..220 '()': ()
+        "#]],
+    )
+}
 
 #[test]
 fn infer_generic_from_later_assignment() {