Fix Ok(()) suggestion when desugaring is involved.

1 files changed, 19 insertions, 2 deletions

diff --git a/compiler/rustc_typeck/src/check/demand.rs b/compiler/rustc_typeck/src/check/demand.rs
index 3b644099ecf..501ff2e3850 100644
--- a/compiler/rustc_typeck/src/check/demand.rs
+++ b/compiler/rustc_typeck/src/check/demand.rs
@@ -276,11 +276,28 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             // we suggest adding a separate return expression instead.
             // (To avoid things like suggesting `Ok(while .. { .. })`.)
             if expr_ty.is_unit() {
+                let mut id = expr.hir_id;
+                let mut parent;
+
+                // Unroll desugaring, to make sure this works for `for` loops etc.
+                loop {
+                    parent = self.tcx.hir().get_parent_node(id);
+                    if let Some(parent_span) = self.tcx.hir().opt_span(parent) {
+                        if parent_span.find_ancestor_inside(expr.span).is_some() {
+                            // The parent node is part of the same span, so is the result of the
+                            // same expansion/desugaring and not the 'real' parent node.
+                            id = parent;
+                            continue;
+                        }
+                    }
+                    break;
+                }
+
                 if let Some(hir::Node::Block(&hir::Block {
                     span: block_span, expr: Some(e), ..
-                })) = self.tcx.hir().find(self.tcx.hir().get_parent_node(expr.hir_id))
+                })) = self.tcx.hir().find(parent)
                 {
-                    if e.hir_id == expr.hir_id {
+                    if e.hir_id == id {
                         if let Some(span) = expr.span.find_ancestor_inside(block_span) {
                             let return_suggestions =
                                 if self.tcx.is_diagnostic_item(sym::Result, expected_adt.did) {