Properly handle attributes on statements

We now collect tokens for the underlying node wrapped by `StmtKind`
instead of storing tokens directly in `Stmt`.

`LazyTokenStream` now supports capturing a trailing semicolon after it
is initially constructed. This allows us to avoid refactoring statement
parsing to wrap the parsing of the semicolon in `parse_tokens`.

Attributes on item statements
(e.g. `fn foo() { #[bar] struct MyStruct; }`) are now treated as
item attributes, not statement attributes, which is consistent with how
we handle attributes on other kinds of statements. The feature-gating
code is adjusted so that proc-macro attributes are still allowed on item
statements on stable.

Two built-in macros (`#[global_allocator]` and `#[test]`) needed to be
adjusted to support being passed `Annotatable::Stmt`.

compiler/rustc_ast/src/ast.rs

@@ -901,29 +901,65 @@ pub struct Stmt {
     pub id: NodeId,
     pub kind: StmtKind,
     pub span: Span,
-    pub tokens: Option<LazyTokenStream>,
 }
 
 impl Stmt {
+    pub fn tokens(&self) -> Option<&LazyTokenStream> {
+        match self.kind {
+            StmtKind::Local(ref local) => local.tokens.as_ref(),
+            StmtKind::Item(ref item) => item.tokens.as_ref(),
+            StmtKind::Expr(ref expr) | StmtKind::Semi(ref expr) => expr.tokens.as_ref(),
+            StmtKind::Empty => None,
+            StmtKind::MacCall(ref mac) => mac.tokens.as_ref(),
+        }
+    }
+
+    pub fn tokens_mut(&mut self) -> Option<&mut LazyTokenStream> {
+        match self.kind {
+            StmtKind::Local(ref mut local) => local.tokens.as_mut(),
+            StmtKind::Item(ref mut item) => item.tokens.as_mut(),
+            StmtKind::Expr(ref mut expr) | StmtKind::Semi(ref mut expr) => expr.tokens.as_mut(),
+            StmtKind::Empty => None,
+            StmtKind::MacCall(ref mut mac) => mac.tokens.as_mut(),
+        }
+    }
+
+    pub fn set_tokens(&mut self, tokens: Option<LazyTokenStream>) {
+        match self.kind {
+            StmtKind::Local(ref mut local) => local.tokens = tokens,
+            StmtKind::Item(ref mut item) => item.tokens = tokens,
+            StmtKind::Expr(ref mut expr) | StmtKind::Semi(ref mut expr) => expr.tokens = tokens,
+            StmtKind::Empty => {}
+            StmtKind::MacCall(ref mut mac) => mac.tokens = tokens,
+        }
+    }
+
     pub fn has_trailing_semicolon(&self) -> bool {
         match &self.kind {
             StmtKind::Semi(_) => true,
             StmtKind::MacCall(mac) => matches!(mac.style, MacStmtStyle::Semicolon),
             _ => false,
         }
     }
+
+    /// Converts a parsed `Stmt` to a `Stmt` with
+    /// a trailing semicolon.
+    ///
+    /// This only modifies the parsed AST struct, not the attached
+    /// `LazyTokenStream`. The parser is responsible for calling
+    /// `CreateTokenStream::add_trailing_semi` when there is actually
+    /// a semicolon in the tokenstream.
     pub fn add_trailing_semicolon(mut self) -> Self {
         self.kind = match self.kind {
             StmtKind::Expr(expr) => StmtKind::Semi(expr),
             StmtKind::MacCall(mac) => {
-                StmtKind::MacCall(mac.map(|MacCallStmt { mac, style: _, attrs }| MacCallStmt {
-                    mac,
-                    style: MacStmtStyle::Semicolon,
-                    attrs,
+                StmtKind::MacCall(mac.map(|MacCallStmt { mac, style: _, attrs, tokens }| {
+                    MacCallStmt { mac, style: MacStmtStyle::Semicolon, attrs, tokens }
                 }))
             }
             kind => kind,
         };
+
         self
     }
 
@@ -963,6 +999,7 @@ pub struct MacCallStmt {
     pub mac: MacCall,
     pub style: MacStmtStyle,
     pub attrs: AttrVec,
+    pub tokens: Option<LazyTokenStream>,
 }
 
 #[derive(Clone, Copy, PartialEq, Encodable, Decodable, Debug)]
@@ -988,6 +1025,7 @@ pub struct Local {
     pub init: Option<P<Expr>>,
     pub span: Span,
     pub attrs: AttrVec,
+    pub tokens: Option<LazyTokenStream>,
 }
 
 /// An arm of a 'match'.

compiler/rustc_ast/src/mut_visit.rs

@@ -576,13 +576,14 @@ pub fn noop_visit_parenthesized_parameter_data<T: MutVisitor>(
 }
 
 pub fn noop_visit_local<T: MutVisitor>(local: &mut P<Local>, vis: &mut T) {
-    let Local { id, pat, ty, init, span, attrs } = local.deref_mut();
+    let Local { id, pat, ty, init, span, attrs, tokens } = local.deref_mut();
     vis.visit_id(id);
     vis.visit_pat(pat);
     visit_opt(ty, |ty| vis.visit_ty(ty));
     visit_opt(init, |init| vis.visit_expr(init));
     vis.visit_span(span);
     visit_thin_attrs(attrs, vis);
+    visit_lazy_tts(tokens, vis);
 }
 
 pub fn noop_visit_attribute<T: MutVisitor>(attr: &mut Attribute, vis: &mut T) {
@@ -1325,16 +1326,12 @@ 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, mut tokens }: Stmt,
+    Stmt { kind, mut span, mut id }: Stmt,
     vis: &mut T,
 ) -> SmallVec<[Stmt; 1]> {
     vis.visit_id(&mut id);
     vis.visit_span(&mut span);
-    visit_lazy_tts(&mut tokens, vis);
-    noop_flat_map_stmt_kind(kind, vis)
-        .into_iter()
-        .map(|kind| Stmt { id, kind, span, tokens: tokens.clone() })
-        .collect()
+    noop_flat_map_stmt_kind(kind, vis).into_iter().map(|kind| Stmt { id, kind, span }).collect()
 }
 
 pub fn noop_flat_map_stmt_kind<T: MutVisitor>(
@@ -1351,9 +1348,10 @@ pub fn noop_flat_map_stmt_kind<T: MutVisitor>(
         StmtKind::Semi(expr) => vis.filter_map_expr(expr).into_iter().map(StmtKind::Semi).collect(),
         StmtKind::Empty => smallvec![StmtKind::Empty],
         StmtKind::MacCall(mut mac) => {
-            let MacCallStmt { mac: mac_, style: _, attrs } = mac.deref_mut();
+            let MacCallStmt { mac: mac_, style: _, attrs, tokens } = mac.deref_mut();
             vis.visit_mac_call(mac_);
             visit_thin_attrs(attrs, vis);
+            visit_lazy_tts(tokens, vis);
             smallvec![StmtKind::MacCall(mac)]
         }
     }

compiler/rustc_ast/src/tokenstream.rs

@@ -121,10 +121,14 @@ where
 }
 
 pub trait CreateTokenStream: sync::Send + sync::Sync {
+    fn add_trailing_semi(&self) -> Box<dyn CreateTokenStream>;
     fn create_token_stream(&self) -> TokenStream;
 }
 
 impl CreateTokenStream for TokenStream {
+    fn add_trailing_semi(&self) -> Box<dyn CreateTokenStream> {
+        panic!("Cannot call `add_trailing_semi` on a `TokenStream`!");
+    }
     fn create_token_stream(&self) -> TokenStream {
         self.clone()
     }
@@ -141,6 +145,13 @@ impl LazyTokenStream {
         LazyTokenStream(Lrc::new(Box::new(inner)))
     }
 
+    /// Extends the captured stream by one token,
+    /// which must be a trailing semicolon. This
+    /// affects the `TokenStream` created by `make_tokenstream`.
+    pub fn add_trailing_semi(&self) -> LazyTokenStream {
+        LazyTokenStream(Lrc::new(self.0.add_trailing_semi()))
+    }
+
     pub fn create_token_stream(&self) -> TokenStream {
         self.0.create_token_stream()
     }

compiler/rustc_ast/src/visit.rs

@@ -686,7 +686,7 @@ pub fn walk_stmt<'a, V: Visitor<'a>>(visitor: &mut V, statement: &'a Stmt) {
         StmtKind::Expr(ref expr) | StmtKind::Semi(ref expr) => visitor.visit_expr(expr),
         StmtKind::Empty => {}
         StmtKind::MacCall(ref mac) => {
-            let MacCallStmt { ref mac, style: _, ref attrs } = **mac;
+            let MacCallStmt { ref mac, style: _, ref attrs, tokens: _ } = **mac;
             visitor.visit_mac_call(mac);
             for attr in attrs.iter() {
                 visitor.visit_attribute(attr);

compiler/rustc_builtin_macros/src/deriving/debug.rs

@@ -132,6 +132,7 @@ fn stmt_let_underscore(cx: &mut ExtCtxt<'_>, sp: Span, expr: P<ast::Expr>) -> as
         id: ast::DUMMY_NODE_ID,
         span: sp,
         attrs: ast::AttrVec::new(),
+        tokens: None,
     });
-    ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Local(local), span: sp, tokens: None }
+    ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Local(local), span: sp }
 }

compiler/rustc_builtin_macros/src/global_allocator.rs

@@ -4,7 +4,7 @@ use rustc_ast::expand::allocator::{
     AllocatorKind, AllocatorMethod, AllocatorTy, ALLOCATOR_METHODS,
 };
 use rustc_ast::ptr::P;
-use rustc_ast::{self as ast, Attribute, Expr, FnHeader, FnSig, Generics, Param};
+use rustc_ast::{self as ast, Attribute, Expr, FnHeader, FnSig, Generics, Param, StmtKind};
 use rustc_ast::{ItemKind, Mutability, Stmt, Ty, TyKind, Unsafe};
 use rustc_expand::base::{Annotatable, ExtCtxt};
 use rustc_span::symbol::{kw, sym, Ident, Symbol};
@@ -14,14 +14,25 @@ pub fn expand(
     ecx: &mut ExtCtxt<'_>,
     _span: Span,
     meta_item: &ast::MetaItem,
-    item: Annotatable,
+    mut item: Annotatable,
 ) -> Vec<Annotatable> {
     check_builtin_macro_attribute(ecx, meta_item, sym::global_allocator);
 
     let not_static = |item: Annotatable| {
         ecx.sess.parse_sess.span_diagnostic.span_err(item.span(), "allocators must be statics");
         vec![item]
     };
+    let orig_item = item.clone();
+    let mut is_stmt = false;
+
+    // Allow using `#[global_allocator]` on an item statement
+    if let Annotatable::Stmt(stmt) = &item {
+        if let StmtKind::Item(item_) = &stmt.kind {
+            item = Annotatable::Item(item_.clone());
+            is_stmt = true;
+        }
+    }
+
     let item = match item {
         Annotatable::Item(item) => match item.kind {
             ItemKind::Static(..) => item,
@@ -41,9 +52,14 @@ pub fn expand(
     let const_ty = ecx.ty(span, TyKind::Tup(Vec::new()));
     let const_body = ecx.expr_block(ecx.block(span, stmts));
     let const_item = ecx.item_const(span, Ident::new(kw::Underscore, span), const_ty, const_body);
+    let const_item = if is_stmt {
+        Annotatable::Stmt(P(ecx.stmt_item(span, const_item)))
+    } else {
+        Annotatable::Item(const_item)
+    };
 
     // Return the original item and the new methods.
-    vec![Annotatable::Item(item), Annotatable::Item(const_item)]
+    vec![orig_item, const_item]
 }
 
 struct AllocFnFactory<'a, 'b> {

compiler/rustc_builtin_macros/src/test.rs

@@ -4,6 +4,7 @@ use crate::util::check_builtin_macro_attribute;
 
 use rustc_ast as ast;
 use rustc_ast::attr;
+use rustc_ast::ptr::P;
 use rustc_ast_pretty::pprust;
 use rustc_expand::base::*;
 use rustc_session::Session;
@@ -78,8 +79,16 @@ pub fn expand_test_or_bench(
         return vec![];
     }
 
-    let item = match item {
-        Annotatable::Item(i) => i,
+    let (item, is_stmt) = match item {
+        Annotatable::Item(i) => (i, false),
+        Annotatable::Stmt(stmt) if matches!(stmt.kind, ast::StmtKind::Item(_)) => {
+            // FIXME: Use an 'if let' guard once they are implemented
+            if let ast::StmtKind::Item(i) = stmt.into_inner().kind {
+                (i, true)
+            } else {
+                unreachable!()
+            }
+        }
         other => {
             cx.struct_span_err(
                 other.span(),
@@ -304,14 +313,25 @@ pub fn expand_test_or_bench(
 
     tracing::debug!("synthetic test item:\n{}\n", pprust::item_to_string(&test_const));
 
-    vec![
-        // Access to libtest under a hygienic name
-        Annotatable::Item(test_extern),
-        // The generated test case
-        Annotatable::Item(test_const),
-        // The original item
-        Annotatable::Item(item),
-    ]
+    if is_stmt {
+        vec![
+            // Access to libtest under a hygienic name
+            Annotatable::Stmt(P(cx.stmt_item(sp, test_extern))),
+            // The generated test case
+            Annotatable::Stmt(P(cx.stmt_item(sp, test_const))),
+            // The original item
+            Annotatable::Stmt(P(cx.stmt_item(sp, item))),
+        ]
+    } else {
+        vec![
+            // Access to libtest under a hygienic name
+            Annotatable::Item(test_extern),
+            // The generated test case
+            Annotatable::Item(test_const),
+            // The original item
+            Annotatable::Item(item),
+        ]
+    }
 }
 
 fn item_path(mod_path: &[Ident], item_ident: &Ident) -> String {

compiler/rustc_expand/src/base.rs

@@ -366,7 +366,6 @@ macro_rules! make_stmts_default {
                 id: ast::DUMMY_NODE_ID,
                 span: e.span,
                 kind: ast::StmtKind::Expr(e),
-                tokens: None
             }]
         })
     };
@@ -609,7 +608,6 @@ 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
         }])
     }
 

compiler/rustc_expand/src/build.rs

@@ -140,12 +140,7 @@ 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),
-            tokens: None,
-        }
+        ast::Stmt { id: ast::DUMMY_NODE_ID, span: expr.span, kind: ast::StmtKind::Expr(expr) }
     }
 
     pub fn stmt_let(&self, sp: Span, mutbl: bool, ident: Ident, ex: P<ast::Expr>) -> ast::Stmt {
@@ -162,13 +157,9 @@ impl<'a> ExtCtxt<'a> {
             id: ast::DUMMY_NODE_ID,
             span: sp,
             attrs: AttrVec::new(),
-        });
-        ast::Stmt {
-            id: ast::DUMMY_NODE_ID,
-            kind: ast::StmtKind::Local(local),
-            span: sp,
             tokens: None,
-        }
+        });
+        ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Local(local), span: sp }
     }
 
     // Generates `let _: Type;`, which is usually used for type assertions.
@@ -180,17 +171,13 @@ impl<'a> ExtCtxt<'a> {
             id: ast::DUMMY_NODE_ID,
             span,
             attrs: AttrVec::new(),
+            tokens: None,
         });
-        ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Local(local), span, tokens: None }
+        ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Local(local), span }
     }
 
     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,
-            tokens: None,
-        }
+        ast::Stmt { id: ast::DUMMY_NODE_ID, kind: ast::StmtKind::Item(item), span: sp }
     }
 
     pub fn block_expr(&self, expr: P<ast::Expr>) -> P<ast::Block> {
@@ -200,7 +187,6 @@ impl<'a> ExtCtxt<'a> {
                 id: ast::DUMMY_NODE_ID,
                 span: expr.span,
                 kind: ast::StmtKind::Expr(expr),
-                tokens: None,
             }],
         )
     }