Fix await suggestion better

3 files changed, 123 insertions, 20 deletions

diff --git a/compiler/rustc_infer/src/infer/error_reporting/mod.rs b/compiler/rustc_infer/src/infer/error_reporting/mod.rs
index 7a0db820c9c..c25ec1356e2 100644
--- a/compiler/rustc_infer/src/infer/error_reporting/mod.rs
+++ b/compiler/rustc_infer/src/infer/error_reporting/mod.rs
@@ -1794,31 +1794,38 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
                 }
             },
             (_, Some(ty)) if ty::TyS::same_type(exp_found.expected, ty) => {
-                let span = match cause.code {
-                    // scrutinee's span
-                    ObligationCauseCode::Pattern { span: Some(span), .. } => span,
-                    _ => exp_span,
-                };
-                diag.span_suggestion_verbose(
-                    span.shrink_to_hi(),
-                    "consider `await`ing on the `Future`",
-                    ".await".to_string(),
-                    Applicability::MaybeIncorrect,
-                );
-            }
-            (Some(ty), _) if ty::TyS::same_type(ty, exp_found.found) => {
-                let span = match cause.code {
-                    // scrutinee's span
-                    ObligationCauseCode::Pattern { span: Some(span), .. } => span,
-                    _ => exp_span,
-                };
                 diag.span_suggestion_verbose(
-                    span.shrink_to_hi(),
+                    exp_span.shrink_to_hi(),
                     "consider `await`ing on the `Future`",
                     ".await".to_string(),
                     Applicability::MaybeIncorrect,
                 );
             }
+            (Some(ty), _) if ty::TyS::same_type(ty, exp_found.found) => match cause.code {
+                ObligationCauseCode::Pattern { span: Some(span), .. }
+                | ObligationCauseCode::IfExpression(box IfExpressionCause { then: span, .. }) => {
+                    diag.span_suggestion_verbose(
+                        span.shrink_to_hi(),
+                        "consider `await`ing on the `Future`",
+                        ".await".to_string(),
+                        Applicability::MaybeIncorrect,
+                    );
+                }
+                ObligationCauseCode::MatchExpressionArm(box MatchExpressionArmCause {
+                    ref prior_arms,
+                    ..
+                }) => {
+                    diag.multipart_suggestion_verbose(
+                        "consider `await`ing on the `Future`",
+                        prior_arms
+                            .iter()
+                            .map(|arm| (arm.shrink_to_hi(), ".await".to_string()))
+                            .collect(),
+                        Applicability::MaybeIncorrect,
+                    );
+                }
+                _ => {}
+            },
             _ => {}
         }
     }
diff --git a/src/test/ui/async-await/suggest-missing-await.rs b/src/test/ui/async-await/suggest-missing-await.rs
index 352a88ac10c..df74df79d9f 100644
--- a/src/test/ui/async-await/suggest-missing-await.rs
+++ b/src/test/ui/async-await/suggest-missing-await.rs
@@ -26,4 +26,32 @@ async fn suggest_await_in_async_fn_return() {
     //~| SUGGESTION .await
 }
 
+#[allow(unused)]
+async fn suggest_await_on_if() {
+    let _x = if true {
+        dummy()
+        //~^ HELP consider `await`ing on the `Future`
+    } else {
+        dummy().await
+        //~^ ERROR `if` and `else` have incompatible types [E0308]
+    };
+}
+
+#[allow(unused)]
+async fn suggest_await_on_previous_match_arms() {
+    let _x = match 0usize {
+        0 => dummy(), //~ HELP consider `await`ing on the `Future`
+        1 => dummy(),
+        2 => dummy().await,
+        //~^ `match` arms have incompatible types [E0308]
+    };
+}
+
+#[allow(unused)]
+async fn suggest_await_on_match_expr() {
+    let _x = match dummy() { //~ HELP consider `await`ing on the `Future`
+        () => {} //~ ERROR mismatched types [E0308]
+    };
+}
+
 fn main() {}
diff --git a/src/test/ui/async-await/suggest-missing-await.stderr b/src/test/ui/async-await/suggest-missing-await.stderr
index 08868a04657..aefe3096fd9 100644
--- a/src/test/ui/async-await/suggest-missing-await.stderr
+++ b/src/test/ui/async-await/suggest-missing-await.stderr
@@ -38,6 +38,74 @@ help: consider using a semicolon here
 LL |     dummy();
    |            +
 
-error: aborting due to 2 previous errors
+error[E0308]: `if` and `else` have incompatible types
+  --> $DIR/suggest-missing-await.rs:35:9
+   |
+LL |       let _x = if true {
+   |  ______________-
+LL | |         dummy()
+   | |         ------- expected because of this
+LL | |
+LL | |     } else {
+LL | |         dummy().await
+   | |         ^^^^^^^^^^^^^ expected opaque type, found `()`
+LL | |
+LL | |     };
+   | |_____- `if` and `else` have incompatible types
+   |
+   = note:   expected type `impl Future`
+           found unit type `()`
+help: consider `await`ing on the `Future`
+   |
+LL |         dummy().await
+   |                ++++++
+
+error[E0308]: `match` arms have incompatible types
+  --> $DIR/suggest-missing-await.rs:45:14
+   |
+LL |       let _x = match 0usize {
+   |  ______________-
+LL | |         0 => dummy(),
+   | |              ------- this is found to be of type `impl Future`
+LL | |         1 => dummy(),
+   | |              ------- this is found to be of type `impl Future`
+LL | |         2 => dummy().await,
+   | |              ^^^^^^^^^^^^^ expected opaque type, found `()`
+LL | |
+LL | |     };
+   | |_____- `match` arms have incompatible types
+   |
+note: while checking the return type of the `async fn`
+  --> $DIR/suggest-missing-await.rs:18:18
+   |
+LL | async fn dummy() {}
+   |                  ^ checked the `Output` of this `async fn`, expected opaque type
+   = note: expected opaque type `impl Future`
+                found unit type `()`
+help: consider `await`ing on the `Future`
+   |
+LL ~         0 => dummy().await,
+LL ~         1 => dummy().await,
+   |
+
+error[E0308]: mismatched types
+  --> $DIR/suggest-missing-await.rs:53:9
+   |
+LL |         () => {}
+   |         ^^ expected opaque type, found `()`
+   |
+note: while checking the return type of the `async fn`
+  --> $DIR/suggest-missing-await.rs:18:18
+   |
+LL | async fn dummy() {}
+   |                  ^ checked the `Output` of this `async fn`, expected opaque type
+   = note: expected opaque type `impl Future`
+                found unit type `()`
+help: consider `await`ing on the `Future`
+   |
+LL |     let _x = match dummy().await {
+   |                           ++++++
+
+error: aborting due to 5 previous errors
 
 For more information about this error, try `rustc --explain E0308`.