suggest await before try when performing trait selection

4 files changed, 92 insertions, 10 deletions

diff --git a/src/libcore/ops/try.rs b/src/libcore/ops/try.rs
index ed8c614b109..ef748bcee6e 100644
--- a/src/libcore/ops/try.rs
+++ b/src/libcore/ops/try.rs
@@ -25,7 +25,7 @@
     )
 )]
 #[doc(alias = "?")]
- #[cfg_attr(not(bootstrap), lang = "try_trait")]
+#[cfg_attr(not(bootstrap), lang = "try_trait")]
 pub trait Try {
     /// The type of this value when viewed as successful.
     #[unstable(feature = "try_trait", issue = "42327")]
diff --git a/src/librustc_trait_selection/traits/error_reporting/mod.rs b/src/librustc_trait_selection/traits/error_reporting/mod.rs
index 405c656bad5..272827cfef8 100644
--- a/src/librustc_trait_selection/traits/error_reporting/mod.rs
+++ b/src/librustc_trait_selection/traits/error_reporting/mod.rs
@@ -400,6 +400,17 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
                         self.suggest_remove_reference(&obligation, &mut err, &trait_ref);
                         self.suggest_semicolon_removal(&obligation, &mut err, span, &trait_ref);
                         self.note_version_mismatch(&mut err, &trait_ref);
+                        //self.sugggest_await_before_try(&mut err, &obligation, &trait_ref);
+                        debug!(
+                            "suggest_await_befor_try: trait_predicate={:?} obligation={:?}, trait_ref={:?}",
+                            trait_predicate, obligation, trait_ref
+                        );
+                        self.suggest_await_befor_try(
+                            &mut err,
+                            &obligation,
+                            trait_ref.self_ty(),
+                            span,
+                        );
                         if self.suggest_impl_trait(&mut err, span, &obligation, &trait_ref) {
                             err.emit();
                             return;
diff --git a/src/librustc_trait_selection/traits/error_reporting/suggestions.rs b/src/librustc_trait_selection/traits/error_reporting/suggestions.rs
index 74dd47a91c2..d0b39d6016a 100644
--- a/src/librustc_trait_selection/traits/error_reporting/suggestions.rs
+++ b/src/librustc_trait_selection/traits/error_reporting/suggestions.rs
@@ -1,8 +1,10 @@
 use super::{
     EvaluationResult, Obligation, ObligationCause, ObligationCauseCode, PredicateObligation,
+    SelectionContext,
 };
 
 use crate::infer::InferCtxt;
+use crate::traits::normalize_projection_type;
 
 use rustc_errors::{error_code, struct_span_err, Applicability, DiagnosticBuilder, Style};
 use rustc_hir as hir;
@@ -150,6 +152,15 @@ pub trait InferCtxtExt<'tcx> {
         T: fmt::Display;
 
     fn suggest_new_overflow_limit(&self, err: &mut DiagnosticBuilder<'_>);
+
+    /// Suggest to await before try: future? => future.await?
+    fn suggest_await_befor_try(
+        &self,
+        err: &mut DiagnosticBuilder<'_>,
+        obligation: &PredicateObligation<'tcx>,
+        ty: Ty<'tcx>,
+        span: Span,
+    );
 }
 
 fn predicate_constraint(generics: &hir::Generics<'_>, pred: String) -> (Span, String) {
@@ -1765,6 +1776,75 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
             suggested_limit, self.tcx.crate_name,
         ));
     }
+
+    fn suggest_await_befor_try(
+        &self,
+        err: &mut DiagnosticBuilder<'_>,
+        obligation: &PredicateObligation<'tcx>,
+        ty: Ty<'tcx>,
+        span: Span,
+    ) {
+        debug!("suggest_await_befor_try: obligation={:?}, span={:?}", obligation, span);
+        let body_hir_id = obligation.cause.body_id;
+        let item_id = self.tcx.hir().get_parent_node(body_hir_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 {
+                // Check for `Future` implementations by constructing a predicate to
+                // prove: `<T as Future>::Output == U`
+                let future_trait = self.tcx.lang_items().future_trait().unwrap();
+                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(ty, self.fresh_substs_for_item(span, item_def_id)),
+                    // `Future::Output`
+                    item_def_id,
+                };
+
+                let cause = ObligationCause::misc(span, body_hir_id);
+                let mut selcx = SelectionContext::new(self);
+
+                let mut obligations = vec![];
+                let normalized_ty = normalize_projection_type(
+                    &mut selcx,
+                    obligation.param_env,
+                    projection_ty,
+                    obligation.cause.clone(),
+                    0,
+                    &mut obligations,
+                );
+
+                debug!("suggest_await_befor_try: normalized_projection_type {:?}", normalized_ty);
+                let try_trait_ref_id = self.tcx.lang_items().try_trait().unwrap();
+                if let Some(try_trait_ref) = self.tcx.impl_trait_ref(try_trait_ref_id) {
+                    let try_predicate = try_trait_ref.without_const().to_predicate();
+                    let try_obligation =
+                        Obligation::new(cause, obligation.param_env, try_predicate);
+                    debug!("suggest_await_befor_try: try_trait_obligation {:?}", try_obligation);
+                    if self.predicate_may_hold(&try_obligation) {
+                        debug!("try_obligation holds");
+                        if let Ok(snippet) = self.tcx.sess.source_map().span_to_snippet(span) {
+                            err.span_suggestion(
+                                span,
+                                "consider using `.await` here",
+                                format!("{}.await", snippet),
+                                Applicability::MaybeIncorrect,
+                            );
+                        }
+                    }
+                }
+            }
+        }
+    }
 }
 
 /// Collect all the returned expressions within the input expression.
diff --git a/src/librustc_typeck/check/mod.rs b/src/librustc_typeck/check/mod.rs
index bd5049446e6..c142e88b7de 100644
--- a/src/librustc_typeck/check/mod.rs
+++ b/src/librustc_typeck/check/mod.rs
@@ -5317,15 +5317,6 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     item_def_id,
                 };
 
-                let cause = traits::ObligationCause::misc(sp, self.body_id);
-                let normalized_ty = self.fulfillment_cx.borrow_mut().normalize_projection_type(
-                    &self.infcx,
-                    self.param_env,
-                    projection_ty,
-                    cause,
-                );
-                debug!("suggest_missing_await: projection_type {:?}", normalized_ty);
-
                 let predicate =
                     ty::Predicate::Projection(ty::Binder::bind(ty::ProjectionPredicate {
                         projection_ty,