Auto merge of #72287 - Aaron1011:feature/min-token-collect, r=petrochenkov

Store tokens inside `ast::Expr`

This is a smaller version of #70091.

We now store captured tokens inside `ast::Expr`, which allows us to avoid some reparsing in `nt_to_tokenstream`. To try to mitigate the performance impact, we only collect tokens when we've seen an outer attribute.

This makes progress towards solving #43081. There are still many things left to do:

* Collect tokens for other AST items.
* Come up with a way to handle inner attributes (we need to be collecting tokens by the time we encounter them)
* Avoid re-parsing when a `#[cfg]` attr is used.

However, this is enough to fix spans for a simple example, which I've included as a test case.

16 files changed, 97 insertions, 18 deletions

diff --git a/src/librustc_ast/ast.rs b/src/librustc_ast/ast.rs
index 7ff83572507..30bb5c0bffa 100644
--- a/src/librustc_ast/ast.rs
+++ b/src/librustc_ast/ast.rs
@@ -1006,11 +1006,12 @@ pub struct Expr {
     pub kind: ExprKind,
     pub span: Span,
     pub attrs: AttrVec,
+    pub tokens: Option<TokenStream>,
 }
 
 // `Expr` is used a lot. Make sure it doesn't unintentionally get bigger.
 #[cfg(target_arch = "x86_64")]
-rustc_data_structures::static_assert_size!(Expr, 96);
+rustc_data_structures::static_assert_size!(Expr, 104);
 
 impl Expr {
     /// Returns `true` if this expression would be valid somewhere that expects a value;
diff --git a/src/librustc_ast/mut_visit.rs b/src/librustc_ast/mut_visit.rs
index 2c575c3e288..7ececb814a6 100644
--- a/src/librustc_ast/mut_visit.rs
+++ b/src/librustc_ast/mut_visit.rs
@@ -1095,7 +1095,10 @@ pub fn noop_visit_anon_const<T: MutVisitor>(AnonConst { id, value }: &mut AnonCo
     vis.visit_expr(value);
 }
 
-pub fn noop_visit_expr<T: MutVisitor>(Expr { kind, id, span, attrs }: &mut Expr, vis: &mut T) {
+pub fn noop_visit_expr<T: MutVisitor>(
+    Expr { kind, id, span, attrs, tokens: _ }: &mut Expr,
+    vis: &mut T,
+) {
     match kind {
         ExprKind::Box(expr) => vis.visit_expr(expr),
         ExprKind::Array(exprs) => visit_exprs(exprs, vis),
diff --git a/src/librustc_ast_lowering/lib.rs b/src/librustc_ast_lowering/lib.rs
index 5d98fdeeaf9..3aab54ea909 100644
--- a/src/librustc_ast_lowering/lib.rs
+++ b/src/librustc_ast_lowering/lib.rs
@@ -1126,6 +1126,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
                                 kind: ExprKind::Path(qself.clone(), path.clone()),
                                 span: ty.span,
                                 attrs: AttrVec::new(),
+                                tokens: None,
                             };
 
                             let ct = self.with_new_scopes(|this| hir::AnonConst {
diff --git a/src/librustc_builtin_macros/asm.rs b/src/librustc_builtin_macros/asm.rs
index c2973924897..224b52b239f 100644
--- a/src/librustc_builtin_macros/asm.rs
+++ b/src/librustc_builtin_macros/asm.rs
@@ -519,6 +519,7 @@ fn expand_preparsed_asm(ecx: &mut ExtCtxt<'_>, sp: Span, args: AsmArgs) -> P<ast
         kind: ast::ExprKind::InlineAsm(inline_asm),
         span: sp,
         attrs: ast::AttrVec::new(),
+        tokens: None,
     })
 }
 
diff --git a/src/librustc_builtin_macros/concat_idents.rs b/src/librustc_builtin_macros/concat_idents.rs
index 8a1741c0654..fdf05ac3880 100644
--- a/src/librustc_builtin_macros/concat_idents.rs
+++ b/src/librustc_builtin_macros/concat_idents.rs
@@ -52,6 +52,7 @@ pub fn expand_concat_idents<'cx>(
                 kind: ast::ExprKind::Path(None, ast::Path::from_ident(self.ident)),
                 span: self.ident.span,
                 attrs: ast::AttrVec::new(),
+                tokens: None,
             }))
         }
 
