From 9aed9697cf0c5042b0465b06071b2508e7884d2a Mon Sep 17 00:00:00 2001
From: yukang <moorekang@gmail.com>
Date: Tue, 11 Jul 2023 17:14:20 +0800
Subject: [PATCH] While let suggestion will work for closure

---
 .../src/infer/error_reporting/suggest.rs      | 81 ++++++++++---------
 tests/ui/inference/issue-113354.fixed         |  4 +
 tests/ui/inference/issue-113354.rs            |  4 +
 tests/ui/inference/issue-113354.stderr        | 14 ++++
 4 files changed, 63 insertions(+), 40 deletions(-)
 create mode 100644 tests/ui/inference/issue-113354.fixed
 create mode 100644 tests/ui/inference/issue-113354.rs
 create mode 100644 tests/ui/inference/issue-113354.stderr

diff --git a/compiler/rustc_infer/src/infer/error_reporting/suggest.rs b/compiler/rustc_infer/src/infer/error_reporting/suggest.rs
index 1422bdc9ea2..b0e7cf23cae 100644
--- a/compiler/rustc_infer/src/infer/error_reporting/suggest.rs
+++ b/compiler/rustc_infer/src/infer/error_reporting/suggest.rs
@@ -464,52 +464,53 @@ impl<'tcx> TypeErrCtxt<'_, 'tcx> {
         span: Span,
     ) -> Option<TypeErrorAdditionalDiags> {
         let hir = self.tcx.hir();
-        if let Some(node) = self.tcx.hir().find_by_def_id(cause.body_id) &&
-            let hir::Node::Item(hir::Item {
-                    kind: hir::ItemKind::Fn(_sig, _, body_id), ..
-                }) = node {
-        let body = hir.body(*body_id);
+        if let Some(body_id) = self.tcx.hir().maybe_body_owned_by(cause.body_id) {
+            let body = hir.body(body_id);
 
-        /// Find the if expression with given span
-        struct IfVisitor {
-            pub result: bool,
-            pub found_if: bool,
-            pub err_span: Span,
-        }
+            /// Find the if expression with given span
+            struct IfVisitor {
+                pub result: bool,
+                pub found_if: bool,
+                pub err_span: Span,
+            }
 
-        impl<'v> Visitor<'v> for IfVisitor {
-            fn visit_expr(&mut self, ex: &'v hir::Expr<'v>) {
-                if self.result { return; }
-                match ex.kind {
-                    hir::ExprKind::If(cond, _, _) => {
-                        self.found_if = true;
-                        walk_expr(self, cond);
-                        self.found_if = false;
+            impl<'v> Visitor<'v> for IfVisitor {
+                fn visit_expr(&mut self, ex: &'v hir::Expr<'v>) {
+                    if self.result {
+                        return;
                     }
-                    _ => walk_expr(self, ex),
+                    match ex.kind {
+                        hir::ExprKind::If(cond, _, _) => {
+                            self.found_if = true;
+                            walk_expr(self, cond);
+                            self.found_if = false;
+                        }
+                        _ => walk_expr(self, ex),
+                    }
+                }
+
+                fn visit_stmt(&mut self, ex: &'v hir::Stmt<'v>) {
+                    if let hir::StmtKind::Local(hir::Local {
+                                    span, pat: hir::Pat{..}, ty: None, init: Some(_), ..
+                                }) = &ex.kind
+                                && self.found_if
+                                && span.eq(&self.err_span) {
+                                self.result = true;
+                            }
+                    walk_stmt(self, ex);
+                }
+
+                fn visit_body(&mut self, body: &'v hir::Body<'v>) {
+                    hir::intravisit::walk_body(self, body);
                 }
             }
 
-            fn visit_stmt(&mut self, ex: &'v hir::Stmt<'v>) {
-                if let hir::StmtKind::Local(hir::Local {
-                        span, pat: hir::Pat{..}, ty: None, init: Some(_), ..
-                    }) = &ex.kind
-                    && self.found_if
-                    && span.eq(&self.err_span) {
-                        self.result = true;
-                }
-                walk_stmt(self, ex);
-            }
-
-            fn visit_body(&mut self, body: &'v hir::Body<'v>) {
-                hir::intravisit::walk_body(self, body);
-            }
-        }
-
-        let mut visitor = IfVisitor { err_span: span, found_if: false, result: false };
-        visitor.visit_body(&body);
-        if visitor.result {
-                return Some(TypeErrorAdditionalDiags::AddLetForLetChains{span: span.shrink_to_lo()});
+            let mut visitor = IfVisitor { err_span: span, found_if: false, result: false };
+            visitor.visit_body(&body);
+            if visitor.result {
+                return Some(TypeErrorAdditionalDiags::AddLetForLetChains {
+                    span: span.shrink_to_lo(),
+                });
             }
         }
         None
diff --git a/tests/ui/inference/issue-113354.fixed b/tests/ui/inference/issue-113354.fixed
new file mode 100644
index 00000000000..804db985ae1
--- /dev/null
+++ b/tests/ui/inference/issue-113354.fixed
@@ -0,0 +1,4 @@
+//run-rustfix
+fn main() {
+    let _ = || { while let Some(_) = Some(1) { } }; //~ ERROR mismatched types
+}
diff --git a/tests/ui/inference/issue-113354.rs b/tests/ui/inference/issue-113354.rs
new file mode 100644
index 00000000000..ec33d1f8b84
--- /dev/null
+++ b/tests/ui/inference/issue-113354.rs
@@ -0,0 +1,4 @@
+//run-rustfix
+fn main() {
+    let _ = || { while Some(_) = Some(1) { } }; //~ ERROR mismatched types
+}
diff --git a/tests/ui/inference/issue-113354.stderr b/tests/ui/inference/issue-113354.stderr
new file mode 100644
index 00000000000..045a5aa7bf0
--- /dev/null
+++ b/tests/ui/inference/issue-113354.stderr
@@ -0,0 +1,14 @@
+error[E0308]: mismatched types
+  --> $DIR/issue-113354.rs:3:24
+   |
+LL |     let _ = || { while Some(_) = Some(1) { } };
+   |                        ^^^^^^^^^^^^^^^^^ expected `bool`, found `()`
+   |
+help: consider adding `let`
+   |
+LL |     let _ = || { while let Some(_) = Some(1) { } };
+   |                        +++
+
+error: aborting due to previous error
+
+For more information about this error, try `rustc --explain E0308`.