Introduce PreexpTokenStream

compiler/rustc_ast/src/ast.rs

@@ -24,7 +24,7 @@ pub use UnsafeSource::*;
 
 use crate::ptr::P;
 use crate::token::{self, CommentKind, DelimToken};
-use crate::tokenstream::{DelimSpan, TokenStream, TokenTree};
+use crate::tokenstream::{DelimSpan, PreexpTokenStream, TokenStream, TokenTree};
 
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::Lrc;
@@ -926,10 +926,8 @@ impl Stmt {
         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,
@@ -973,6 +971,7 @@ pub struct MacCallStmt {
     pub mac: MacCall,
     pub style: MacStmtStyle,
     pub attrs: AttrVec,
+    pub tokens: Option<PreexpTokenStream>,
 }
 
 #[derive(Clone, Copy, PartialEq, Encodable, Decodable, Debug)]
@@ -998,6 +997,7 @@ pub struct Local {
     pub init: Option<P<Expr>>,
     pub span: Span,
     pub attrs: AttrVec,
+    pub tokens: Option<PreexpTokenStream>,
 }
 
 /// An arm of a 'match'.
@@ -1066,7 +1066,7 @@ pub struct Expr {
     pub kind: ExprKind,
     pub span: Span,
     pub attrs: AttrVec,
-    pub tokens: Option<TokenStream>,
+    pub tokens: Option<PreexpTokenStream>,
 }
 
 // `Expr` is used a lot. Make sure it doesn't unintentionally get bigger.
@@ -2436,6 +2436,7 @@ pub struct Attribute {
     /// or the construct this attribute is contained within (inner).
     pub style: AttrStyle,
     pub span: Span,
+    pub tokens: Option<TokenStream>,
 }
 
 #[derive(Clone, Encodable, Decodable, Debug)]
@@ -2582,7 +2583,7 @@ pub struct Item<K = ItemKind> {
     ///
     /// Note that the tokens here do not include the outer attributes, but will
     /// include inner attributes.
-    pub tokens: Option<TokenStream>,
+    pub tokens: Option<PreexpTokenStream>,
 }
 
 impl Item {

compiler/rustc_ast/src/attr/mod.rs

@@ -8,7 +8,7 @@ use crate::ast::{Path, PathSegment};
 use crate::mut_visit::visit_clobber;
 use crate::ptr::P;
 use crate::token::{self, CommentKind, Token};
-use crate::tokenstream::{DelimSpan, TokenStream, TokenTree, TreeAndSpacing};
+use crate::tokenstream::{DelimSpan, PreexpTokenStream, TokenStream, TokenTree, TreeAndSpacing};
 
 use rustc_index::bit_set::GrowableBitSet;
 use rustc_span::source_map::{BytePos, Spanned};
@@ -334,7 +334,7 @@ pub fn mk_attr(style: AttrStyle, path: Path, args: MacArgs, span: Span) -> Attri
 }
 
 pub fn mk_attr_from_item(style: AttrStyle, item: AttrItem, span: Span) -> Attribute {
-    Attribute { kind: AttrKind::Normal(item), id: mk_attr_id(), style, span }
+    Attribute { kind: AttrKind::Normal(item), id: mk_attr_id(), style, span, tokens: None }
 }
 
 /// Returns an inner attribute with the given value and span.
@@ -353,7 +353,13 @@ pub fn mk_doc_comment(
     data: Symbol,
     span: Span,
 ) -> Attribute {
-    Attribute { kind: AttrKind::DocComment(comment_kind, data), id: mk_attr_id(), style, span }
+    Attribute {
+        kind: AttrKind::DocComment(comment_kind, data),
+        id: mk_attr_id(),
+        style,
+        span,
+        tokens: None,
+    }
 }
 
 pub fn list_contains_name(items: &[NestedMetaItem], name: Symbol) -> bool {
@@ -585,6 +591,7 @@ impl NestedMetaItem {
 pub trait HasAttrs: Sized {
     fn attrs(&self) -> &[Attribute];
     fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>));
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut PreexpTokenStream));
 }
 
 impl<T: HasAttrs> HasAttrs for Spanned<T> {
@@ -594,6 +601,9 @@ impl<T: HasAttrs> HasAttrs for Spanned<T> {
     fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>)) {
         self.node.visit_attrs(f);
     }
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut PreexpTokenStream)) {
+        self.node.visit_tokens(f)
+    }
 }
 
 impl HasAttrs for Vec<Attribute> {
@@ -603,6 +613,7 @@ impl HasAttrs for Vec<Attribute> {
     fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>)) {
         f(self)
     }
+    fn visit_tokens(&mut self, _f: impl FnOnce(&mut PreexpTokenStream)) {}
 }
 
 impl HasAttrs for AttrVec {
@@ -616,6 +627,7 @@ impl HasAttrs for AttrVec {
             vec.into()
         });
     }