diff --git a/src/librustc_builtin_macros/llvm_asm.rs b/src/librustc_builtin_macros/llvm_asm.rs
index e12fcd98f9d..0f4efc153b9 100644
--- a/src/librustc_builtin_macros/llvm_asm.rs
+++ b/src/librustc_builtin_macros/llvm_asm.rs
@@ -61,6 +61,7 @@ pub fn expand_llvm_asm<'cx>(
         kind: ast::ExprKind::LlvmInlineAsm(P(inline_asm)),
         span: cx.with_def_site_ctxt(sp),
         attrs: ast::AttrVec::new(),
+        tokens: None,
     }))
 }
 
diff --git a/src/librustc_expand/base.rs b/src/librustc_expand/base.rs
index 0137080938f..649aac488fc 100644
--- a/src/librustc_expand/base.rs
+++ b/src/librustc_expand/base.rs
@@ -594,6 +594,7 @@ impl DummyResult {
             kind: if is_error { ast::ExprKind::Err } else { ast::ExprKind::Tup(Vec::new()) },
             span: sp,
             attrs: ast::AttrVec::new(),
+            tokens: None,
         })
     }
 
diff --git a/src/librustc_expand/build.rs b/src/librustc_expand/build.rs
index be2c52a85eb..6185e014d3c 100644
--- a/src/librustc_expand/build.rs
+++ b/src/librustc_expand/build.rs
@@ -70,7 +70,13 @@ impl<'a> ExtCtxt<'a> {
     pub fn anon_const(&self, span: Span, kind: ast::ExprKind) -> ast::AnonConst {
         ast::AnonConst {
             id: ast::DUMMY_NODE_ID,
-            value: P(ast::Expr { id: ast::DUMMY_NODE_ID, kind, span, attrs: AttrVec::new() }),
+            value: P(ast::Expr {
+                id: ast::DUMMY_NODE_ID,
+                kind,
+                span,
+                attrs: AttrVec::new(),
+                tokens: None,
+            }),
         }
     }
 
@@ -205,7 +211,7 @@ impl<'a> ExtCtxt<'a> {
     }
 
     pub fn expr(&self, span: Span, kind: ast::ExprKind) -> P<ast::Expr> {
-        P(ast::Expr { id: ast::DUMMY_NODE_ID, kind, span, attrs: AttrVec::new() })
+        P(ast::Expr { id: ast::DUMMY_NODE_ID, kind, span, attrs: AttrVec::new(), tokens: None })
     }
 
     pub fn expr_path(&self, path: ast::Path) -> P<ast::Expr> {
diff --git a/src/librustc_expand/placeholders.rs b/src/librustc_expand/placeholders.rs
index 23f7a5b28fe..b4ffd714fef 100644
--- a/src/librustc_expand/placeholders.rs
+++ b/src/librustc_expand/placeholders.rs
@@ -34,6 +34,7 @@ pub fn placeholder(
             span,
             attrs: ast::AttrVec::new(),
             kind: ast::ExprKind::MacCall(mac_placeholder()),
+            tokens: None,
         })
     };
     let ty = || P(ast::Ty { id, kind: ast::TyKind::MacCall(mac_placeholder()), span });
diff --git a/src/librustc_interface/util.rs b/src/librustc_interface/util.rs
index a15da94a215..7eaaff05fb5 100644
--- a/src/librustc_interface/util.rs
+++ b/src/librustc_interface/util.rs
@@ -713,6 +713,7 @@ impl<'a> MutVisitor for ReplaceBodyWithLoop<'a, '_> {
                 kind: ast::ExprKind::Block(P(b), None),
                 span: rustc_span::DUMMY_SP,
                 attrs: AttrVec::new(),
+                tokens: None,
             });
 
             ast::Stmt {
@@ -728,6 +729,7 @@ impl<'a> MutVisitor for ReplaceBodyWithLoop<'a, '_> {
             id: self.resolver.next_node_id(),
             span: rustc_span::DUMMY_SP,
             attrs: AttrVec::new(),
+            tokens: None,
         });
 
         let loop_stmt = ast::Stmt {
diff --git a/src/librustc_parse/lib.rs b/src/librustc_parse/lib.rs
index 0c817a71281..8ca3f6c5768 100644
--- a/src/librustc_parse/lib.rs
+++ b/src/librustc_parse/lib.rs
@@ -272,6 +272,12 @@ pub fn nt_to_tokenstream(nt: &Nonterminal, sess: &ParseSess, span: Span) -> Toke
             Some(tokenstream::TokenTree::token(token::Lifetime(ident.name), ident.span).into())
         }
         Nonterminal::NtTT(ref tt) => Some(tt.clone().into()),
