Compatible variants suggestion with desugaring

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) {

src/test/ui/async-await/proper-span-for-type-error.fixed

@@ -5,7 +5,8 @@
 async fn a() {}
 
 async fn foo() -> Result<(), i32> {
-    Ok(a().await) //~ ERROR mismatched types
+    a().await;
+    Ok(()) //~ ERROR mismatched types
 }
 
 fn main() {}

src/test/ui/async-await/proper-span-for-type-error.stderr

@@ -6,10 +6,11 @@ LL |     a().await
    |
    = note:   expected enum `Result<(), i32>`
            found unit type `()`
-help: try wrapping the expression in `Ok`
+help: try adding an expression at the end of the block
+   |
+LL ~     a().await;
+LL ~     Ok(())
    |
-LL |     Ok(a().await)
-   |     +++         +
 
 error: aborting due to previous error
 

src/test/ui/did_you_mean/compatible-variants.rs

@@ -23,6 +23,21 @@ fn b() -> Result<(), ()> {
     //~| HELP try adding an expression
 }
 
+fn c() -> Option<()> {
+    for _ in [1, 2] {
+        //~^ ERROR mismatched types
+        f();
+    }
+    //~^ HELP try adding an expression
+}
+
+fn d() -> Option<()> {
+    c()?
+    //~^ ERROR incompatible types
+    //~| HELP try removing this `?`
+    //~| HELP try adding an expression
+}
+
 fn main() {
     let _: Option<()> = while false {};
     //~^ ERROR mismatched types

src/test/ui/did_you_mean/compatible-variants.stderr

@@ -37,7 +37,52 @@ LL +     Ok(())
    |
 
 error[E0308]: mismatched types
-  --> $DIR/compatible-variants.rs:27:25
+  --> $DIR/compatible-variants.rs:27:5
+   |
+LL |   fn c() -> Option<()> {
+   |             ---------- expected `Option<()>` because of return type
+LL | /     for _ in [1, 2] {
+LL | |
+LL | |         f();
+LL | |     }
+   | |_____^ expected enum `Option`, found `()`
+   |
+   = note:   expected enum `Option<()>`
+           found unit type `()`
+help: try adding an expression at the end of the block
+   |
+LL ~     }
+LL +     None
+   |
+LL ~     }
+LL +     Some(())
+   |
+
+error[E0308]: `?` operator has incompatible types
+  --> $DIR/compatible-variants.rs:35:5
+   |
+LL |     c()?
+   |     ^^^^ expected enum `Option`, found `()`
+   |
+   = note: `?` operator cannot convert from `()` to `Option<()>`
+   = note:   expected enum `Option<()>`
+           found unit type `()`
+help: try removing this `?`
+   |
+LL -     c()?
+LL +     c()
+   | 
+help: try adding an expression at the end of the block
+   |
+LL ~     c()?;
+LL +     None
+   |
+LL ~     c()?;
+LL +     Some(())
+   |
+
+error[E0308]: mismatched types
+  --> $DIR/compatible-variants.rs:42:25
    |
 LL |     let _: Option<()> = while false {};
    |            ----------   ^^^^^^^^^^^^^^ expected enum `Option`, found `()`
@@ -52,7 +97,7 @@ LL |     let _: Option<()> = Some(while false {});
    |                         +++++              +
 
 error[E0308]: mismatched types
-  --> $DIR/compatible-variants.rs:31:9
+  --> $DIR/compatible-variants.rs:46:9
    |
 LL |         while false {}
    |         ^^^^^^^^^^^^^^ expected enum `Option`, found `()`
@@ -69,7 +114,7 @@ LL +         Some(())
    |
 
 error[E0308]: mismatched types
-  --> $DIR/compatible-variants.rs:35:31
+  --> $DIR/compatible-variants.rs:50:31
    |
 LL |     let _: Result<i32, i32> = 1;
    |            ----------------   ^ expected enum `Result`, found integer
@@ -86,7 +131,7 @@ LL |     let _: Result<i32, i32> = Err(1);
    |                               ++++ +
 
 error[E0308]: mismatched types
-  --> $DIR/compatible-variants.rs:38:26
+  --> $DIR/compatible-variants.rs:53:26
    |
 LL |     let _: Option<i32> = 1;
    |            -----------   ^ expected enum `Option`, found integer
@@ -101,7 +146,7 @@ LL |     let _: Option<i32> = Some(1);
    |                          +++++ +
 
 error[E0308]: mismatched types
-  --> $DIR/compatible-variants.rs:41:28
+  --> $DIR/compatible-variants.rs:56:28
    |
 LL |     let _: Hey<i32, i32> = 1;
    |            -------------   ^ expected enum `Hey`, found integer
@@ -118,7 +163,7 @@ LL |     let _: Hey<i32, i32> = Hey::B(1);
    |                            +++++++ +
 
 error[E0308]: mismatched types
-  --> $DIR/compatible-variants.rs:44:29
+  --> $DIR/compatible-variants.rs:59:29
    |
 LL |     let _: Hey<i32, bool> = false;
    |            --------------   ^^^^^ expected enum `Hey`, found `bool`
@@ -133,7 +178,7 @@ LL |     let _: Hey<i32, bool> = Hey::B(false);
    |                             +++++++     +
 
 error[E0308]: mismatched types
-  --> $DIR/compatible-variants.rs:48:19
+  --> $DIR/compatible-variants.rs:63:19
    |
 LL |     let _ = Foo { bar };
    |                   ^^^ expected enum `Option`, found `i32`
@@ -145,6 +190,6 @@ help: try wrapping the expression in `Some`
 LL |     let _ = Foo { bar: Some(bar) };
    |                   ++++++++++   +
 
-error: aborting due to 9 previous errors
+error: aborting due to 11 previous errors
 
 For more information about this error, try `rustc --explain E0308`.