Add more `.await` suggestions on E0308

12 files changed, 167 insertions, 188 deletions

diff --git a/compiler/rustc_infer/src/infer/error_reporting/mod.rs b/compiler/rustc_infer/src/infer/error_reporting/mod.rs
index 9e93cfc27d0..fcf21c61d8f 100644
--- a/compiler/rustc_infer/src/infer/error_reporting/mod.rs
+++ b/compiler/rustc_infer/src/infer/error_reporting/mod.rs
@@ -50,7 +50,6 @@ use super::region_constraints::GenericKind;
 use super::{InferCtxt, RegionVariableOrigin, SubregionOrigin, TypeTrace, ValuePairs};
 
 use crate::infer;
-use crate::infer::OriginalQueryValues;
 use crate::traits::error_reporting::report_object_safety_error;
 use crate::traits::{
     IfExpressionCause, MatchExpressionArmCause, ObligationCause, ObligationCauseCode,
@@ -65,7 +64,6 @@ use rustc_hir::def_id::DefId;
 use rustc_hir::lang_items::LangItem;
 use rustc_hir::{Item, ItemKind, Node};
 use rustc_middle::ty::error::TypeError;
-use rustc_middle::ty::ParamEnvAnd;
 use rustc_middle::ty::{
     self,
     subst::{Subst, SubstsRef},
@@ -1621,6 +1619,16 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
             Mismatch::Variable(exp_found) => Some(exp_found),
             Mismatch::Fixed(_) => None,
         };
+        let exp_found = match terr {
+            // `terr` has more accurate type information than `exp_found` in match expressions.
+            ty::error::TypeError::Sorts(terr)
+                if exp_found.map_or(false, |ef| terr.found == ef.found) =>
+            {
+                Some(*terr)
+            }
+            _ => exp_found,
+        };
+        debug!("exp_found {:?} terr {:?}", exp_found, terr);
         if let Some(exp_found) = exp_found {
             self.suggest_as_ref_where_appropriate(span, &exp_found, diag);
             self.suggest_await_on_expect_found(cause, span, &exp_found, diag);
@@ -1642,6 +1650,53 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
         self.note_error_origin(diag, cause, exp_found);
     }
 
+    fn get_impl_future_output_ty(&self, ty: Ty<'tcx>) -> Option<Ty<'tcx>> {
+        if let ty::Opaque(def_id, substs) = ty.kind() {
+            let future_trait = self.tcx.require_lang_item(LangItem::Future, None);
+            // Future::Output
+            let item_def_id = self
+                .tcx
+                .associated_items(future_trait)
+                .in_definition_order()
+                .next()
+                .unwrap()
+                .def_id;
+
+            let bounds = self.tcx.explicit_item_bounds(*def_id);
+
+            for (predicate, _) in bounds {
+                let predicate = predicate.subst(self.tcx, substs);
+                if let ty::PredicateAtom::Projection(projection_predicate) =
+                    predicate.skip_binders()
+                {
+                    if projection_predicate.projection_ty.item_def_id == item_def_id {
+                        // We don't account for multiple `Future::Output = Ty` contraints.
+                        return Some(projection_predicate.ty);
+                    }
+                }
+            }
+        }
+        None
+    }
+
+    /// A possible error is to forget to add `.await` when using futures:
+    ///
+    /// ```
+    /// async fn make_u32() -> u32 {
+    ///     22
+    /// }
+    ///
+    /// fn take_u32(x: u32) {}
+    ///
+    /// async fn foo() {
+    ///     let x = make_u32();
+    ///     take_u32(x);
+    /// }
+    /// ```
+    ///
+    /// This routine checks if the found type `T` implements `Future<Output=U>` where `U` is the
+    /// expected type. If this is the case, and we are inside of an async body, it suggests adding
+    /// `.await` to the tail of the expression.
     fn suggest_await_on_expect_found(
         &self,
         cause: &ObligationCause<'tcx>,
@@ -1651,50 +1706,76 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
     ) {
         debug!(
             "suggest_await_on_expect_found: exp_span={:?}, expected_ty={:?}, found_ty={:?}",
-            exp_span, exp_found.expected, exp_found.found
+            exp_span, exp_found.expected, exp_found.found,
         );
 
-        if let ty::Opaque(def_id, _) = *exp_found.expected.kind() {
-            let future_trait = self.tcx.require_lang_item(LangItem::Future, None);
-            // Future::Output
-            let item_def_id = self
-                .tcx
-                .associated_items(future_trait)
-                .in_definition_order()
-                .next()
-                .unwrap()
-                .def_id;
+        if let ObligationCauseCode::CompareImplMethodObligation { .. } = &cause.code {
+            return;
+        }
 
-            let projection_ty = self.tcx.projection_ty_from_predicates((def_id, item_def_id));
-            if let Some(projection_ty) = projection_ty {
-                let projection_query = self.canonicalize_query(
-                    &ParamEnvAnd { param_env: self.tcx.param_env(def_id), value: projection_ty },
-                    &mut OriginalQueryValues::default(),
-                );
-                if let Ok(resp) = self.tcx.normalize_projection_ty(projection_query) {
-                    let normalized_ty = resp.value.value.normalized_ty;
-                    debug!("suggest_await_on_expect_found: normalized={:?}", normalized_ty);
-                    if ty::TyS::same_type(normalized_ty, exp_found.found) {
-                        let span = if let ObligationCauseCode::Pattern {
-                            span,
-                            origin_expr: _,
-                            root_ty: _,
-                        } = cause.code
-                        {
-                            // scrutinee's span
-                            span.unwrap_or(exp_span)
-                        } else {
-                            exp_span
-                        };
-                        diag.span_suggestion_verbose(
-                            span.shrink_to_hi(),
-                            "consider awaiting on the future",
-                            ".await".to_string(),
+        match (
+            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 {
+                ObligationCauseCode::IfExpression(box IfExpressionCause { then, .. }) => {
+                    diag.multipart_suggestion(
+                        "consider `await`ing on both `Future`s",
+                        vec![
+                            (then.shrink_to_hi(), ".await".to_string()),
+                            (exp_span.shrink_to_hi(), ".await".to_string()),
+                        ],
+                        Applicability::MaybeIncorrect,
+                    );
+                }
+                ObligationCauseCode::MatchExpressionArm(box MatchExpressionArmCause {
+                    prior_arms,
+                    ..
+                }) => {
+                    if let [.., arm_span] = &prior_arms[..] {
+                        diag.multipart_suggestion(
+                            "consider `await`ing on both `Future`s",
+                            vec![
+                                (arm_span.shrink_to_hi(), ".await".to_string()),
+                                (exp_span.shrink_to_hi(), ".await".to_string()),
+                            ],
                             Applicability::MaybeIncorrect,
                         );
+                    } else {
+                        diag.help("consider `await`ing on both `Future`s");
                     }
                 }
+                _ => {
+                    diag.help("consider `await`ing on both `Future`s");
+                }
+            },
+            (_, 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(),
+                    "consider `await`ing on the `Future`",
+                    ".await".to_string(),
+                    Applicability::MaybeIncorrect,
+                );
             }
+            _ => {}
         }
     }
 
diff --git a/compiler/rustc_middle/src/ty/error.rs b/compiler/rustc_middle/src/ty/error.rs
index 235f8749cf9..5ec0ec0c56a 100644
--- a/compiler/rustc_middle/src/ty/error.rs
+++ b/compiler/rustc_middle/src/ty/error.rs
@@ -334,26 +334,15 @@ impl<'tcx> TyCtxt<'tcx> {
         debug!("note_and_explain_type_err err={:?} cause={:?}", err, cause);
         match err {
             Sorts(values) => {
-                let expected_str = values.expected.sort_string(self);
-                let found_str = values.found.sort_string(self);
-                if expected_str == found_str && expected_str == "closure" {
-                    db.note("no two closures, even if identical, have the same type");
-                    db.help("consider boxing your closure and/or using it as a trait object");
-                }
-                if expected_str == found_str && expected_str == "opaque type" {
-                    // Issue #63167
-                    db.note("distinct uses of `impl Trait` result in different opaque types");
-                    let e_str = values.expected.to_string();
-                    let f_str = values.found.to_string();
-                    if e_str == f_str && &e_str == "impl std::future::Future" {
-                        // FIXME: use non-string based check.
-                        db.help(
-                            "if both `Future`s have the same `Output` type, consider \
-                                 `.await`ing on both of them",
-                        );
-                    }
-                }
                 match (values.expected.kind(), values.found.kind()) {
+                    (ty::Closure(..), ty::Closure(..)) => {
+                        db.note("no two closures, even if identical, have the same type");
+                        db.help("consider boxing your closure and/or using it as a trait object");
+                    }
+                    (ty::Opaque(..), ty::Opaque(..)) => {
+                        // Issue #63167
+                        db.note("distinct uses of `impl Trait` result in different opaque types");
+                    }
                     (ty::Float(_), ty::Infer(ty::IntVar(_))) => {
                         if let Ok(
                             // Issue #53280
@@ -382,12 +371,12 @@ impl<'tcx> TyCtxt<'tcx> {
                         }
                         db.note(
                             "a type parameter was expected, but a different one was found; \
-                                 you might be missing a type parameter or trait bound",
+                             you might be missing a type parameter or trait bound",
                         );
                         db.note(
                             "for more information, visit \
-                                 https://doc.rust-lang.org/book/ch10-02-traits.html\
-                                 #traits-as-parameters",
+                             https://doc.rust-lang.org/book/ch10-02-traits.html\
+                             #traits-as-parameters",
                         );
                     }
                     (ty::Projection(_), ty::Projection(_)) => {
@@ -471,8 +460,8 @@ impl<T> Trait<T> for X {
                         }
                         db.note(
                             "for more information, visit \
-                                 https://doc.rust-lang.org/book/ch10-02-traits.html\
-                                 #traits-as-parameters",
+                             https://doc.rust-lang.org/book/ch10-02-traits.html\
+                             #traits-as-parameters",
                         );
                     }
                     (ty::Param(p), ty::Closure(..) | ty::Generator(..)) => {
diff --git a/compiler/rustc_typeck/src/check/demand.rs b/compiler/rustc_typeck/src/check/demand.rs
index b8143787a2d..241803fab1e 100644
--- a/compiler/rustc_typeck/src/check/demand.rs
+++ b/compiler/rustc_typeck/src/check/demand.rs
@@ -33,7 +33,6 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             return;
         }
         self.suggest_boxing_when_appropriate(err, expr, expected, expr_ty);
-        self.suggest_missing_await(err, expr, expected, expr_ty);
         self.suggest_missing_parentheses(err, expr);
         self.note_need_for_fn_pointer(err, expected, expr_ty);
         self.note_internal_mutation_in_method(err, expr, expected, expr_ty);
diff --git a/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs b/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs
index 18db8d63850..a8ad9f4fdf8 100644
--- a/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs
+++ b/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs
@@ -3,7 +3,6 @@ use crate::astconv::AstConv;
 
 use rustc_ast::util::parser::ExprPrecedence;
 use rustc_span::{self, Span};
-use rustc_trait_selection::traits;
 
 use rustc_errors::{Applicability, DiagnosticBuilder};
 use rustc_hir as hir;
@@ -13,7 +12,6 @@ use rustc_hir::{ExprKind, ItemKind, Node};
 use rustc_infer::infer;
 use rustc_middle::ty::{self, Ty};
 use rustc_span::symbol::kw;
-use rustc_trait_selection::traits::query::evaluate_obligation::InferCtxtExt as _;
 
 use std::iter;
 
@@ -433,83 +431,6 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         }
     }
 
-    /// A possible error is to forget to add `.await` when using futures:
-    ///
-    /// ```
-    /// async fn make_u32() -> u32 {
-    ///     22
-    /// }
-    ///
-    /// fn take_u32(x: u32) {}
-    ///
-    /// async fn foo() {
-    ///     let x = make_u32();
-    ///     take_u32(x);
-    /// }
-    /// ```
-    ///
-    /// This routine checks if the found type `T` implements `Future<Output=U>` where `U` is the
-    /// expected type. If this is the case, and we are inside of an async body, it suggests adding
-    /// `.await` to the tail of the expression.
-    pub(in super::super) fn suggest_missing_await(
-        &self,
-        err: &mut DiagnosticBuilder<'_>,
-        expr: &hir::Expr<'_>,
-        expected: Ty<'tcx>,
-        found: Ty<'tcx>,
-    ) {
-        debug!("suggest_missing_await: expr={:?} expected={:?}, found={:?}", expr, expected, found);
-        // `.await` is not permitted outside of `async` bodies, so don't bother to suggest if the
-        // body isn't `async`.
-        let item_id = self.tcx().hir().get_parent_node(self.body_id);
-        if let Some(body_id) = self.tcx().hir().maybe_body_owned_by(item_id) {
-            let body = self.tcx().hir().body(body_id);
-            if let Some(hir::GeneratorKind::Async(_)) = body.generator_kind {
-                let sp = expr.span;
-                // Check for `Future` implementations by constructing a predicate to
-                // prove: `<T as Future>::Output == U`
-                let future_trait = self.tcx.require_lang_item(LangItem::Future, Some(sp));
-                let item_def_id = self
-                    .tcx
-                    .associated_items(future_trait)
-                    .in_definition_order()
-                    .next()
-                    .unwrap()
-                    .def_id;
-                // `<T as Future>::Output`
-                let projection_ty = ty::ProjectionTy {
-                    // `T`
-                    substs: self
-                        .tcx
-                        .mk_substs_trait(found, self.fresh_substs_for_item(sp, item_def_id)),
-                    // `Future::Output`
-                    item_def_id,
-                };
-
-                let predicate = ty::PredicateAtom::Projection(ty::ProjectionPredicate {
-                    projection_ty,
-                    ty: expected,
-                })
-                .potentially_quantified(self.tcx, ty::PredicateKind::ForAll);
-                let obligation = traits::Obligation::new(self.misc(sp), self.param_env, predicate);
-
-                debug!("suggest_missing_await: trying obligation {:?}", obligation);
-
-                if self.infcx.predicate_may_hold(&obligation) {
-                    debug!("suggest_missing_await: obligation held: {:?}", obligation);
-                    err.span_suggestion_verbose(
-                        sp.shrink_to_hi(),
-                        "consider `await`ing on the `Future`",
-                        ".await".to_string(),
-                        Applicability::MaybeIncorrect,
-                    );
-                } else {
-                    debug!("suggest_missing_await: obligation did not hold: {:?}", obligation)
-                }
-            }
-        }
-    }
-
     pub(in super::super) fn suggest_missing_parentheses(
         &self,
         err: &mut DiagnosticBuilder<'_>,
diff --git a/src/test/ui/async-await/dont-suggest-missing-await.stderr b/src/test/ui/async-await/dont-suggest-missing-await.stderr
index e70ed9badbd..14e72c2b1e7 100644
--- a/src/test/ui/async-await/dont-suggest-missing-await.stderr
+++ b/src/test/ui/async-await/dont-suggest-missing-await.stderr
@@ -9,6 +9,10 @@ LL |         take_u32(x)
    |
    = note:     expected type `u32`
            found opaque type `impl Future`
+help: consider `await`ing on the `Future`
+   |
+LL |         take_u32(x.await)
+   |                   ^^^^^^
 
 error: aborting due to previous error
 
diff --git a/src/test/ui/async-await/issue-61076.stderr b/src/test/ui/async-await/issue-61076.stderr
index afba889f014..df54ac88ace 100644
--- a/src/test/ui/async-await/issue-61076.stderr
+++ b/src/test/ui/async-await/issue-61076.stderr
@@ -53,6 +53,10 @@ LL |         Tuple(_) => {}
    |
    = note: expected opaque type `impl Future`
                    found struct `Tuple`
+help: consider `await`ing on the `Future`
+   |
+LL |     match tuple().await {
+   |                  ^^^^^^
 
 error: aborting due to 6 previous errors
 
diff --git a/src/test/ui/async-await/suggest-missing-await.fixed b/src/test/ui/async-await/suggest-missing-await.fixed
deleted file mode 100644
index e548fda7cb4..00000000000
--- a/src/test/ui/async-await/suggest-missing-await.fixed
+++ /dev/null
@@ -1,30 +0,0 @@
-// edition:2018
-// run-rustfix
-
-fn take_u32(_x: u32) {}
-
-async fn make_u32() -> u32 {
-    22
-}
-
-#[allow(unused)]
-async fn suggest_await_in_async_fn() {
-    let x = make_u32();
-    take_u32(x.await)
-    //~^ ERROR mismatched types [E0308]
-    //~| HELP consider `await`ing on the `Future`
-    //~| SUGGESTION .await
-}
-
-async fn dummy() {}
-
-#[allow(unused)]
-async fn suggest_await_in_async_fn_return() {
-    dummy().await;
-    //~^ ERROR mismatched types [E0308]
-    //~| HELP try adding a semicolon
-    //~| HELP consider `await`ing on the `Future`
-    //~| SUGGESTION .await
-}
-
-fn main() {}
diff --git a/src/test/ui/async-await/suggest-missing-await.rs b/src/test/ui/async-await/suggest-missing-await.rs
index 464a9602119..d629054911d 100644
--- a/src/test/ui/async-await/suggest-missing-await.rs
+++ b/src/test/ui/async-await/suggest-missing-await.rs
@@ -1,5 +1,4 @@
 // edition:2018
-// run-rustfix
 
 fn take_u32(_x: u32) {}
 
diff --git a/src/test/ui/async-await/suggest-missing-await.stderr b/src/test/ui/async-await/suggest-missing-await.stderr
index d729420e930..46615dae7e2 100644
--- a/src/test/ui/async-await/suggest-missing-await.stderr
+++ b/src/test/ui/async-await/suggest-missing-await.stderr
@@ -1,5 +1,5 @@
 error[E0308]: mismatched types
-  --> $DIR/suggest-missing-await.rs:13:14
+  --> $DIR/suggest-missing-await.rs:12:14
    |
 LL | async fn make_u32() -> u32 {
    |                        --- the `Output` of this `async fn`'s found opaque type
@@ -15,7 +15,7 @@ LL |     take_u32(x.await)
    |               ^^^^^^
 
 error[E0308]: mismatched types
-  --> $DIR/suggest-missing-await.rs:23:5
+  --> $DIR/suggest-missing-await.rs:22:5
    |
 LL | async fn dummy() {}
    |                  - the `Output` of this `async fn`'s found opaque type
@@ -25,14 +25,14 @@ LL |     dummy()
    |
    = note: expected unit type `()`
             found opaque type `impl Future`
-help: try adding a semicolon
-   |
-LL |     dummy();
-   |            ^
 help: consider `await`ing on the `Future`
    |
 LL |     dummy().await
    |            ^^^^^^
+help: try adding a semicolon
+   |
+LL |     dummy();
+   |            ^
 
 error: aborting due to 2 previous errors
 
diff --git a/src/test/ui/suggestions/match-prev-arm-needing-semi.rs b/src/test/ui/suggestions/match-prev-arm-needing-semi.rs
index 81e6abf7d7b..b8ac030b0bb 100644
--- a/src/test/ui/suggestions/match-prev-arm-needing-semi.rs
+++ b/src/test/ui/suggestions/match-prev-arm-needing-semi.rs
@@ -46,6 +46,7 @@ async fn async_extra_semicolon_different() {
 async fn async_different_futures() {
     let _ = match true { //~ NOTE `match` arms have incompatible types
         true => async_dummy(), //~ NOTE this is found to be
+        //~| HELP consider `await`ing on both `Future`s
         false => async_dummy2(), //~ ERROR `match` arms have incompatible types
         //~^ NOTE expected opaque type, found a different opaque type
         //~| NOTE expected type `impl Future`
diff --git a/src/test/ui/suggestions/match-prev-arm-needing-semi.stderr b/src/test/ui/suggestions/match-prev-arm-needing-semi.stderr
index 8f4b860589e..7a4f74a1994 100644
--- a/src/test/ui/suggestions/match-prev-arm-needing-semi.stderr
+++ b/src/test/ui/suggestions/match-prev-arm-needing-semi.stderr
@@ -20,14 +20,14 @@ LL | |     };
    |
    = note:     expected type `()`
            found opaque type `impl Future`
-help: consider removing this semicolon and boxing the expression
-   |
-LL |             async_dummy()
-   |                         --
 help: consider `await`ing on the `Future`
    |
 LL |         false => async_dummy().await,
    |                               ^^^^^^
+help: consider removing this semicolon and boxing the expression
+   |
+LL |             async_dummy()
+   |                         --
 
 error[E0308]: `match` arms have incompatible types
   --> $DIR/match-prev-arm-needing-semi.rs:39:18
@@ -51,17 +51,17 @@ LL | |     };
    |
    = note:     expected type `()`
            found opaque type `impl Future`
-help: consider removing this semicolon and boxing the expression
-   |
-LL |             async_dummy()
-   |                         --
 help: consider `await`ing on the `Future`
    |
 LL |         false => async_dummy2().await,
    |                                ^^^^^^
+help: consider removing this semicolon and boxing the expression
+   |
+LL |             async_dummy()
+   |                         --
 
 error[E0308]: `match` arms have incompatible types
-  --> $DIR/match-prev-arm-needing-semi.rs:49:18
+  --> $DIR/match-prev-arm-needing-semi.rs:50:18
    |
 LL |   async fn async_dummy2() {}
    |                           - the `Output` of this `async fn`'s found opaque type
@@ -81,6 +81,11 @@ LL | |     };
    = note:     expected type `impl Future` (opaque type at <$DIR/match-prev-arm-needing-semi.rs:16:24>)
            found opaque type `impl Future` (opaque type at <$DIR/match-prev-arm-needing-semi.rs:17:25>)
    = note: distinct uses of `impl Trait` result in different opaque types
+help: consider `await`ing on both `Future`s
+   |
+LL |         true => async_dummy().await,
+LL |         false => async_dummy2().await,
+   |
 
 error[E0308]: `match` arms have incompatible types
   --> $DIR/match-prev-arm-needing-semi.rs:11:18
diff --git a/src/test/ui/suggestions/opaque-type-error.stderr b/src/test/ui/suggestions/opaque-type-error.stderr
index a7c2b82942f..d74076cbc9b 100644
--- a/src/test/ui/suggestions/opaque-type-error.stderr
+++ b/src/test/ui/suggestions/opaque-type-error.stderr
@@ -16,6 +16,12 @@ LL | |     }.await
    = note:     expected type `impl Future` (opaque type at <$DIR/opaque-type-error.rs:8:19>)
            found opaque type `impl Future` (opaque type at <$DIR/opaque-type-error.rs:12:19>)
    = note: distinct uses of `impl Trait` result in different opaque types
+help: consider `await`ing on both `Future`s
+   |
+LL |         thing_one().await
+LL |     } else {
+LL |         thing_two().await
+   |
 
 error: aborting due to previous error