+        Nonterminal::NtExpr(ref expr) => {
+            if expr.tokens.is_none() {
+                debug!("missing tokens for expr {:?}", expr);
+            }
+            prepend_attrs(sess, &expr.attrs, expr.tokens.as_ref(), span)
+        }
         _ => None,
     };
 
@@ -311,6 +317,8 @@ pub fn nt_to_tokenstream(nt: &Nonterminal, sess: &ParseSess, span: Span) -> Toke
             "cached tokens found, but they're not \"probably equal\", \
                 going with stringified version"
         );
+        info!("cached tokens: {:?}", tokens);
+        info!("reparsed tokens: {:?}", tokens_for_real);
     }
     tokens_for_real
 }
diff --git a/src/librustc_parse/parser/diagnostics.rs b/src/librustc_parse/parser/diagnostics.rs
index 93c7faf22a7..f05d0186138 100644
--- a/src/librustc_parse/parser/diagnostics.rs
+++ b/src/librustc_parse/parser/diagnostics.rs
@@ -95,6 +95,7 @@ impl RecoverQPath for Expr {
             kind: ExprKind::Path(qself, path),
             attrs: AttrVec::new(),
             id: ast::DUMMY_NODE_ID,
+            tokens: None,
         }
     }
 }
diff --git a/src/librustc_parse/parser/expr.rs b/src/librustc_parse/parser/expr.rs
index ca497a3b06f..e0c37284839 100644
--- a/src/librustc_parse/parser/expr.rs
+++ b/src/librustc_parse/parser/expr.rs
@@ -4,6 +4,7 @@ use super::{BlockMode, Parser, PathStyle, Restrictions, TokenType};
 use super::{SemiColonMode, SeqSep, TokenExpectType};
 use crate::maybe_recover_from_interpolated_ty_qpath;
 
+use log::debug;
 use rustc_ast::ast::{self, AttrStyle, AttrVec, CaptureBy, Field, Lit, UnOp, DUMMY_NODE_ID};
 use rustc_ast::ast::{AnonConst, BinOp, BinOpKind, FnDecl, FnRetTy, MacCall, Param, Ty, TyKind};
 use rustc_ast::ast::{Arm, Async, BlockCheckMode, Expr, ExprKind, Label, Movability, RangeLimits};
@@ -431,19 +432,23 @@ impl<'a> Parser<'a> {
     /// Parses a prefix-unary-operator expr.
     fn parse_prefix_expr(&mut self, attrs: Option<AttrVec>) -> PResult<'a, P<Expr>> {
         let attrs = self.parse_or_use_outer_attributes(attrs)?;
-        let lo = self.token.span;
-        // Note: when adding new unary operators, don't forget to adjust TokenKind::can_begin_expr()
-        let (hi, ex) = match self.token.uninterpolate().kind {
-            token::Not => self.parse_unary_expr(lo, UnOp::Not), // `!expr`
-            token::Tilde => self.recover_tilde_expr(lo),        // `~expr`
-            token::BinOp(token::Minus) => self.parse_unary_expr(lo, UnOp::Neg), // `-expr`
-            token::BinOp(token::Star) => self.parse_unary_expr(lo, UnOp::Deref), // `*expr`
-            token::BinOp(token::And) | token::AndAnd => self.parse_borrow_expr(lo),
-            token::Ident(..) if self.token.is_keyword(kw::Box) => self.parse_box_expr(lo),
-            token::Ident(..) if self.is_mistaken_not_ident_negation() => self.recover_not_expr(lo),
-            _ => return self.parse_dot_or_call_expr(Some(attrs)),
-        }?;
-        Ok(self.mk_expr(lo.to(hi), ex, attrs))
+        self.maybe_collect_tokens(!attrs.is_empty(), |this| {
+            let lo = this.token.span;
+            // Note: when adding new unary operators, don't forget to adjust TokenKind::can_begin_expr()
+            let (hi, ex) = match this.token.uninterpolate().kind {
+                token::Not => this.parse_unary_expr(lo, UnOp::Not), // `!expr`
+                token::Tilde => this.recover_tilde_expr(lo),        // `~expr`
+                token::BinOp(token::Minus) => this.parse_unary_expr(lo, UnOp::Neg), // `-expr`
+                token::BinOp(token::Star) => this.parse_unary_expr(lo, UnOp::Deref), // `*expr`
+                token::BinOp(token::And) | token::AndAnd => this.parse_borrow_expr(lo),
+                token::Ident(..) if this.token.is_keyword(kw::Box) => this.parse_box_expr(lo),
+                token::Ident(..) if this.is_mistaken_not_ident_negation() => {
+                    this.recover_not_expr(lo)
+                }
+                _ => return this.parse_dot_or_call_expr(Some(attrs)),
+            }?;
+            Ok(this.mk_expr(lo.to(hi), ex, attrs))
+        })
     }
 
     fn parse_prefix_expr_common(&mut self, lo: Span) -> PResult<'a, (Span, P<Expr>)> {