+    fn visit_tokens(&mut self, _f: impl FnOnce(&mut PreexpTokenStream)) {}
 }
 
 impl<T: HasAttrs + 'static> HasAttrs for P<T> {
@@ -625,14 +637,30 @@ impl<T: HasAttrs + 'static> HasAttrs for P<T> {
     fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>)) {
         (**self).visit_attrs(f);
     }
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut PreexpTokenStream)) {
+        (**self).visit_tokens(f)
+    }
+}
+
+impl<T: HasAttrs> HasAttrs for Option<T> {
+    fn attrs(&self) -> &[Attribute] {
+        self.as_ref().map_or(&[], |inner| inner.attrs())
+    }
+    fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>)) {
+        self.as_mut().map(|inner| inner.visit_attrs(f));
+    }
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut PreexpTokenStream)) {
+        self.as_mut().map(|inner| inner.visit_tokens(f));
+    }
 }
 
 impl HasAttrs for StmtKind {
     fn attrs(&self) -> &[Attribute] {
         match *self {
             StmtKind::Local(ref local) => local.attrs(),
             StmtKind::Expr(ref expr) | StmtKind::Semi(ref expr) => expr.attrs(),
-            StmtKind::Empty | StmtKind::Item(..) => &[],
+            StmtKind::Item(..) => &[],
+            StmtKind::Empty => &[],
             StmtKind::MacCall(ref mac) => mac.attrs.attrs(),
         }
     }
@@ -641,12 +669,24 @@ impl HasAttrs for StmtKind {
         match self {
             StmtKind::Local(local) => local.visit_attrs(f),
             StmtKind::Expr(expr) | StmtKind::Semi(expr) => expr.visit_attrs(f),
-            StmtKind::Empty | StmtKind::Item(..) => {}
+            StmtKind::Item(item) => item.visit_attrs(f),
+            StmtKind::Empty => {}
             StmtKind::MacCall(mac) => {
                 mac.attrs.visit_attrs(f);
             }
         }
     }
+
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut PreexpTokenStream)) {
+        match self {
+            StmtKind::Local(local) => local.visit_tokens(f),
+            StmtKind::Expr(expr) | StmtKind::Semi(expr) => expr.visit_tokens(f),
+            // FIXME: Is this correct?
+            StmtKind::Item(item) => item.visit_tokens(f),
+            StmtKind::Empty => {}
+            StmtKind::MacCall(mac) => mac.attrs.visit_tokens(f),
+        }
+    }
 }
 
 impl HasAttrs for Stmt {
@@ -657,6 +697,9 @@ impl HasAttrs for Stmt {
     fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>)) {
         self.kind.visit_attrs(f);
     }
+    fn visit_tokens(&mut self, f: impl FnOnce(&mut PreexpTokenStream)) {
+        self.kind.visit_tokens(f);
+    }
 }
 
 macro_rules! derive_has_attrs {
@@ -669,11 +712,36 @@ macro_rules! derive_has_attrs {
             fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>)) {
                 self.attrs.visit_attrs(f);
             }
+            fn visit_tokens(&mut self, _f: impl FnOnce(&mut PreexpTokenStream)) {}
+        }
+    )* }
+}
+
+macro_rules! derive_has_attrs_with_tokens {
+    ($($ty:path),*) => { $(
+        impl HasAttrs for $ty {
+            fn attrs(&self) -> &[Attribute] {
+                &self.attrs
+            }
+
+            fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<Attribute>)) {
+                self.attrs.visit_attrs(f);
+            }
+
+            fn visit_tokens(&mut self, f: impl FnOnce(&mut PreexpTokenStream)) {
+                if let Some(tokens) = self.tokens.as_mut() {
+                    f(tokens)
+                }
+            }
         }
     )* }
 }
 
 derive_has_attrs! {
-    Item, Expr, Local, ast::AssocItem, ast::ForeignItem, ast::StructField, ast::Arm,
+    ast::StructField, ast::Arm,
     ast::Field, ast::FieldPat, ast::Variant, ast::Param, GenericParam
 }
+
+derive_has_attrs_with_tokens! {
+    Expr, Item, Local, ast::AssocItem, ast::ForeignItem
+}

compiler/rustc_ast/src/mut_visit.rs

