[WIP] Implement token-based handling of attributes

compiler/rustc_ast/src/ast.rs

@@ -905,26 +905,6 @@ pub struct Stmt {
 }
 
 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 has_trailing_semicolon(&self) -> bool {
         match &self.kind {
             StmtKind::Semi(_) => true,
@@ -979,6 +959,38 @@ pub enum StmtKind {
     MacCall(P<MacCallStmt>),
 }
 
+impl StmtKind {
+    pub fn tokens(&self) -> Option<&LazyTokenStream> {
+        match self {
+            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 {
+            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 {
+            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,
+        }
+    }
+}
+
 #[derive(Clone, Encodable, Decodable, Debug)]
 pub struct MacCallStmt {
     pub mac: MacCall,

compiler/rustc_ast/src/ast_like.rs

@@ -1,35 +1,42 @@
 use super::ptr::P;
-use super::tokenstream::LazyTokenStream;
+use super::tokenstream::{AttributesData, LazyTokenStream};
 use super::{Arm, Field, FieldPat, GenericParam, Param, StructField, 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 rustc_span::sym;
 
 /// 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 {
+    const SUPPORTS_INNER_ATTRS: bool;
     fn attrs(&self) -> &[Attribute];
     fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>));
     /// Called by `Parser::collect_tokens` to store the collected
     /// tokens inside an AST node
-    fn finalize_tokens(&mut self, _tokens: LazyTokenStream) {
+    fn finalize_tokens(&mut self, _tokens: LazyTokenStream) -> Option<AttributesData> {
         // This default impl makes this trait easier to implement
         // in tools like `rust-analyzer`
         panic!("`finalize_tokens` is not supported!")
     }
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut Option<LazyTokenStream>));
 }
 
 impl<T: AstLike + 'static> AstLike for P<T> {
+    const SUPPORTS_INNER_ATTRS: bool = T::SUPPORTS_INNER_ATTRS;
     fn attrs(&self) -> &[Attribute] {
         (**self).attrs()
     }
     fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>)) {
         (**self).visit_attrs(f);
     }
-    fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+    fn finalize_tokens(&mut self, tokens: LazyTokenStream) -> Option<AttributesData> {
         (**self).finalize_tokens(tokens)
     }
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut Option<LazyTokenStream>)) {
+        (**self).visit_tokens(f);
+    }
 }
 
 fn visit_attrvec(attrs: &mut AttrVec, f: impl FnOnce(&mut Vec<Attribute>)) {
@@ -41,6 +48,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_INNER_ATTRS: bool = true;
+
     fn attrs(&self) -> &[Attribute] {
         match *self {
             StmtKind::Local(ref local) => local.attrs(),
@@ -60,53 +71,84 @@ impl AstLike for StmtKind {
             StmtKind::MacCall(mac) => visit_attrvec(&mut mac.attrs, f),
         }
     }
-    fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
-        let stmt_tokens = match self {
-            StmtKind::Local(ref mut local) => &mut local.tokens,
-            StmtKind::Item(ref mut item) => &mut item.tokens,
-            StmtKind::Expr(ref mut expr) | StmtKind::Semi(ref mut expr) => &mut expr.tokens,
-            StmtKind::Empty => return,
-            StmtKind::MacCall(ref mut mac) => &mut mac.tokens,
+    fn finalize_tokens(&mut self, tokens: LazyTokenStream) -> Option<AttributesData> {
+        match self {
+            StmtKind::Local(ref mut local) => local.finalize_tokens(tokens),
+            StmtKind::MacCall(ref mut mac) => mac.finalize_tokens(tokens),
+            StmtKind::Expr(ref mut expr) | StmtKind::Semi(ref mut expr) => {
+                expr.finalize_tokens(tokens)
+            }
+            StmtKind::Item(ref mut item) => item.finalize_tokens(tokens),
+            StmtKind::Empty => None,
+        }
+    }
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut Option<LazyTokenStream>)) {
+        let tokens = match self {
+            StmtKind::Local(ref mut local) => Some(&mut local.tokens),
+            StmtKind::Item(ref mut item) => Some(&mut item.tokens),
+            StmtKind::Expr(ref mut expr) | StmtKind::Semi(ref mut expr) => Some(&mut expr.tokens),
+            StmtKind::Empty => None,
+            StmtKind::MacCall(ref mut mac) => Some(&mut mac.tokens),
         };
-        if stmt_tokens.is_none() {
-            *stmt_tokens = Some(tokens);
+        if let Some(tokens) = tokens {
+            f(tokens);
         }
     }
 }
 
 impl AstLike for Stmt {
+    const SUPPORTS_INNER_ATTRS: bool = StmtKind::SUPPORTS_INNER_ATTRS;
+
     fn attrs(&self) -> &[Attribute] {
         self.kind.attrs()
     }
 
     fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>)) {
         self.kind.visit_attrs(f);
     }
