Provide structured suggestion for type mismatch in loop

compiler/rustc_hir_typeck/src/coercion.rs

@@ -36,7 +36,9 @@
 //! ```
 
 use crate::FnCtxt;
-use rustc_errors::{struct_span_err, Diagnostic, DiagnosticBuilder, ErrorGuaranteed, MultiSpan};
+use rustc_errors::{
+    struct_span_err, Applicability, Diagnostic, DiagnosticBuilder, ErrorGuaranteed, MultiSpan,
+};
 use rustc_hir as hir;
 use rustc_hir::def_id::DefId;
 use rustc_hir::intravisit::{self, Visitor};
@@ -53,8 +55,7 @@ use rustc_middle::ty::adjustment::{
 use rustc_middle::ty::error::TypeError;
 use rustc_middle::ty::relate::RelateResult;
 use rustc_middle::ty::visit::TypeVisitableExt;
-use rustc_middle::ty::GenericArgsRef;
-use rustc_middle::ty::{self, Ty, TypeAndMut};
+use rustc_middle::ty::{self, GenericArgsRef, Ty, TyCtxt, TypeAndMut};
 use rustc_session::parse::feature_err;
 use rustc_span::symbol::sym;
 use rustc_span::{self, DesugaringKind};
@@ -1660,12 +1661,15 @@ impl<'tcx, 'exprs, E: AsCoercionSite> CoerceMany<'tcx, 'exprs, E> {
                         None,
                         Some(coercion_error),
                     );
-                }
-
-                if visitor.ret_exprs.len() > 0
-                    && let Some(expr) = expression
-                {
-                    self.note_unreachable_loop_return(&mut err, expr, &visitor.ret_exprs);
+                    if visitor.ret_exprs.len() > 0 {
+                        self.note_unreachable_loop_return(
+                            &mut err,
+                            fcx.tcx,
+                            &expr,
+                            &visitor.ret_exprs,
+                            expected,
+                        );
+                    }
                 }
 
                 let reported = err.emit_unless(unsized_return);
@@ -1678,8 +1682,10 @@ impl<'tcx, 'exprs, E: AsCoercionSite> CoerceMany<'tcx, 'exprs, E> {
     fn note_unreachable_loop_return(
         &self,
         err: &mut Diagnostic,
+        tcx: TyCtxt<'tcx>,
         expr: &hir::Expr<'tcx>,
         ret_exprs: &Vec<&'tcx hir::Expr<'tcx>>,
+        ty: Ty<'tcx>,
     ) {
         let hir::ExprKind::Loop(_, _, _, loop_span) = expr.kind else {
             return;
@@ -1704,10 +1710,77 @@ impl<'tcx, 'exprs, E: AsCoercionSite> CoerceMany<'tcx, 'exprs, E> {
                 ret_exprs.len() - MAXITER
             ));
         }
-        err.help(
-            "return a value for the case when the loop has zero elements to iterate on, or \
-           consider changing the return type to account for that possibility",
-        );
+        let hir = tcx.hir();
+        let item = hir.get_parent_item(expr.hir_id);
+        let ret_msg = "return a value for the case when the loop has zero elements to iterate on";
+        let ret_ty_msg =
+            "otherwise consider changing the return type to account for that possibility";
+        if let Some(node) = hir.find(item.into())
+            && let Some(body_id) = node.body_id()
+            && let Some(sig) = node.fn_sig()
+            && let hir::ExprKind::Block(block, _) = hir.body(body_id).value.kind
+            && !ty.is_never()
+        {
+            let indentation = if let None = block.expr
+                && let [.., last] = &block.stmts[..]
+            {
+                tcx.sess.source_map().indentation_before(last.span).unwrap_or_else(String::new)
+            } else if let Some(expr) = block.expr {
+                tcx.sess.source_map().indentation_before(expr.span).unwrap_or_else(String::new)
+            } else {
+                String::new()
+            };
+            if let None = block.expr
+                && let [.., last] = &block.stmts[..]
+            {
+                err.span_suggestion_verbose(
+                    last.span.shrink_to_hi(),
+                    ret_msg,
+                    format!("\n{indentation}/* `{ty}` value */"),
+                    Applicability::MaybeIncorrect,
+                );
+            } else if let Some(expr) = block.expr {
+                err.span_suggestion_verbose(
+                    expr.span.shrink_to_hi(),
+                    ret_msg,
+                    format!("\n{indentation}/* `{ty}` value */"),
+                    Applicability::MaybeIncorrect,
+                );
+            }
+            let mut sugg = match sig.decl.output {
+                hir::FnRetTy::DefaultReturn(span) => {
+                    vec![(span, " -> Option<()>".to_string())]
+                }
+                hir::FnRetTy::Return(ty) => {
+                    vec![
+                        (ty.span.shrink_to_lo(), "Option<".to_string()),
+                        (ty.span.shrink_to_hi(), ">".to_string()),
+                    ]
+                }
+            };
+            for ret_expr in ret_exprs {
+                match ret_expr.kind {
+                    hir::ExprKind::Ret(Some(expr)) => {
+                        sugg.push((expr.span.shrink_to_lo(), "Some(".to_string()));
+                        sugg.push((expr.span.shrink_to_hi(), ")".to_string()));
+                    }
+                    hir::ExprKind::Ret(None) => {
+                        sugg.push((ret_expr.span.shrink_to_hi(), " Some(())".to_string()));
+                    }
+                    _ => {}
+                }
+            }
+            if let None = block.expr
+                && let [.., last] = &block.stmts[..]
+            {
+                sugg.push((last.span.shrink_to_hi(), format!("\n{indentation}None")));
+            } else if let Some(expr) = block.expr {
+                sugg.push((expr.span.shrink_to_hi(), format!("\n{indentation}None")));
+            }
+            err.multipart_suggestion(ret_ty_msg, sugg, Applicability::MaybeIncorrect);
+        } else {
+            err.help(format!("{ret_msg}, {ret_ty_msg}"));
+        }
     }
 
     fn report_return_mismatched_types<'a>(

