fix: Diagnose non-value return and break type mismatches

3 files changed, 62 insertions, 23 deletions

diff --git a/crates/hir-ty/src/infer/coerce.rs b/crates/hir-ty/src/infer/coerce.rs
index e9c94449cfc..05a476f632d 100644
--- a/crates/hir-ty/src/infer/coerce.rs
+++ b/crates/hir-ty/src/infer/coerce.rs
@@ -50,6 +50,13 @@ fn success(
     Ok(InferOk { goals, value: (adj, target) })
 }
 
+pub(super) enum CoercionCause {
+    // FIXME: Make better use of this. Right now things like return and break without a value
+    // use it to point to themselves, causing us to report a mismatch on those expressions even
+    // though technically they themselves are `!`
+    Expr(ExprId),
+}
+
 #[derive(Clone, Debug)]
 pub(super) struct CoerceMany {
     expected_ty: Ty,
@@ -90,8 +97,12 @@ impl CoerceMany {
         }
     }
 
-    pub(super) fn coerce_forced_unit(&mut self, ctx: &mut InferenceContext<'_>) {
-        self.coerce(ctx, None, &ctx.result.standard_types.unit.clone())
+    pub(super) fn coerce_forced_unit(
+        &mut self,
+        ctx: &mut InferenceContext<'_>,
+        cause: CoercionCause,
+    ) {
+        self.coerce(ctx, None, &ctx.result.standard_types.unit.clone(), cause)
     }
 
     /// Merge two types from different branches, with possible coercion.
@@ -106,6 +117,7 @@ impl CoerceMany {
         ctx: &mut InferenceContext<'_>,
         expr: Option<ExprId>,
         expr_ty: &Ty,
+        cause: CoercionCause,
     ) {
         let expr_ty = ctx.resolve_ty_shallow(expr_ty);
         self.expected_ty = ctx.resolve_ty_shallow(&self.expected_ty);
@@ -153,11 +165,13 @@ impl CoerceMany {
         } else if let Ok(res) = ctx.coerce(expr, &self.merged_ty(), &expr_ty) {
             self.final_ty = Some(res);
         } else {
-            if let Some(id) = expr {
-                ctx.result.type_mismatches.insert(
-                    id.into(),
-                    TypeMismatch { expected: self.merged_ty(), actual: expr_ty.clone() },
-                );
+            match cause {
+                CoercionCause::Expr(id) => {
+                    ctx.result.type_mismatches.insert(
+                        id.into(),
+                        TypeMismatch { expected: self.merged_ty(), actual: expr_ty.clone() },
+                    );
+                }
             }
             cov_mark::hit!(coerce_merge_fail_fallback);
         }
diff --git a/crates/hir-ty/src/infer/expr.rs b/crates/hir-ty/src/infer/expr.rs
index 96218e4fb09..353bf6568c1 100644
--- a/crates/hir-ty/src/infer/expr.rs
+++ b/crates/hir-ty/src/infer/expr.rs
@@ -26,7 +26,10 @@ use crate::{
     autoderef::{builtin_deref, deref_by_trait, Autoderef},
     consteval,
     infer::{
-        coerce::CoerceMany, find_continuable, pat::contains_explicit_ref_binding, BreakableKind,
+        coerce::{CoerceMany, CoercionCause},
+        find_continuable,
+        pat::contains_explicit_ref_binding,
+        BreakableKind,
     },
     lang_items::lang_items_for_bin_op,
     lower::{
@@ -132,24 +135,28 @@ impl<'a> InferenceContext<'a> {
                 );
 
                 let condition_diverges = mem::replace(&mut self.diverges, Diverges::Maybe);
-                let mut both_arms_diverge = Diverges::Always;
 
                 let then_ty = self.infer_expr_inner(then_branch, expected);
-                both_arms_diverge &= mem::replace(&mut self.diverges, Diverges::Maybe);
+                let then_diverges = mem::replace(&mut self.diverges, Diverges::Maybe);
                 let mut coerce = CoerceMany::new(expected.coercion_target_type(&mut self.table));
-                coerce.coerce(self, Some(then_branch), &then_ty);
+                coerce.coerce(self, Some(then_branch), &then_ty, CoercionCause::Expr(then_branch));
                 match else_branch {
                     Some(else_branch) => {
                         let else_ty = self.infer_expr_inner(else_branch, expected);
-                        coerce.coerce(self, Some(else_branch), &else_ty);
+                        let else_diverges = mem::replace(&mut self.diverges, Diverges::Maybe);
+                        coerce.coerce(
+                            self,
+                            Some(else_branch),
+                            &else_ty,
+                            CoercionCause::Expr(else_branch),
+                        );
+                        self.diverges = condition_diverges | then_diverges & else_diverges;
                     }
                     None => {
-                        coerce.coerce_forced_unit(self);
+                        coerce.coerce_forced_unit(self, CoercionCause::Expr(tgt_expr));
+                        self.diverges = condition_diverges;
                     }
                 }
-                both_arms_diverge &= self.diverges;
-
-                self.diverges = condition_diverges | both_arms_diverge;
 
                 coerce.complete(self)
             }
@@ -444,7 +451,7 @@ impl<'a> InferenceContext<'a> {
 
                         let arm_ty = self.infer_expr_inner(arm.expr, &expected);
                         all_arms_diverge &= self.diverges;
-                        coerce.coerce(self, Some(arm.expr), &arm_ty);
+                        coerce.coerce(self, Some(arm.expr), &arm_ty, CoercionCause::Expr(arm.expr));
                     }
 
                     self.diverges = matchee_diverges | all_arms_diverge;
@@ -492,7 +499,11 @@ impl<'a> InferenceContext<'a> {
                 match find_breakable(&mut self.breakables, label) {
                     Some(ctxt) => match ctxt.coerce.take() {
                         Some(mut coerce) => {
-                            coerce.coerce(self, expr, &val_ty);
+                            let cause = match expr {
+                                Some(expr) => CoercionCause::Expr(expr),
+                                None => CoercionCause::Expr(tgt_expr),
+                            };
+                            coerce.coerce(self, expr, &val_ty, cause);
 
                             // Avoiding borrowck
                             let ctxt = find_breakable(&mut self.breakables, label)
@@ -512,7 +523,7 @@ impl<'a> InferenceContext<'a> {
                 }
                 self.result.standard_types.never.clone()
             }
-            &Expr::Return { expr } => self.infer_expr_return(expr),
+            &Expr::Return { expr } => self.infer_expr_return(tgt_expr, expr),
             Expr::Yield { expr } => {
                 if let Some((resume_ty, yield_ty)) = self.resume_yield_tys.clone() {
                     if let Some(expr) = expr {
@@ -952,7 +963,7 @@ impl<'a> InferenceContext<'a> {
                 let mut coerce = CoerceMany::new(elem_ty);
                 for &expr in elements.iter() {
                     let cur_elem_ty = self.infer_expr_inner(expr, &expected);
-                    coerce.coerce(self, Some(expr), &cur_elem_ty);
+                    coerce.coerce(self, Some(expr), &cur_elem_ty, CoercionCause::Expr(expr));
                 }
                 (
                     coerce.complete(self),
@@ -997,18 +1008,18 @@ impl<'a> InferenceContext<'a> {
             .expected_ty();
         let return_expr_ty = 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);
+        coerce_many.coerce(self, Some(expr), &return_expr_ty, CoercionCause::Expr(expr));
         self.return_coercion = Some(coerce_many);
     }
 
-    fn infer_expr_return(&mut self, expr: Option<ExprId>) -> Ty {
+    fn infer_expr_return(&mut self, ret: ExprId, expr: Option<ExprId>) -> Ty {
         match self.return_coercion {
             Some(_) => {
                 if let Some(expr) = expr {
                     self.infer_return(expr);
                 } else {
                     let mut coerce = self.return_coercion.take().unwrap();
-                    coerce.coerce_forced_unit(self);
+                    coerce.coerce_forced_unit(self, CoercionCause::Expr(ret));
                     self.return_coercion = Some(coerce);
                 }
             }
diff --git a/crates/ide-diagnostics/src/handlers/type_mismatch.rs b/crates/ide-diagnostics/src/handlers/type_mismatch.rs
index c462a16e362..a5359741ac9 100644
--- a/crates/ide-diagnostics/src/handlers/type_mismatch.rs
+++ b/crates/ide-diagnostics/src/handlers/type_mismatch.rs
@@ -681,4 +681,18 @@ struct Bar {
 "#,
         );
     }
+
+    #[test]
+    fn return_no_value() {
+        check_diagnostics(
+            r#"
+fn f() -> i32 {
+    return;
+ // ^^^^^^ error: expected i32, found ()
+    0
+}
+fn g() { return; }
+"#,
+        );
+    }
 }