Tweak await span

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);