@@ -567,7 +567,7 @@ 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));
@@ -577,7 +577,7 @@ pub fn noop_visit_local<T: MutVisitor>(local: &mut P<Local>, vis: &mut T) {
 }
 
 pub fn noop_visit_attribute<T: MutVisitor>(attr: &mut Attribute, vis: &mut T) {
-    let Attribute { kind, id: _, style: _, span } = attr;
+    let Attribute { kind, id: _, style: _, span, tokens: _ } = attr;
     match kind {
         AttrKind::Normal(AttrItem { path, args, tokens: _ }) => {
             vis.visit_path(path);
@@ -1311,7 +1311,7 @@ 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(mac_);
             visit_thin_attrs(attrs, vis);
             smallvec![StmtKind::MacCall(mac)]

compiler/rustc_ast/src/tokenstream.rs

@@ -403,6 +403,10 @@ impl Cursor {
     pub fn look_ahead(&self, n: usize) -> Option<TokenTree> {
         self.stream.0[self.index..].get(n).map(|(tree, _)| tree.clone())
     }
+
+    pub fn index(&self) -> usize {
+        self.index
+    }
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Encodable, Decodable, HashStable_Generic)]
@@ -428,3 +432,85 @@ impl DelimSpan {
         self.open.with_hi(self.close.hi())
     }
 }
+
+#[derive(Clone, Debug, Default, Encodable, Decodable)]
+pub struct PreexpTokenStream(pub Lrc<Vec<(PreexpTokenTree, Spacing)>>);
+
+#[derive(Clone, Debug, Encodable, Decodable)]
+pub enum PreexpTokenTree {
+    Token(Token),
+    Delimited(DelimSpan, DelimToken, PreexpTokenStream),
+    OuterAttributes(AttributesData),
+}
+
+impl PreexpTokenStream {
+    pub fn new(tokens: Vec<(PreexpTokenTree, Spacing)>) -> PreexpTokenStream {
+        PreexpTokenStream(Lrc::new(tokens))
+    }
+
+    pub fn replace_attributes(&mut self, f: impl FnOnce(&mut AttributesData)) {
+        if let &[(PreexpTokenTree::OuterAttributes(ref data), joint)] = &**self.0 {
+            let mut data = data.clone();
+            f(&mut data);
+            *self = PreexpTokenStream::new(vec![(PreexpTokenTree::OuterAttributes(data), joint)]);
+        } else {
+            panic!("Expected a single PreexpTokenTree::OuterAttributes, found {:?}", self);
+        }
+    }
+
+    pub fn to_tokenstream(self) -> TokenStream {
+        let trees: Vec<_> = self
+            .0
+            .iter()
+            .flat_map(|tree| match &tree.0 {
+                PreexpTokenTree::Token(inner) => {
+                    smallvec![(TokenTree::Token(inner.clone()), tree.1)].into_iter()
+                }
+                PreexpTokenTree::Delimited(span, delim, stream) => smallvec![(
+                    TokenTree::Delimited(*span, *delim, stream.clone().to_tokenstream()),
+                    tree.1,
+                )]
+                .into_iter(),
+                PreexpTokenTree::OuterAttributes(data) => {
+                    let flat: SmallVec<[_; 1]> = data
+                        .attrs
+                        .iter()
+                        .filter(|attr| attr.style == crate::AttrStyle::Outer)
+                        .flat_map(|attr| {
+                            attr.tokens.as_ref().expect("Missing tokens").0.iter().cloned()
+                        })
+                        .chain(data.tokens.clone().to_tokenstream().0.iter().cloned())
+                        .collect();
+                    flat.into_iter()
+                }
+            })
+            .collect();
+        TokenStream::new(trees)
+    }
+
+    pub fn from_tokenstream(stream: TokenStream) -> PreexpTokenStream {
+        let trees: Vec<_> = stream
+            .0
+            .iter()
+            .cloned()
+            .map(|tree| {
+                let new_tree = match tree.0 {
+                    TokenTree::Token(token) => PreexpTokenTree::Token(token),
+                    TokenTree::Delimited(sp, delim, inner) => PreexpTokenTree::Delimited(
+                        sp,
+                        delim,
+                        PreexpTokenStream::from_tokenstream(inner),
+                    ),
+                };
+                (new_tree, tree.1)
+            })
+            .collect();
+        PreexpTokenStream::new(trees)
+    }
+}
+
+#[derive(Clone, Debug, Encodable, Decodable)]
+pub struct AttributesData {
+    pub attrs: Vec<crate::Attribute>,
+    pub tokens: PreexpTokenStream,
+}

compiler/rustc_ast/src/visit.rs

@@ -692,7 +692,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(mac);
             for attr in attrs.iter() {
                 visitor.visit_attribute(attr);

compiler/rustc_ast_lowering/src/lib.rs

@@ -972,7 +972,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             AttrKind::DocComment(comment_kind, data) => AttrKind::DocComment(comment_kind, data),
         };
 
-        Attribute { kind, id: attr.id, style: attr.style, span: attr.span }
+        Attribute { kind, id: attr.id, style: attr.style, span: attr.span, tokens: None }
     }
 
     fn lower_mac_args(&mut self, args: &MacArgs) -> MacArgs {

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 }
 }