-    fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+    fn finalize_tokens(&mut self, tokens: LazyTokenStream) -> Option<AttributesData> {
         self.kind.finalize_tokens(tokens)
     }
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut Option<LazyTokenStream>)) {
+        self.kind.visit_tokens(f)
+    }
 }
 
 impl AstLike for Attribute {
+    const SUPPORTS_INNER_ATTRS: bool = false;
+
     fn attrs(&self) -> &[Attribute] {
         &[]
     }
     fn visit_attrs(&mut self, _f: impl FnOnce(&mut Vec<Attribute>)) {}
-    fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+    fn finalize_tokens(&mut self, tokens: LazyTokenStream) -> Option<AttributesData> {
         match &mut self.kind {
             AttrKind::Normal(_, attr_tokens) => {
                 if attr_tokens.is_none() {
                     *attr_tokens = Some(tokens);
                 }
+                None
             }
             AttrKind::DocComment(..) => {
                 panic!("Called finalize_tokens on doc comment attr {:?}", self)
             }
         }
     }
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut Option<LazyTokenStream>)) {
+        match &mut self.kind {
+            AttrKind::Normal(_, attr_tokens) => {
+                f(attr_tokens);
+            }
+            AttrKind::DocComment(..) => {
+                panic!("Called visit_tokens on doc comment attr {:?}", self)
+            }
+        }
+    }
 }
 
 impl<T: AstLike> AstLike for Option<T> {
+    const SUPPORTS_INNER_ATTRS: bool = T::SUPPORTS_INNER_ATTRS;
+
     fn attrs(&self) -> &[Attribute] {
         self.as_ref().map(|inner| inner.attrs()).unwrap_or(&[])
     }
@@ -115,13 +157,25 @@ impl<T: AstLike> AstLike for Option<T> {
             inner.visit_attrs(f);
         }
     }
-    fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+    fn finalize_tokens(&mut self, tokens: LazyTokenStream) -> Option<AttributesData> {
+        self.as_mut().and_then(|inner| inner.finalize_tokens(tokens))
+    }
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut Option<LazyTokenStream>)) {
         if let Some(inner) = self {
-            inner.finalize_tokens(tokens);
+            inner.visit_tokens(f);
         }
     }
 }
 
+// NOTE: Builtin attributes like `cfg` and `cfg_attr` cannot be renamed via imports.
+// Therefore, the absence of a literal `cfg` or `cfg_attr` guarantees that
+// we don't need to do any eager expansion.
+pub fn has_cfg_or_cfg_any(attrs: &[Attribute]) -> bool {
+    attrs.iter().any(|attr| {
+        attr.ident().map_or(false, |ident| ident.name == sym::cfg || ident.name == sym::cfg_attr)
+    })
+}
+
 /// Helper trait for the macros below. Abstracts over
 /// the two types of attribute fields that AST nodes
 /// may have (`Vec<Attribute>` or `AttrVec`)
