fix: Fix async block type inference using containing function return type

Author: Veykril

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

crates/hir_ty/src/infer/expr.rs

@@ -157,9 +157,17 @@ impl<'a> InferenceContext<'a> {
                 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))

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() {