Use clearer message when obligation is caused by await expr

5 files changed, 126 insertions, 46 deletions

diff --git a/src/librustc_ast_lowering/expr.rs b/src/librustc_ast_lowering/expr.rs
index 9984eb4e282..2fc58efbf0e 100644
--- a/src/librustc_ast_lowering/expr.rs
+++ b/src/librustc_ast_lowering/expr.rs
@@ -590,6 +590,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             await_span,
             self.allow_gen_future.clone(),
         );
+        let expr = self.lower_expr(expr);
 
         let pinned_ident = Ident::with_dummy_span(sym::pinned);
         let (pinned_pat, pinned_pat_hid) =
@@ -671,7 +672,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             let unit = self.expr_unit(span);
             let yield_expr = self.expr(
                 span,
-                hir::ExprKind::Yield(unit, hir::YieldSource::Await),
+                hir::ExprKind::Yield(unit, hir::YieldSource::Await { expr: Some(expr.hir_id) }),
                 ThinVec::new(),
             );
             let yield_expr = self.arena.alloc(yield_expr);
@@ -704,7 +705,6 @@ impl<'hir> LoweringContext<'_, 'hir> {
         // match <expr> {
         //     mut pinned => loop { .. }
         // }
-        let expr = self.lower_expr(expr);
         hir::ExprKind::Match(expr, arena_vec![self; pinned_arm], hir::MatchSource::AwaitDesugar)
     }
 
diff --git a/src/librustc_hir/hir.rs b/src/librustc_hir/hir.rs
index b719d576d6f..440fdc66286 100644
--- a/src/librustc_hir/hir.rs
+++ b/src/librustc_hir/hir.rs
@@ -1736,15 +1736,24 @@ pub struct Destination {
 #[derive(Copy, Clone, PartialEq, Eq, Debug, RustcEncodable, RustcDecodable, HashStable_Generic)]
 pub enum YieldSource {
     /// An `<expr>.await`.
-    Await,
+    Await { expr: Option<HirId> },
     /// A plain `yield`.
     Yield,
 }
 
+impl YieldSource {
+    pub fn is_await(&self) -> bool {
+        match self {
+            YieldSource::Await { .. } => true,
+            YieldSource::Yield => false,
+        }
+    }
+}
+
 impl fmt::Display for YieldSource {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.write_str(match self {
-            YieldSource::Await => "`await`",
+            YieldSource::Await { .. } => "`await`",
             YieldSource::Yield => "`yield`",
         })
     }
@@ -1755,7 +1764,7 @@ impl From<GeneratorKind> for YieldSource {
         match kind {
             // Guess based on the kind of the current generator.
             GeneratorKind::Gen => Self::Yield,
-            GeneratorKind::Async(_) => Self::Await,
+            GeneratorKind::Async(_) => Self::Await { expr: None },
         }
     }
 }
diff --git a/src/librustc_trait_selection/traits/error_reporting/suggestions.rs b/src/librustc_trait_selection/traits/error_reporting/suggestions.rs
index 0f4fbc41d16..f81992e19c9 100644
--- a/src/librustc_trait_selection/traits/error_reporting/suggestions.rs
+++ b/src/librustc_trait_selection/traits/error_reporting/suggestions.rs
@@ -127,13 +127,14 @@ pub trait InferCtxtExt<'tcx> {
         scope_span: &Option<Span>,
         expr: Option<hir::HirId>,
         snippet: String,
-        inner_generator: DefId,
+        inner_generator_body: Option<&hir::Body<'_>>,
         outer_generator: Option<DefId>,
         trait_ref: ty::TraitRef<'_>,
         target_ty: Ty<'tcx>,
         tables: &ty::TypeckTables<'_>,
         obligation: &PredicateObligation<'tcx>,
         next_code: Option<&ObligationCauseCode<'tcx>>,
+        from_awaited_ty: Option<Span>,
     );
 
     fn note_obligation_cause_code<T>(
@@ -1203,6 +1204,17 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
             }
         };
 
+        let generator_body = self.tcx
+            .hir()
+            .as_local_hir_id(generator_did)
+            .and_then(|hir_id| self.tcx.hir().maybe_body_owned_by(hir_id))
+            .map(|body_id| self.tcx.hir().body(body_id));
+        let mut visitor = AwaitsVisitor::default();
+        if let Some(body) = generator_body {
+            visitor.visit_body(body);
+        }
+        debug!("maybe_note_obligation_cause_for_async_await: awaits = {:?}", visitor.awaits);
+
         // Look for a type inside the generator interior that matches the target type to get
         // a span.
         let target_ty_erased = self.tcx.erase_regions(&target_ty);