@@ -142,8 +196,13 @@ impl VecOrAttrVec for AttrVec {
 }
 
 macro_rules! derive_has_tokens_and_attrs {
-    ($($ty:path),*) => { $(
+    (
+        const SUPPORTS_INNER_ATTRS: bool = $inner_attrs:literal;
+        $($ty:path),*
+    ) => { $(
         impl AstLike for $ty {
+            const SUPPORTS_INNER_ATTRS: bool = $inner_attrs;
+
             fn attrs(&self) -> &[Attribute] {
                 &self.attrs
             }
@@ -152,19 +211,31 @@ macro_rules! derive_has_tokens_and_attrs {
                 VecOrAttrVec::visit(&mut self.attrs, f)
             }
 
-            fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+            fn finalize_tokens(&mut self, tokens: LazyTokenStream) -> Option<AttributesData> {
                 if self.tokens.is_none() {
                     self.tokens = Some(tokens);
                 }
 
+                if has_cfg_or_cfg_any(&self.attrs) {
+                    Some(AttributesData { attrs: self.attrs.clone().into(), tokens: self.tokens.clone().unwrap() })
+                } else {
+                    None
+                }
             }
+
+            fn visit_tokens(&mut self, f: impl FnOnce(&mut Option<LazyTokenStream>)) {
+                f(&mut self.tokens)
+            }
+
         }
     )* }
 }
 
 macro_rules! derive_has_attrs_no_tokens {
     ($($ty:path),*) => { $(
         impl AstLike for $ty {
+            const SUPPORTS_INNER_ATTRS: bool = false;
+
             fn attrs(&self) -> &[Attribute] {
                 &self.attrs
             }
@@ -173,35 +244,56 @@ macro_rules! derive_has_attrs_no_tokens {
                 VecOrAttrVec::visit(&mut self.attrs, f)
             }
 
-            fn finalize_tokens(&mut self, _tokens: LazyTokenStream) {}
+            fn finalize_tokens(&mut self, tokens: LazyTokenStream) -> Option<AttributesData> {
+                if has_cfg_or_cfg_any(&self.attrs) {
+                    Some(AttributesData { attrs: self.attrs.clone().into(), tokens })
+                } else {
+                    None
+                }
+            }
+            fn visit_tokens(&mut self, _f: impl FnOnce(&mut Option<LazyTokenStream>)) {}
+
         }
     )* }
 }
 
 macro_rules! derive_has_tokens_no_attrs {
     ($($ty:path),*) => { $(
         impl AstLike for $ty {
+            const SUPPORTS_INNER_ATTRS: bool = false;
+
             fn attrs(&self) -> &[Attribute] {
                 &[]
             }
 
             fn visit_attrs(&mut self, _f: impl FnOnce(&mut Vec<Attribute>)) {
             }
 
-            fn finalize_tokens(&mut self, tokens: LazyTokenStream) {
+            fn finalize_tokens(&mut self, tokens: LazyTokenStream) -> Option<AttributesData> {
                 if self.tokens.is_none() {
                     self.tokens = Some(tokens);
                 }
-
+                None
+            }
+            fn visit_tokens(&mut self, f: impl FnOnce(&mut Option<LazyTokenStream>)) {
+                f(&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_INNER_ATTRS: bool = true;
+    Item, AssocItem
+}
+
 derive_has_tokens_and_attrs! {
-    Item, Expr, Local, AssocItem, ForeignItem
+    const SUPPORTS_INNER_ATTRS: bool = false;
+    ForeignItem, Expr, Local, MacCallStmt
 }
 
 // These ast nodes only support inert attributes, so they don't

compiler/rustc_ast/src/attr/mod.rs

@@ -6,7 +6,10 @@ 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::{
+    DelimSpan, LazyTokenStream, PreexpTokenStream, PreexpTokenTree, Spacing, TokenStream,
+    TokenTree, TreeAndSpacing,
+};
 
 use rustc_index::bit_set::GrowableBitSet;
 use rustc_span::source_map::BytePos;
@@ -309,14 +312,18 @@ impl Attribute {
         }
     }
 
-    pub fn tokens(&self) -> TokenStream {
+    pub fn tokens(&self) -> PreexpTokenStream {
         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) => PreexpTokenStream::from((
+                PreexpTokenTree::Token(Token::new(
+                    token::DocComment(comment_kind, self.style, data),
+                    self.span,
+                )),
+                Spacing::Alone,
             )),
         }
     }