Implement token-based handling of attributes during expansion

compiler/rustc_ast/src/ast_like.rs

@@ -1,20 +1,32 @@
 use super::ptr::P;
+use super::token::Nonterminal;
 use super::tokenstream::LazyTokenStream;
 use super::{Arm, ExprField, FieldDef, GenericParam, Param, PatField, Variant};
-use super::{AssocItem, Expr, ForeignItem, Item, Local};
+use super::{AssocItem, Expr, ForeignItem, Item, Local, MacCallStmt};
 use super::{AttrItem, AttrKind, Block, Pat, Path, Ty, Visibility};
 use super::{AttrVec, Attribute, Stmt, StmtKind};
 
+use std::fmt::Debug;
+
 /// An `AstLike` represents an AST node (or some wrapper around
 /// and AST node) which stores some combination of attributes
 /// and tokens.
-pub trait AstLike: Sized {
+pub trait AstLike: Sized + Debug {
+    /// This is `true` if this `AstLike` might support 'custom' (proc-macro) inner
+    /// attributes. Attributes like `#![cfg]` and `#![cfg_attr]` are not
+    /// considered 'custom' attributes
+    ///
+    /// If this is `false`, then this `AstLike` definitely does
+    /// not support 'custom' inner attributes, which enables some optimizations
+    /// during token collection.
+    const SUPPORTS_CUSTOM_INNER_ATTRS: bool;
     fn attrs(&self) -> &[Attribute];
     fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>));
     fn tokens_mut(&mut self) -> Option<&mut Option<LazyTokenStream>>;
 }
 
 impl<T: AstLike + 'static> AstLike for P<T> {
+    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = T::SUPPORTS_CUSTOM_INNER_ATTRS;
     fn attrs(&self) -> &[Attribute] {
         (**self).attrs()
     }
@@ -26,6 +38,55 @@ impl<T: AstLike + 'static> AstLike for P<T> {
     }
 }
 
