Attach tokens to `ast::Stmt`

We currently only attach tokens when parsing a `:stmt` matcher for a
`macro_rules!` macro. Proc-macro attributes on statements are still
unstable, and need additional work.

11 files changed, 56 insertions, 15 deletions

diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
index e2703989646..dee3a16f9b1 100644
--- a/compiler/rustc_ast/src/ast.rs
+++ b/compiler/rustc_ast/src/ast.rs
@@ -918,6 +918,7 @@ pub struct Stmt {
     pub id: NodeId,
     pub kind: StmtKind,
     pub span: Span,
+    pub tokens: Option<TokenStream>,
 }
 
 impl Stmt {
diff --git a/compiler/rustc_ast/src/mut_visit.rs b/compiler/rustc_ast/src/mut_visit.rs
index d37bba646ef..425ef83b57a 100644
--- a/compiler/rustc_ast/src/mut_visit.rs
+++ b/compiler/rustc_ast/src/mut_visit.rs
@@ -1286,12 +1286,15 @@ pub fn noop_filter_map_expr<T: MutVisitor>(mut e: P<Expr>, vis: &mut T) -> Optio
 }
 
 pub fn noop_flat_map_stmt<T: MutVisitor>(
-    Stmt { kind, mut span, mut id }: Stmt,
+    Stmt { kind, mut span, mut id, tokens }: Stmt,
     vis: &mut T,
 ) -> SmallVec<[Stmt; 1]> {
     vis.visit_id(&mut id);
     vis.visit_span(&mut span);
-    noop_flat_map_stmt_kind(kind, vis).into_iter().map(|kind| Stmt { id, kind, span }).collect()
+    noop_flat_map_stmt_kind(kind, vis)
+        .into_iter()
+        .map(|kind| Stmt { id, kind, span, tokens: tokens.clone() })
+        .collect()
 }
 
 pub fn noop_flat_map_stmt_kind<T: MutVisitor>(
diff --git a/compiler/rustc_builtin_macros/src/deriving/debug.rs b/compiler/rustc_builtin_macros/src/deriving/debug.rs
index 120e859f2b1..d84b3956475 100644
--- a/compiler/rustc_builtin_macros/src/deriving/debug.rs
+++ b/compiler/rustc_builtin_macros/src/deriving/debug.rs
@@ -133,5 +133,5 @@ fn stmt_let_underscore(cx: &mut ExtCtxt<'_>, sp: Span, expr: P<ast::Expr>) -> as
         span: sp,
         attrs: ast::AttrVec::new(),
     });
-    ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Local(local), span: sp }
+    ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Local(local), span: sp, tokens: None }
 }
diff --git a/compiler/rustc_expand/src/base.rs b/compiler/rustc_expand/src/base.rs
index 5bdc4760dcc..926e3dbfc52 100644
--- a/compiler/rustc_expand/src/base.rs
+++ b/compiler/rustc_expand/src/base.rs
@@ -400,6 +400,7 @@ macro_rules! make_stmts_default {
                 id: ast::DUMMY_NODE_ID,
                 span: e.span,
                 kind: ast::StmtKind::Expr(e),
+                tokens: None
             }]
         })
     };
@@ -642,6 +643,7 @@ impl MacResult for DummyResult {
             id: ast::DUMMY_NODE_ID,
             kind: ast::StmtKind::Expr(DummyResult::raw_expr(self.span, self.is_error)),
             span: self.span,
+            tokens: None
         }])
     }
 
