Suggest await on cases involving infer

3 files changed, 85 insertions, 4 deletions

diff --git a/compiler/rustc_infer/src/infer/error_reporting/mod.rs b/compiler/rustc_infer/src/infer/error_reporting/mod.rs
index c25ec1356e2..3c2f9900080 100644
--- a/compiler/rustc_infer/src/infer/error_reporting/mod.rs
+++ b/compiler/rustc_infer/src/infer/error_reporting/mod.rs
@@ -310,6 +310,34 @@ pub fn unexpected_hidden_region_diagnostic(
     err
 }
 
+/// Structurally compares two types, modulo any inference variables.
+///
+/// Returns `true` if two types are equal, or if one type is an inference variable compatible
+/// with the other type. A TyVar inference type is compatible with any type, and an IntVar or
+/// FloatVar inference type are compatible with themselves or their concrete types (Int and
+/// Float types, respectively). When comparing two ADTs, these rules apply recursively.
+pub fn same_type_modulo_infer(a: Ty<'tcx>, b: Ty<'ctx>) -> bool {
+    match (&a.kind(), &b.kind()) {
+        (&ty::Adt(did_a, substs_a), &ty::Adt(did_b, substs_b)) => {
+            if did_a != did_b {
+                return false;
+            }
+
+            substs_a.types().zip(substs_b.types()).all(|(a, b)| same_type_modulo_infer(a, b))
+        }
+        (&ty::Int(_), &ty::Infer(ty::InferTy::IntVar(_)))
+        | (&ty::Infer(ty::InferTy::IntVar(_)), &ty::Int(_) | &ty::Infer(ty::InferTy::IntVar(_)))
+        | (&ty::Float(_), &ty::Infer(ty::InferTy::FloatVar(_)))
+        | (
+            &ty::Infer(ty::InferTy::FloatVar(_)),
+            &ty::Float(_) | &ty::Infer(ty::InferTy::FloatVar(_)),
+        )
+        | (&ty::Infer(ty::InferTy::TyVar(_)), _)
+        | (_, &ty::Infer(ty::InferTy::TyVar(_))) => true,
+        _ => a == b,
+    }
+}
+
 impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
     pub fn report_region_errors(&self, errors: &Vec<RegionResolutionError<'tcx>>) {
         debug!("report_region_errors(): {} errors to start", errors.len());
@@ -1761,7 +1789,7 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
             self.get_impl_future_output_ty(exp_found.expected),
             self.get_impl_future_output_ty(exp_found.found),
         ) {
-            (Some(exp), Some(found)) if ty::TyS::same_type(exp, found) => match &cause.code {
+            (Some(exp), Some(found)) if same_type_modulo_infer(exp, found) => match &cause.code {
                 ObligationCauseCode::IfExpression(box IfExpressionCause { then, .. }) => {
                     diag.multipart_suggestion(
                         "consider `await`ing on both `Future`s",
@@ -1793,7 +1821,7 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
                     diag.help("consider `await`ing on both `Future`s");
                 }
             },
-            (_, Some(ty)) if ty::TyS::same_type(exp_found.expected, ty) => {
+            (_, Some(ty)) if same_type_modulo_infer(exp_found.expected, ty) => {
                 diag.span_suggestion_verbose(
                     exp_span.shrink_to_hi(),
                     "consider `await`ing on the `Future`",
@@ -1801,7 +1829,7 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
                     Applicability::MaybeIncorrect,
                 );
             }
-            (Some(ty), _) if ty::TyS::same_type(ty, exp_found.found) => match cause.code {
+            (Some(ty), _) if same_type_modulo_infer(ty, exp_found.found) => match cause.code {
                 ObligationCauseCode::Pattern { span: Some(span), .. }
                 | ObligationCauseCode::IfExpression(box IfExpressionCause { then: span, .. }) => {
                     diag.span_suggestion_verbose(
diff --git a/src/test/ui/async-await/suggest-missing-await.rs b/src/test/ui/async-await/suggest-missing-await.rs
index df74df79d9f..c7c5b51e733 100644
--- a/src/test/ui/async-await/suggest-missing-await.rs
+++ b/src/test/ui/async-await/suggest-missing-await.rs
@@ -54,4 +54,21 @@ async fn suggest_await_on_match_expr() {
     };
 }
 
+async fn dummy_result() -> Result<(), ()> {
+    Ok(())
+}
+
+#[allow(unused)]
+async fn suggest_await_in_generic_pattern() {
+    match dummy_result() {
+        //~^ HELP consider `await`ing on the `Future`
+        //~| HELP consider `await`ing on the `Future`
+        //~| SUGGESTION .await
+        Ok(_) => {}
+        //~^ ERROR mismatched types [E0308]
+        Err(_) => {}
+        //~^ 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 bea50b3bfc0..5757f39af18 100644
--- a/src/test/ui/async-await/suggest-missing-await.stderr
+++ b/src/test/ui/async-await/suggest-missing-await.stderr
@@ -106,6 +106,42 @@ help: consider `await`ing on the `Future`
 LL |     let _x = match dummy().await {
    |                           ++++++
 
-error: aborting due to 5 previous errors
+error[E0308]: mismatched types
+  --> $DIR/suggest-missing-await.rs:67:9
+   |
+LL |         Ok(_) => {}
+   |         ^^^^^ expected opaque type, found enum `Result`
+   |
+note: while checking the return type of the `async fn`
+  --> $DIR/suggest-missing-await.rs:57:28
+   |
+LL | async fn dummy_result() -> Result<(), ()> {
+   |                            ^^^^^^^^^^^^^^ checked the `Output` of this `async fn`, expected opaque type
+   = note: expected opaque type `impl Future`
+                     found enum `Result<_, _>`
+help: consider `await`ing on the `Future`
+   |
+LL |     match dummy_result().await {
+   |                         ++++++
+
+error[E0308]: mismatched types
+  --> $DIR/suggest-missing-await.rs:69:9
+   |
+LL |         Err(_) => {}
+   |         ^^^^^^ expected opaque type, found enum `Result`
+   |
+note: while checking the return type of the `async fn`
+  --> $DIR/suggest-missing-await.rs:57:28
+   |
+LL | async fn dummy_result() -> Result<(), ()> {
+   |                            ^^^^^^^^^^^^^^ checked the `Output` of this `async fn`, expected opaque type
+   = note: expected opaque type `impl Future`
+                     found enum `Result<_, _>`
+help: consider `await`ing on the `Future`
+   |
+LL |     match dummy_result().await {
+   |                         ++++++
+
+error: aborting due to 7 previous errors
 
 For more information about this error, try `rustc --explain E0308`.