Tweak await span

Author: compiler-errors

compiler/rustc_ast/src/ast.rs

@@ -1430,8 +1430,8 @@ pub enum ExprKind {
     /// The async block used to have a `NodeId`, which was removed in favor of
     /// using the parent `NodeId` of the parent `Expr`.
     Async(CaptureBy, P<Block>),
-    /// An await expression (`my_future.await`).
-    Await(P<Expr>),
+    /// An await expression (`my_future.await`). Span is of await keyword.
+    Await(P<Expr>, Span),
 
     /// A try block (`try { ... }`).
     TryBlock(P<Block>),

compiler/rustc_ast/src/mut_visit.rs

@@ -1415,7 +1415,10 @@ pub fn noop_visit_expr<T: MutVisitor>(
         ExprKind::Async(_capture_by, body) => {
             vis.visit_block(body);
         }
-        ExprKind::Await(expr) => vis.visit_expr(expr),
+        ExprKind::Await(expr, await_kw_span) => {
+            vis.visit_expr(expr);
+            vis.visit_span(await_kw_span);
+        }
         ExprKind::Assign(el, er, _) => {
             vis.visit_expr(el);
             vis.visit_expr(er);

compiler/rustc_ast/src/util/parser.rs

@@ -388,7 +388,7 @@ pub fn contains_exterior_struct_lit(value: &ast::Expr) -> bool {
             // X { y: 1 } + X { y: 2 }
             contains_exterior_struct_lit(lhs) || contains_exterior_struct_lit(rhs)
         }
-        ast::ExprKind::Await(x)
+        ast::ExprKind::Await(x, _)
         | ast::ExprKind::Unary(_, x)
         | ast::ExprKind::Cast(x, _)
         | ast::ExprKind::Type(x, _)

compiler/rustc_ast/src/visit.rs

@@ -864,7 +864,7 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
         ExprKind::Async(_, body) => {
             visitor.visit_block(body);
         }
-        ExprKind::Await(expr) => visitor.visit_expr(expr),
+        ExprKind::Await(expr, _) => visitor.visit_expr(expr),
         ExprKind::Assign(lhs, rhs, _) => {
             visitor.visit_expr(lhs);
             visitor.visit_expr(rhs);

compiler/rustc_ast_lowering/src/expr.rs

@@ -185,20 +185,14 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     hir::AsyncGeneratorKind::Block,
                     |this| this.with_new_scopes(|this| this.lower_block_expr(block)),
                 ),
-                ExprKind::Await(expr) => {
-                    let dot_await_span = if expr.span.hi() < e.span.hi() {
-                        let span_with_whitespace = self
-                            .tcx
-                            .sess
-                            .source_map()
-                            .span_extend_while(expr.span, char::is_whitespace)
-                            .unwrap_or(expr.span);
-                        span_with_whitespace.shrink_to_hi().with_hi(e.span.hi())
+                ExprKind::Await(expr, await_kw_span) => {
+                    let await_kw_span = if expr.span.hi() < await_kw_span.hi() {
+                        *await_kw_span
                     } else {
                         // this is a recovered `await expr`
                         e.span
                     };
-                    self.lower_expr_await(dot_await_span, expr)
+                    self.lower_expr_await(await_kw_span, expr)
                 }
                 ExprKind::Closure(box Closure {
                     binder,
@@ -710,18 +704,18 @@ impl<'hir> LoweringContext<'_, 'hir> {
     ///     }
     /// }
     /// ```
-    fn lower_expr_await(&mut self, dot_await_span: Span, expr: &Expr) -> hir::ExprKind<'hir> {
-        let full_span = expr.span.to(dot_await_span);
+    fn lower_expr_await(&mut self, await_kw_span: Span, expr: &Expr) -> hir::ExprKind<'hir> {
+        let full_span = expr.span.to(await_kw_span);
         match self.generator_kind {
             Some(hir::GeneratorKind::Async(_)) => {}
             Some(hir::GeneratorKind::Gen) | None => {
                 self.tcx.sess.emit_err(AwaitOnlyInAsyncFnAndBlocks {
-                    dot_await_span,
+                    dot_await_span: await_kw_span,
                     item_span: self.current_item,
                 });
             }
         }
-        let span = self.mark_span_with_reason(DesugaringKind::Await, dot_await_span, None);
+        let span = self.mark_span_with_reason(DesugaringKind::Await, await_kw_span, None);
         let gen_future_span = self.mark_span_with_reason(
             DesugaringKind::Await,
             full_span,

compiler/rustc_ast_lowering/src/format.rs

@@ -583,7 +583,7 @@ fn may_contain_yield_point(e: &ast::Expr) -> bool {
 
     impl Visitor<'_> for MayContainYieldPoint {
         fn visit_expr(&mut self, e: &ast::Expr) {
-            if let ast::ExprKind::Await(_) | ast::ExprKind::Yield(_) = e.kind {
+            if let ast::ExprKind::Await(_, _) | ast::ExprKind::Yield(_) = e.kind {
                 self.0 = true;
             } else {
                 visit::walk_expr(self, e);

compiler/rustc_ast_pretty/src/pprust/state/expr.rs

@@ -447,7 +447,7 @@ impl<'a> State<'a> {
                 self.ibox(0);
                 self.print_block_with_attrs(blk, attrs);
             }
-            ast::ExprKind::Await(expr) => {
+            ast::ExprKind::Await(expr, _) => {
                 self.print_expr_maybe_paren(expr, parser::PREC_POSTFIX);
                 self.word(".await");
             }

compiler/rustc_builtin_macros/src/assert/context.rs

@@ -288,7 +288,7 @@ impl<'cx, 'a> Context<'cx, 'a> {
             ExprKind::Assign(_, _, _)
             | ExprKind::AssignOp(_, _, _)
             | ExprKind::Async(_, _)
-            | ExprKind::Await(_)
+            | ExprKind::Await(_, _)
             | ExprKind::Block(_, _)
             | ExprKind::Break(_, _)
             | ExprKind::Closure(_)

compiler/rustc_parse/src/parser/diagnostics.rs

@@ -1646,7 +1646,7 @@ impl<'a> Parser<'a> {
             // Avoid knock-down errors as we don't know whether to interpret this as `foo().await?`
             // or `foo()?.await` (the very reason we went with postfix syntax 😅).
             ExprKind::Try(_) => ExprKind::Err,
-            _ => ExprKind::Await(expr),
+            _ => ExprKind::Await(expr, await_sp),
         };
         let expr = self.mk_expr(lo.to(sp), kind);
         self.maybe_recover_from_bad_qpath(expr)

compiler/rustc_parse/src/parser/expr.rs

@@ -859,7 +859,7 @@ impl<'a> Parser<'a> {
                     ExprKind::Field(_, _) => "a field access",
                     ExprKind::MethodCall(_) => "a method call",
                     ExprKind::Call(_, _) => "a function call",
-                    ExprKind::Await(_) => "`.await`",
+                    ExprKind::Await(_, _) => "`.await`",
                     ExprKind::Err => return Ok(with_postfix),
                     _ => unreachable!("parse_dot_or_call_expr_with_ shouldn't produce this"),
                 }
@@ -3256,7 +3256,7 @@ impl<'a> Parser<'a> {
 
     fn mk_await_expr(&mut self, self_arg: P<Expr>, lo: Span) -> P<Expr> {
         let span = lo.to(self.prev_token.span);
-        let await_expr = self.mk_expr(span, ExprKind::Await(self_arg));
+        let await_expr = self.mk_expr(span, ExprKind::Await(self_arg, self.prev_token.span));
         self.recover_from_await_method_call();
         await_expr
     }