diff --git a/compiler/rustc_expand/src/build.rs b/compiler/rustc_expand/src/build.rs
index 70603622bed..a5a7ee6c9a3 100644
--- a/compiler/rustc_expand/src/build.rs
+++ b/compiler/rustc_expand/src/build.rs
@@ -158,7 +158,12 @@ impl<'a> ExtCtxt<'a> {
     }
 
     pub fn stmt_expr(&self, expr: P<ast::Expr>) -> ast::Stmt {
-        ast::Stmt { id: ast::DUMMY_NODE_ID, span: expr.span, kind: ast::StmtKind::Expr(expr) }
+        ast::Stmt {
+            id: ast::DUMMY_NODE_ID,
+            span: expr.span,
+            kind: ast::StmtKind::Expr(expr),
+            tokens: None,
+        }
     }
 
     pub fn stmt_let(&self, sp: Span, mutbl: bool, ident: Ident, ex: P<ast::Expr>) -> ast::Stmt {
@@ -176,7 +181,12 @@ impl<'a> ExtCtxt<'a> {
             span: sp,
             attrs: AttrVec::new(),
         });
-        ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Local(local), span: sp }
+        ast::Stmt {
+            id: ast::DUMMY_NODE_ID,
+            kind: ast::StmtKind::Local(local),
+            span: sp,
+            tokens: None,
+        }
     }
 
     // Generates `let _: Type;`, which is usually used for type assertions.
@@ -189,11 +199,16 @@ impl<'a> ExtCtxt<'a> {
             span,
             attrs: AttrVec::new(),
         });
-        ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Local(local), span }
+        ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Local(local), span, tokens: None }
     }
 
     pub fn stmt_item(&self, sp: Span, item: P<ast::Item>) -> ast::Stmt {
-        ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Item(item), span: sp }
+        ast::Stmt {
+            id: ast::DUMMY_NODE_ID,
+            kind: ast::StmtKind::Item(item),
+            span: sp,
+            tokens: None,
+        }
     }
 
     pub fn block_expr(&self, expr: P<ast::Expr>) -> P<ast::Block> {
@@ -203,6 +218,7 @@ impl<'a> ExtCtxt<'a> {
                 id: ast::DUMMY_NODE_ID,
                 span: expr.span,
                 kind: ast::StmtKind::Expr(expr),
+                tokens: None,
             }],
         )
     }
diff --git a/compiler/rustc_expand/src/expand.rs b/compiler/rustc_expand/src/expand.rs
index 47c070f1651..e5cfb866938 100644
--- a/compiler/rustc_expand/src/expand.rs
+++ b/compiler/rustc_expand/src/expand.rs
@@ -1396,10 +1396,10 @@ impl<'a, 'b> MutVisitor for InvocationCollector<'a, 'b> {
         }
 
         // The placeholder expander gives ids to statements, so we avoid folding the id here.
-        let ast::Stmt { id, kind, span } = stmt;
+        let ast::Stmt { id, kind, span, tokens } = stmt;
         noop_flat_map_stmt_kind(kind, self)
             .into_iter()
-            .map(|kind| ast::Stmt { id, kind, span })
+            .map(|kind| ast::Stmt { id, kind, span, tokens: tokens.clone() })
             .collect()
     }
 
diff --git a/compiler/rustc_expand/src/placeholders.rs b/compiler/rustc_expand/src/placeholders.rs
index 036c00bb753..4c9271a58df 100644
--- a/compiler/rustc_expand/src/placeholders.rs
+++ b/compiler/rustc_expand/src/placeholders.rs
@@ -105,7 +105,7 @@ pub fn placeholder(
                 style: ast::MacStmtStyle::Braces,
                 attrs: ast::AttrVec::new(),
             });
