From ad78b50a86c6ac908df1c385642d46d3183740d6 Mon Sep 17 00:00:00 2001
From: Tor Hovland <tor.hovland@bekk.no>
Date: Sat, 24 Apr 2021 13:57:41 +0200
Subject: [PATCH] Implemented suggestion.

---
 compiler/rustc_typeck/src/check/coercion.rs   |  4 +-
 .../src/check/fn_ctxt/suggestions.rs          | 78 ++++++++++++++++++-
 2 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/compiler/rustc_typeck/src/check/coercion.rs b/compiler/rustc_typeck/src/check/coercion.rs
index 427f967a9b6..236fec94bdb 100644
--- a/compiler/rustc_typeck/src/check/coercion.rs
+++ b/compiler/rustc_typeck/src/check/coercion.rs
@@ -1494,7 +1494,9 @@ impl<'tcx, 'exprs, E: AsCoercionSite> CoerceMany<'tcx, 'exprs, E> {
         if let (Some((expr, _)), Some((fn_decl, _, _))) =
             (expression, fcx.get_node_fn_decl(parent_item))
         {
-            fcx.suggest_missing_return_expr(&mut err, expr, fn_decl, expected, found, parent_id);
+            fcx.suggest_missing_break_or_return_expr(
+                &mut err, expr, fn_decl, expected, found, id, parent_id,
+            );
         }
 
         if let (Some(sp), Some(fn_output)) = (fcx.ret_coercion_span.get(), fn_output) {
diff --git a/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs b/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs
index b7583344845..c3417247725 100644
--- a/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs
+++ b/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs
@@ -8,7 +8,7 @@ use rustc_errors::{Applicability, DiagnosticBuilder};
 use rustc_hir as hir;
 use rustc_hir::def::{CtorOf, DefKind};
 use rustc_hir::lang_items::LangItem;
-use rustc_hir::{ExprKind, ItemKind, Node};
+use rustc_hir::{ExprKind, ItemKind, Node, StmtKind};
 use rustc_infer::infer;
 use rustc_middle::lint::in_external_macro;
 use rustc_middle::ty::{self, Binder, Ty};
@@ -55,7 +55,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             pointing_at_return_type =
                 self.suggest_missing_return_type(err, &fn_decl, expected, found, can_suggest);
             let fn_id = self.tcx.hir().get_return_block(blk_id).unwrap();
-            self.suggest_missing_return_expr(err, expr, &fn_decl, expected, found, fn_id);
+            self.suggest_missing_break_or_return_expr(
+                err, expr, &fn_decl, expected, found, blk_id, fn_id,
+            );
         }
         pointing_at_return_type
     }
@@ -472,7 +474,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         }
     }
 
-    pub(in super::super) fn suggest_missing_return_expr(
+    pub(in super::super) fn suggest_missing_break_or_return_expr(
         &self,
         err: &mut DiagnosticBuilder<'_>,
         expr: &'tcx hir::Expr<'tcx>,
@@ -480,14 +482,30 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         expected: Ty<'tcx>,
         found: Ty<'tcx>,
         id: hir::HirId,
+        fn_id: hir::HirId,
     ) {
         if !expected.is_unit() {
             return;
         }
         let found = self.resolve_vars_with_obligations(found);
+
+        if self.in_loop(id) {
+            if self.in_local_statement(id) {
+                err.multipart_suggestion(
+                    "you might have meant to break the loop with this value",
+                    vec![
+                        (expr.span.shrink_to_lo(), "break ".to_string()),
+                        (expr.span.shrink_to_hi(), ";".to_string()),
+                    ],
+                    Applicability::MaybeIncorrect,
+                );
+                return;
+            }
+        }
+
         if let hir::FnRetTy::Return(ty) = fn_decl.output {
             let ty = <dyn AstConv<'_>>::ast_ty_to_ty(self, ty);
-            let bound_vars = self.tcx.late_bound_vars(id);
+            let bound_vars = self.tcx.late_bound_vars(fn_id);
             let ty = self.tcx.erase_late_bound_regions(Binder::bind_with_vars(ty, bound_vars));
             let ty = self.normalize_associated_types_in(expr.span, ty);
             if self.can_coerce(found, ty) {
@@ -514,4 +532,56 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             self.tcx.sess.parse_sess.expr_parentheses_needed(err, *sp, None);
         }
     }
+
+    fn in_loop(&self, id: hir::HirId) -> bool {
+        if self.is_loop(id) {
+            return true;
+        }
+
+        for (parent_id, _) in self.tcx.hir().parent_iter(id) {
+            if self.is_loop(parent_id) {
+                return true;
+            }
+        }
+
+        false
+    }
+
+    fn is_loop(&self, id: hir::HirId) -> bool {
+        let node = self.tcx.hir().get(id);
+
+        if let Node::Expr(expr) = node {
+            if let ExprKind::Loop(..) = expr.kind {
+                return true;
+            }
+        }
+
+        false
+    }
+
+    fn in_local_statement(&self, id: hir::HirId) -> bool {
+        if self.is_local_statement(id) {
+            return true;
+        }
+
+        for (parent_id, _) in self.tcx.hir().parent_iter(id) {
+            if self.is_local_statement(parent_id) {
+                return true;
+            }
+        }
+
+        false
+    }
+
+    fn is_local_statement(&self, id: hir::HirId) -> bool {
+        let node = self.tcx.hir().get(id);
+
+        if let Node::Stmt(stmt) = node {
+            if let StmtKind::Local(..) = stmt.kind {
+                return true;
+            }
+        }
+
+        false
+    }
 }