@@ -998,6 +1003,21 @@ impl<'a> Parser<'a> {
         }
     }
 
+    fn maybe_collect_tokens(
+        &mut self,
+        has_outer_attrs: bool,
+        f: impl FnOnce(&mut Self) -> PResult<'a, P<Expr>>,
+    ) -> PResult<'a, P<Expr>> {
+        if has_outer_attrs {
+            let (mut expr, tokens) = self.collect_tokens(f)?;
+            debug!("maybe_collect_tokens: Collected tokens for {:?} (tokens {:?}", expr, tokens);
+            expr.tokens = Some(tokens);
+            Ok(expr)
+        } else {
+            f(self)
+        }
+    }
+
     fn parse_lit_expr(&mut self, attrs: AttrVec) -> PResult<'a, P<Expr>> {
         let lo = self.token.span;
         match self.parse_opt_lit() {
@@ -2169,7 +2189,7 @@ impl<'a> Parser<'a> {
     }
 
     crate fn mk_expr(&self, span: Span, kind: ExprKind, attrs: AttrVec) -> P<Expr> {
-        P(Expr { kind, span, attrs, id: DUMMY_NODE_ID })
+        P(Expr { kind, span, attrs, id: DUMMY_NODE_ID, tokens: None })
     }
 
     pub(super) fn mk_expr_err(&self, span: Span) -> P<Expr> {
diff --git a/src/test/ui-fulldeps/pprust-expr-roundtrip.rs b/src/test/ui-fulldeps/pprust-expr-roundtrip.rs
index a3d31d25774..cef600bed5f 100644
--- a/src/test/ui-fulldeps/pprust-expr-roundtrip.rs
+++ b/src/test/ui-fulldeps/pprust-expr-roundtrip.rs
@@ -56,6 +56,7 @@ fn expr(kind: ExprKind) -> P<Expr> {
         kind,
         span: DUMMY_SP,
         attrs: ThinVec::new(),
+        tokens: None
     })
 }
 
@@ -200,6 +201,7 @@ impl MutVisitor for AddParens {
                 kind: ExprKind::Paren(e),
                 span: DUMMY_SP,
                 attrs: ThinVec::new(),
+                tokens: None
             })
         });
     }
diff --git a/src/test/ui/proc-macro/keep-expr-tokens.rs b/src/test/ui/proc-macro/keep-expr-tokens.rs
new file mode 100644
index 00000000000..888785363cf
--- /dev/null
+++ b/src/test/ui/proc-macro/keep-expr-tokens.rs
@@ -0,0 +1,15 @@
+// aux-build:test-macros.rs
+
+#![feature(stmt_expr_attributes)]
+#![feature(proc_macro_hygiene)]
+
+extern crate test_macros;
+
+use test_macros::recollect_attr;
+
+fn main() {
+    #[test_macros::recollect_attr]
+    for item in missing_fn() {} //~ ERROR cannot find
+
+    (#[recollect_attr] #[recollect_attr] ((#[recollect_attr] bad))); //~ ERROR cannot
+}
diff --git a/src/test/ui/proc-macro/keep-expr-tokens.stderr b/src/test/ui/proc-macro/keep-expr-tokens.stderr
new file mode 100644
index 00000000000..2be8c0184da
--- /dev/null
+++ b/src/test/ui/proc-macro/keep-expr-tokens.stderr
@@ -0,0 +1,15 @@
+error[E0425]: cannot find function `missing_fn` in this scope
+  --> $DIR/keep-expr-tokens.rs:12:17
+   |
+LL |     for item in missing_fn() {}
+   |                 ^^^^^^^^^^ not found in this scope
+
+error[E0425]: cannot find value `bad` in this scope
+  --> $DIR/keep-expr-tokens.rs:14:62
+   |
+LL |     (#[recollect_attr] #[recollect_attr] ((#[recollect_attr] bad)));
+   |                                                              ^^^ not found in this scope
+
+error: aborting due to 2 previous errors
+
+For more information about this error, try `rustc --explain E0425`.