@@ -1232,8 +1244,26 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
                 );
                 eq
             })
-            .map(|ty::GeneratorInteriorTypeCause { span, scope_span, expr, .. }| {
-                (span, source_map.span_to_snippet(*span), scope_span, expr)
+            .map(|cause| {
+                // Check to see if any awaited expressions have the target type.
+                let from_awaited_ty = visitor.awaits.into_iter()
+                    .map(|id| self.tcx.hir().expect_expr(id))
+                    .find(|expr| {
+                        let ty = tables.expr_ty_adjusted(&expr);
+                        // Compare types using the same logic as above.
+                        let ty_erased = self.tcx.erase_late_bound_regions(&ty::Binder::bind(ty));
+                        let ty_erased = self.tcx.erase_regions(&ty_erased);
+                        let eq = ty::TyS::same_type(ty_erased, target_ty_erased);
+                        debug!(
+                            "maybe_note_obligation_cause_for_async_await: await_expr={:?} \
+                            await_ty_erased={:?}  target_ty_erased={:?} eq={:?}",
+                            expr, ty_erased, target_ty_erased, eq
+                        );
+                        eq
+                    })
+                    .map(|expr| expr.span);
+                let ty::GeneratorInteriorTypeCause { span, scope_span, expr, .. } = cause;
+                (span, source_map.span_to_snippet(*span), scope_span, expr, from_awaited_ty)
             });
 
         debug!(
@@ -1241,20 +1271,21 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
                 generator_interior_types={:?} target_span={:?}",
             target_ty, tables.generator_interior_types, target_span
         );
-        if let Some((target_span, Ok(snippet), scope_span, expr)) = target_span {
+        if let Some((target_span, Ok(snippet), scope_span, expr, from_awaited_ty)) = target_span {
             self.note_obligation_cause_for_async_await(
                 err,
                 *target_span,
                 scope_span,
                 *expr,
                 snippet,
-                generator_did,
+                generator_body,
                 outer_generator,
                 trait_ref,
                 target_ty,
                 tables,
                 obligation,
                 next_code,
+                from_awaited_ty,
             );
             true
         } else {
@@ -1271,22 +1302,18 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
         scope_span: &Option<Span>,
         expr: Option<hir::HirId>,
         snippet: String,
-        inner_generator: DefId,
+        inner_generator_body: Option<&hir::Body<'_>>,
         outer_generator: Option<DefId>,
         trait_ref: ty::TraitRef<'_>,
         target_ty: Ty<'tcx>,
         tables: &ty::TypeckTables<'_>,
         obligation: &PredicateObligation<'tcx>,
         next_code: Option<&ObligationCauseCode<'tcx>>,
+        from_awaited_ty: Option<Span>,
     ) {
         let source_map = self.tcx.sess.source_map();
 
-        let is_async = self
-            .tcx
-            .hir()
-            .as_local_hir_id(inner_generator)
-            .and_then(|hir_id| self.tcx.hir().maybe_body_owned_by(hir_id))
-            .map(|body_id| self.tcx.hir().body(body_id))
+        let is_async = inner_generator_body
             .and_then(|body| body.generator_kind())
             .map(|generator_kind| match generator_kind {
                 hir::GeneratorKind::Async(..) => true,
@@ -1345,33 +1372,57 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
             )
         };
 
-        // Look at the last interior type to get a span for the `.await`.
-        let await_span = tables.generator_interior_types.iter().map(|t| t.span).last().unwrap();
-        let mut span = MultiSpan::from_span(await_span);
-        span.push_span_label(
-            await_span,
-            format!("{} occurs here, with `{}` maybe used later", await_or_yield, snippet),
-        );
+        let push_target_span = |span: &mut MultiSpan| {
+            if target_ty.is_impl_trait() {
+                // It's not very useful to tell the user the type if it's opaque.
+                span.push_span_label(target_span, "created here".to_string());
+            } else {
+                span.push_span_label(target_span, format!("has type `{}`", target_ty));
+            }
+        };
 
-        if target_ty.is_impl_trait() {
-            // It's not very useful to tell the user the type if it's opaque.
-            span.push_span_label(target_span, "created here".to_string());
-        } else {
-            span.push_span_label(target_span, format!("has type `{}`", target_ty));
-        }
+        if let Some(await_span) = from_awaited_ty {
+            // The type causing this obligation is one being awaited at await_span.
+            let mut span = MultiSpan::from_span(await_span);
+            span.push_span_label(
+                await_span,
+                "await occurs here".to_string(),
+            );
+
+            push_target_span(&mut span);
 
-        // If available, use the scope span to annotate the drop location.
-        if let Some(scope_span) = scope_span {
+            err.span_note(
+                span,
+                &format!("{} as this value is used in an await", trait_explanation),
+            );
+        } else {
+            // Look at the last interior type to get a span for the `.await`.
+            debug!(
+                "note_obligation_cause_for_async_await generator_interior_types: {:#?}",
+                tables.generator_interior_types
+            );
+            let await_span = tables.generator_interior_types.iter().map(|t| t.span).last().unwrap();
+            let mut span = MultiSpan::from_span(await_span);
             span.push_span_label(
-                source_map.end_point(*scope_span),
-                format!("`{}` is later dropped here", snippet),
+                await_span,
+                format!("{} occurs here, with `{}` maybe used later", await_or_yield, snippet),
             );
-        }
 
-        err.span_note(
-            span,
-            &format!("{} as this value is used across an {}", trait_explanation, await_or_yield),
-        );
+            push_target_span(&mut span);
+
+            // If available, use the scope span to annotate the drop location.
+            if let Some(scope_span) = scope_span {
+                span.push_span_label(
+                    source_map.end_point(*scope_span),
+                    format!("`{}` is later dropped here", snippet),
+                );
+            }
+
+            err.span_note(
+                span,
+                &format!("{} as this value is used across an {}", trait_explanation, await_or_yield),
+            );
+        }
 
         if let Some(expr_id) = expr {
             let expr = hir.expect_expr(expr_id);
@@ -1716,6 +1767,29 @@ impl<'v> Visitor<'v> for ReturnsVisitor<'v> {
     }
 }
 
+/// Collect all the awaited expressions within the input expression.
+#[derive(Default)]
+struct AwaitsVisitor {
+    awaits: Vec<hir::HirId>,
+}
+
+impl<'v> Visitor<'v> for AwaitsVisitor {
+    type Map = hir::intravisit::ErasedMap<'v>;
+
+    fn nested_visit_map(&mut self) -> hir::intravisit::NestedVisitorMap<Self::Map> {
+        hir::intravisit::NestedVisitorMap::None
+    }
+
+    fn visit_expr(&mut self, ex: &'v hir::Expr<'v>) {
+        match ex.kind {
+            hir::ExprKind::Yield(_, hir::YieldSource::Await { expr: Some(id) }) =>
+                self.awaits.push(id),
+            _ => (),
+        }
+        hir::intravisit::walk_expr(self, ex)
+    }
+}
+
 pub trait NextTypeParamName {
     fn next_type_param_name(&self, name: Option<&str>) -> String;
 }
diff --git a/src/librustc_typeck/check/expr.rs b/src/librustc_typeck/check/expr.rs
index 57e2349bb2e..7cb51b4d6d8 100644
--- a/src/librustc_typeck/check/expr.rs
+++ b/src/librustc_typeck/check/expr.rs
@@ -1797,7 +1797,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             // we know that the yield type must be `()`; however, the context won't contain this
             // information. Hence, we check the source of the yield expression here and check its
             // value's type against `()` (this check should always hold).
-            None if src == &hir::YieldSource::Await => {
+            None if src.is_await() => {
                 self.check_expr_coercable_to_type(&value, self.tcx.mk_unit());
                 self.tcx.mk_unit()
             }
diff --git a/src/test/ui/async-await/issue-68112.stderr b/src/test/ui/async-await/issue-68112.stderr
index eec2171024a..f79908110c4 100644
--- a/src/test/ui/async-await/issue-68112.stderr
+++ b/src/test/ui/async-await/issue-68112.stderr
@@ -8,16 +8,13 @@ LL |     require_send(send_fut);
    |     ^^^^^^^^^^^^ future created by async block is not `Send`
    |
    = help: the trait `std::marker::Sync` is not implemented for `std::cell::RefCell<i32>`
-note: future is not `Send` as this value is used across an await
-  --> $DIR/issue-68112.rs:32:9
+note: future is not `Send` as this value is used in an await
+  --> $DIR/issue-68112.rs:31:17
    |
 LL |         let non_send_fut = make_non_send_future1();
    |             ------------ created here
 LL |         let _ = non_send_fut.await;
-LL |         ready(0).await;
-   |         ^^^^^^^^ await occurs here, with `non_send_fut` maybe used later
-LL |     };
-   |     - `non_send_fut` is later dropped here
+   |                 ^^^^^^^^^^^^ await occurs here
 
 error[E0277]: `std::cell::RefCell<i32>` cannot be shared between threads safely
   --> $DIR/issue-68112.rs:49:5