tests/ui/for-loop-while/break-while-condition.stderr

@@ -38,7 +38,7 @@ LL |             while false {
    |             ^^^^^^^^^^^ this might have zero elements to iterate on
 LL |                 return
    |                 ------ if the loop doesn't execute, this value would never get returned
-   = help: return a value for the case when the loop has zero elements to iterate on, or consider changing the return type to account for that possibility
+   = help: return a value for the case when the loop has zero elements to iterate on, otherwise consider changing the return type to account for that possibility
 
 error: aborting due to 3 previous errors
 

tests/ui/typeck/issue-100285.rs

@@ -1,6 +1,5 @@
-fn foo(n: i32) -> i32 {
-    for i in 0..0 {
-    //~^ ERROR: mismatched types [E0308]
+fn foo(n: i32) -> i32 { //~ HELP otherwise consider changing the return type to account for that possibility
+    for i in 0..0 { //~ ERROR mismatched types [E0308]
        if n < 0 {
         return i;
         } else if n < 10 {
@@ -15,8 +14,7 @@ fn foo(n: i32) -> i32 {
           return 5;
         }
 
-    }
-    //~| help: return a value for the case when the loop has zero elements to iterate on, or consider changing the return type to account for that possibility
+    } //~ HELP return a value for the case when the loop has zero elements to iterate on
 }
 
 fn main() {}

tests/ui/typeck/issue-100285.stderr

@@ -4,9 +4,9 @@ error[E0308]: mismatched types
 LL |   fn foo(n: i32) -> i32 {
    |                     --- expected `i32` because of return type
 LL | /     for i in 0..0 {
-LL | |
 LL | |        if n < 0 {
 LL | |         return i;
+LL | |         } else if n < 10 {
 ...  |
 LL | |
 LL | |     }
@@ -17,7 +17,7 @@ note: the function expects a value to always be returned, but loops might run ze
    |
 LL |     for i in 0..0 {
    |     ^^^^^^^^^^^^^ this might have zero elements to iterate on
-...
+LL |        if n < 0 {
 LL |         return i;
    |         -------- if the loop doesn't execute, this value would never get returned
 LL |         } else if n < 10 {
@@ -27,7 +27,32 @@ LL |         } else if n < 20 {
 LL |           return 2;
    |           -------- if the loop doesn't execute, this value would never get returned
    = note: if the loop doesn't execute, 3 other values would never get returned
-   = help: return a value for the case when the loop has zero elements to iterate on, or consider changing the return type to account for that possibility
+help: return a value for the case when the loop has zero elements to iterate on
+   |
+LL ~     }
+LL ~     /* `i32` value */
+   |
+help: otherwise consider changing the return type to account for that possibility
+   |
+LL ~ fn foo(n: i32) -> Option<i32> {
+LL |     for i in 0..0 {
+LL |        if n < 0 {
+LL ~         return Some(i);
+LL |         } else if n < 10 {
+LL ~           return Some(1);
+LL |         } else if n < 20 {
+LL ~           return Some(2);
+LL |         } else if n < 30 {
+LL ~           return Some(3);
+LL |         } else if n < 40 {
+LL ~           return Some(4);
+LL |         } else {
+LL ~           return Some(5);
+LL |         }
+LL | 
+LL ~     }
+LL ~     None
+   |
 
 error: aborting due to previous error
 

tests/ui/typeck/issue-98982.rs

@@ -1,9 +1,7 @@
-fn foo() -> i32 {
-    for i in 0..0 {
-    //~^ ERROR: mismatched types [E0308]
+fn foo() -> i32 { //~ HELP otherwise consider changing the return type to account for that possibility
+    for i in 0..0 { //~ ERROR mismatched types [E0308]
         return i;
-    }
-    //~| help: return a value for the case when the loop has zero elements to iterate on, or consider changing the return type to account for that possibility
+    } //~ HELP return a value for the case when the loop has zero elements to iterate on
 }
 
 fn main() {}

tests/ui/typeck/issue-98982.stderr

@@ -4,7 +4,6 @@ error[E0308]: mismatched types
 LL |   fn foo() -> i32 {
    |               --- expected `i32` because of return type
 LL | /     for i in 0..0 {
-LL | |
 LL | |         return i;
 LL | |     }
    | |_____^ expected `i32`, found `()`
@@ -14,10 +13,21 @@ note: the function expects a value to always be returned, but loops might run ze
    |
 LL |     for i in 0..0 {
    |     ^^^^^^^^^^^^^ this might have zero elements to iterate on
-LL |
 LL |         return i;
    |         -------- if the loop doesn't execute, this value would never get returned
-   = help: return a value for the case when the loop has zero elements to iterate on, or consider changing the return type to account for that possibility
+help: return a value for the case when the loop has zero elements to iterate on
+   |
+LL ~     }
+LL ~     /* `i32` value */
+   |
+help: otherwise consider changing the return type to account for that possibility
+   |
+LL ~ fn foo() -> Option<i32> {
+LL |     for i in 0..0 {
+LL ~         return Some(i);
+LL ~     }
+LL ~     None
+   |
 
 error: aborting due to previous error