+impl AstLike for crate::token::Nonterminal {
+    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = true;
+    fn attrs(&self) -> &[Attribute] {
+        match self {
+            Nonterminal::NtItem(item) => item.attrs(),
+            Nonterminal::NtStmt(stmt) => stmt.attrs(),
+            Nonterminal::NtExpr(expr) | Nonterminal::NtLiteral(expr) => expr.attrs(),
+            Nonterminal::NtPat(_)
+            | Nonterminal::NtTy(_)
+            | Nonterminal::NtMeta(_)
+            | Nonterminal::NtPath(_)
+            | Nonterminal::NtVis(_)
+            | Nonterminal::NtTT(_)
+            | Nonterminal::NtBlock(_)
+            | Nonterminal::NtIdent(..)
+            | Nonterminal::NtLifetime(_) => &[],
+        }
+    }
+    fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>)) {
+        match self {
+            Nonterminal::NtItem(item) => item.visit_attrs(f),
+            Nonterminal::NtStmt(stmt) => stmt.visit_attrs(f),
+            Nonterminal::NtExpr(expr) | Nonterminal::NtLiteral(expr) => expr.visit_attrs(f),
+            Nonterminal::NtPat(_)
+            | Nonterminal::NtTy(_)
+            | Nonterminal::NtMeta(_)
+            | Nonterminal::NtPath(_)
+            | Nonterminal::NtVis(_)
+            | Nonterminal::NtTT(_)
+            | Nonterminal::NtBlock(_)
+            | Nonterminal::NtIdent(..)
+            | Nonterminal::NtLifetime(_) => {}
+        }
+    }
+    fn tokens_mut(&mut self) -> Option<&mut Option<LazyTokenStream>> {
+        match self {
+            Nonterminal::NtItem(item) => item.tokens_mut(),
+            Nonterminal::NtStmt(stmt) => stmt.tokens_mut(),
+            Nonterminal::NtExpr(expr) | Nonterminal::NtLiteral(expr) => expr.tokens_mut(),
+            Nonterminal::NtPat(pat) => pat.tokens_mut(),
+            Nonterminal::NtTy(ty) => ty.tokens_mut(),
+            Nonterminal::NtMeta(attr_item) => attr_item.tokens_mut(),
+            Nonterminal::NtPath(path) => path.tokens_mut(),
+            Nonterminal::NtVis(vis) => vis.tokens_mut(),
+            _ => panic!("Called tokens_mut on {:?}", self),
+        }
+    }
+}
+
 fn visit_attrvec(attrs: &mut AttrVec, f: impl FnOnce(&mut Vec<Attribute>)) {
     crate::mut_visit::visit_clobber(attrs, |attrs| {
         let mut vec = attrs.into();
@@ -35,6 +96,10 @@ fn visit_attrvec(attrs: &mut AttrVec, f: impl FnOnce(&mut Vec<Attribute>)) {
 }
 
 impl AstLike for StmtKind {
+    // This might be an `StmtKind::Item`, which contains
+    // an item that supports inner attrs
+    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = true;
+
     fn attrs(&self) -> &[Attribute] {
         match self {
             StmtKind::Local(local) => local.attrs(),
@@ -66,6 +131,8 @@ impl AstLike for StmtKind {
 }
 
 impl AstLike for Stmt {
+    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = StmtKind::SUPPORTS_CUSTOM_INNER_ATTRS;
+
     fn attrs(&self) -> &[Attribute] {
         self.kind.attrs()
     }
@@ -79,6 +146,8 @@ impl AstLike for Stmt {
 }
 
 impl AstLike for Attribute {
+    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = false;
+
     fn attrs(&self) -> &[Attribute] {
         &[]
     }
@@ -94,6 +163,8 @@ impl AstLike for Attribute {
 }
 
 impl<T: AstLike> AstLike for Option<T> {
+    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = T::SUPPORTS_CUSTOM_INNER_ATTRS;
+
     fn attrs(&self) -> &[Attribute] {
         self.as_ref().map(|inner| inner.attrs()).unwrap_or(&[])
     }
@@ -127,8 +198,13 @@ impl VecOrAttrVec for AttrVec {
 }
 
 macro_rules! derive_has_tokens_and_attrs {
-    ($($ty:path),*) => { $(
+    (
+        const SUPPORTS_CUSTOM_INNER_ATTRS: bool = $inner_attrs:literal;
+        $($ty:path),*
+    ) => { $(
         impl AstLike for $ty {
+            const SUPPORTS_CUSTOM_INNER_ATTRS: bool = $inner_attrs;
+
             fn attrs(&self) -> &[Attribute] {
                 &self.attrs
             }
@@ -140,13 +216,16 @@ macro_rules! derive_has_tokens_and_attrs {
             fn tokens_mut(&mut self) -> Option<&mut Option<LazyTokenStream>> {
                 Some(&mut self.tokens)
             }
+
         }
     )* }
 }
 
 macro_rules! derive_has_attrs_no_tokens {
     ($($ty:path),*) => { $(
         impl AstLike for $ty {
+            const SUPPORTS_CUSTOM_INNER_ATTRS: bool = false;
+
             fn attrs(&self) -> &[Attribute] {
                 &self.attrs
             }
@@ -165,23 +244,32 @@ macro_rules! derive_has_attrs_no_tokens {
 macro_rules! derive_has_tokens_no_attrs {
     ($($ty:path),*) => { $(
         impl AstLike for $ty {
+            const SUPPORTS_CUSTOM_INNER_ATTRS: bool = false;
+
             fn attrs(&self) -> &[Attribute] {
                 &[]
             }
 
             fn visit_attrs(&mut self, _f: impl FnOnce(&mut Vec<Attribute>)) {}
-
             fn tokens_mut(&mut self) -> Option<&mut Option<LazyTokenStream>> {
                 Some(&mut self.tokens)
             }
         }
     )* }
 }
 
-// These AST nodes support both inert and active
-// attributes, so they also have tokens.
+// These ast nodes support both active and inert attributes,
+// so they have tokens collected to pass to proc macros
+derive_has_tokens_and_attrs! {
+    // Both `Item` and `AssocItem` can have bodies, which
+    // can contain inner attributes
+    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = true;
+    Item, AssocItem, ForeignItem
+}
+
 derive_has_tokens_and_attrs! {
-    Item, Expr, Local, AssocItem, ForeignItem
+    const SUPPORTS_CUSTOM_INNER_ATTRS: bool = false;
+    Local, MacCallStmt, Expr
 }
 
 // These ast nodes only support inert attributes, so they don't

compiler/rustc_ast/src/attr/mod.rs

@@ -6,7 +6,9 @@ use crate::ast::{Lit, LitKind};
 use crate::ast::{MacArgs, MacDelimiter, MetaItem, MetaItemKind, NestedMetaItem};
 use crate::ast::{Path, PathSegment};
 use crate::token::{self, CommentKind, Token};
-use crate::tokenstream::{DelimSpan, LazyTokenStream, TokenStream, TokenTree, TreeAndSpacing};
+use crate::tokenstream::{AttrAnnotatedTokenStream, AttrAnnotatedTokenTree};
+use crate::tokenstream::{DelimSpan, Spacing, TokenTree, TreeAndSpacing};
+use crate::tokenstream::{LazyTokenStream, TokenStream};
 
 use rustc_index::bit_set::GrowableBitSet;
 use rustc_span::source_map::BytePos;
@@ -268,14 +270,18 @@ impl Attribute {
         }
     }
 
-    pub fn tokens(&self) -> TokenStream {
+    pub fn tokens(&self) -> AttrAnnotatedTokenStream {
         match self.kind {
             AttrKind::Normal(_, ref tokens) => tokens
                 .as_ref()
                 .unwrap_or_else(|| panic!("attribute is missing tokens: {:?}", self))
                 .create_token_stream(),
-            AttrKind::DocComment(comment_kind, data) => TokenStream::from(TokenTree::Token(
-                Token::new(token::DocComment(comment_kind, self.style, data), self.span),
+            AttrKind::DocComment(comment_kind, data) => AttrAnnotatedTokenStream::from((
+                AttrAnnotatedTokenTree::Token(Token::new(
+                    token::DocComment(comment_kind, self.style, data),
+                    self.span,
+                )),
+                Spacing::Alone,
             )),
         }
     }

compiler/rustc_ast/src/mut_visit.rs

@@ -630,6 +630,33 @@ pub fn noop_flat_map_param<T: MutVisitor>(mut param: Param, vis: &mut T) -> Smal
     smallvec![param]
 }
 
+// No `noop_` prefix because there isn't a corresponding method in `MutVisitor`.
+pub fn visit_attr_annotated_tt<T: MutVisitor>(tt: &mut AttrAnnotatedTokenTree, vis: &mut T) {
+    match tt {
+        AttrAnnotatedTokenTree::Token(token) => {
+            visit_token(token, vis);
+        }
+        AttrAnnotatedTokenTree::Delimited(DelimSpan { open, close }, _delim, tts) => {
+            vis.visit_span(open);
+            vis.visit_span(close);
+            visit_attr_annotated_tts(tts, vis);
+        }
+        AttrAnnotatedTokenTree::Attributes(data) => {
+            for attr in &mut *data.attrs {
+                match &mut attr.kind {
+                    AttrKind::Normal(_, attr_tokens) => {
+                        visit_lazy_tts(attr_tokens, vis);
+                    }
+                    AttrKind::DocComment(..) => {
+                        vis.visit_span(&mut attr.span);
+                    }
+                }
+            }
+            visit_lazy_tts_opt_mut(Some(&mut data.tokens), vis);
+        }
+    }
+}
+
 // No `noop_` prefix because there isn't a corresponding method in `MutVisitor`.
 pub fn visit_tt<T: MutVisitor>(tt: &mut TokenTree, vis: &mut T) {
     match tt {
@@ -652,16 +679,30 @@ pub fn visit_tts<T: MutVisitor>(TokenStream(tts): &mut TokenStream, vis: &mut T)
     }
 }
 
-pub fn visit_lazy_tts<T: MutVisitor>(lazy_tts: &mut Option<LazyTokenStream>, vis: &mut T) {
+pub fn visit_attr_annotated_tts<T: MutVisitor>(
+    AttrAnnotatedTokenStream(tts): &mut AttrAnnotatedTokenStream,
+    vis: &mut T,
+) {
+    if vis.token_visiting_enabled() && !tts.is_empty() {
+        let tts = Lrc::make_mut(tts);
+        visit_vec(tts, |(tree, _is_joint)| visit_attr_annotated_tt(tree, vis));
+    }
+}
+
+pub fn visit_lazy_tts_opt_mut<T: MutVisitor>(lazy_tts: Option<&mut LazyTokenStream>, vis: &mut T) {
     if vis.token_visiting_enabled() {
-        visit_opt(lazy_tts, |lazy_tts| {
+        if let Some(lazy_tts) = lazy_tts {
             let mut tts = lazy_tts.create_token_stream();
-            visit_tts(&mut tts, vis);
+            visit_attr_annotated_tts(&mut tts, vis);
             *lazy_tts = LazyTokenStream::new(tts);
-        })
+        }
     }
 }
 
+pub fn visit_lazy_tts<T: MutVisitor>(lazy_tts: &mut Option<LazyTokenStream>, vis: &mut T) {
+    visit_lazy_tts_opt_mut(lazy_tts.as_mut(), vis);
+}
+
 // No `noop_` prefix because there isn't a corresponding method in `MutVisitor`.
 // Applies ident visitor if it's an ident; applies other visits to interpolated nodes.
 // In practice the ident part is not actually used by specific visitors right now,