-            ast::Stmt { id, span, kind: ast::StmtKind::MacCall(mac) }
+            ast::Stmt { id, span, kind: ast::StmtKind::MacCall(mac), tokens: None }
         }]),
         AstFragmentKind::Arms => AstFragment::Arms(smallvec![ast::Arm {
             attrs: Default::default(),
diff --git a/compiler/rustc_interface/src/util.rs b/compiler/rustc_interface/src/util.rs
index 01441cafe8a..f15eb413833 100644
--- a/compiler/rustc_interface/src/util.rs
+++ b/compiler/rustc_interface/src/util.rs
@@ -710,6 +710,7 @@ impl<'a> MutVisitor for ReplaceBodyWithLoop<'a, '_> {
                 id: resolver.next_node_id(),
                 kind: ast::StmtKind::Expr(expr),
                 span: rustc_span::DUMMY_SP,
+                tokens: None,
             }
         }
 
@@ -726,6 +727,7 @@ impl<'a> MutVisitor for ReplaceBodyWithLoop<'a, '_> {
             id: self.resolver.next_node_id(),
             span: rustc_span::DUMMY_SP,
             kind: ast::StmtKind::Expr(loop_expr),
+            tokens: None,
         };
 
         if self.within_static_or_const {
diff --git a/compiler/rustc_parse/src/lib.rs b/compiler/rustc_parse/src/lib.rs
index 3807daa7015..dedb9850b5a 100644
--- a/compiler/rustc_parse/src/lib.rs
+++ b/compiler/rustc_parse/src/lib.rs
@@ -269,6 +269,13 @@ pub fn nt_to_tokenstream(nt: &Nonterminal, sess: &ParseSess, span: Span) -> Toke
             prepend_attrs(sess, &item.attrs, item.tokens.as_ref(), span)
         }
         Nonterminal::NtBlock(ref block) => block.tokens.clone(),
+        Nonterminal::NtStmt(ref stmt) => {
+            // FIXME: We currently only collect tokens for `:stmt`
+            // matchers in `macro_rules!` macros. When we start collecting
+            // tokens for attributes on statements, we will need to prepend
+            // attributes here
+            stmt.tokens.clone()
+        }
         Nonterminal::NtPat(ref pat) => pat.tokens.clone(),
         Nonterminal::NtTy(ref ty) => ty.tokens.clone(),
         Nonterminal::NtIdent(ident, is_raw) => {
diff --git a/compiler/rustc_parse/src/parser/nonterminal.rs b/compiler/rustc_parse/src/parser/nonterminal.rs
index f98f2a89dc6..15660fd574c 100644
--- a/compiler/rustc_parse/src/parser/nonterminal.rs
+++ b/compiler/rustc_parse/src/parser/nonterminal.rs
@@ -119,10 +119,20 @@ impl<'a> Parser<'a> {
                 }
                 token::NtBlock(block)
             }
-            NonterminalKind::Stmt => match self.parse_stmt()? {
-                Some(s) => token::NtStmt(s),
-                None => return Err(self.struct_span_err(self.token.span, "expected a statement")),
-            },
+            NonterminalKind::Stmt => {
+                let (stmt, tokens) = self.collect_tokens(|this| this.parse_stmt())?;
+                match stmt {
+                    Some(mut s) => {
+                        if s.tokens.is_none() {
+                            s.tokens = Some(tokens);
+                        }
+                        token::NtStmt(s)
+                    }
+                    None => {
+                        return Err(self.struct_span_err(self.token.span, "expected a statement"));
+                    }
+                }
+            }
             NonterminalKind::Pat => {
                 let (mut pat, tokens) = self.collect_tokens(|this| this.parse_pat(None))?;
                 // We have have eaten an NtPat, which could already have tokens
diff --git a/compiler/rustc_parse/src/parser/stmt.rs b/compiler/rustc_parse/src/parser/stmt.rs
index 6cc42487684..64b959e8325 100644
--- a/compiler/rustc_parse/src/parser/stmt.rs
+++ b/compiler/rustc_parse/src/parser/stmt.rs
@@ -415,7 +415,7 @@ impl<'a> Parser<'a> {
     }
 
     pub(super) fn mk_stmt(&self, span: Span, kind: StmtKind) -> Stmt {
-        Stmt { id: DUMMY_NODE_ID, kind, span }
+        Stmt { id: DUMMY_NODE_ID, kind, span, tokens: None }
     }
 
     fn mk_stmt_err(&self, span: Span) -> Stmt {