Force token collection to run when parsing nonterminals

Fixes #81007

Previously, we would fail to collect tokens in the proper place when
only builtin attributes were present. As a result, we would end up with
attribute tokens in the collected `TokenStream`, leading to duplication
when we attempted to prepend the attributes from the AST node.

We now explicitly track when token collection must be performed due to
nomterminal parsing.

5 files changed, 89 insertions, 63 deletions

diff --git a/compiler/rustc_parse/src/parser/expr.rs b/compiler/rustc_parse/src/parser/expr.rs
index 6db415ead41..47869f775fe 100644
--- a/compiler/rustc_parse/src/parser/expr.rs
+++ b/compiler/rustc_parse/src/parser/expr.rs
@@ -472,7 +472,11 @@ 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 needs_tokens = super::attr::maybe_needs_tokens(&attrs);
+        // FIXME: Use super::attr::maybe_needs_tokens(&attrs) once we come up
+        // with a good way of passing `force_tokens` through from `parse_nonterminal`.
+        // Checking !attrs.is_empty() is correct, but will cause us to unnecessarily
+        // capture tokens in some circumstances.
+        let needs_tokens = !attrs.is_empty();
         let do_parse = |this: &mut Parser<'a>| {
             let lo = this.token.span;
             // Note: when adding new unary operators, don't forget to adjust TokenKind::can_begin_expr()
diff --git a/compiler/rustc_parse/src/parser/item.rs b/compiler/rustc_parse/src/parser/item.rs
index 810ae61307c..28067f0216c 100644
--- a/compiler/rustc_parse/src/parser/item.rs
+++ b/compiler/rustc_parse/src/parser/item.rs
@@ -1,8 +1,8 @@
 use super::diagnostics::{dummy_arg, ConsumeClosingDelim, Error};
 use super::ty::{AllowPlus, RecoverQPath, RecoverReturnSign};
-use super::{FollowedByType, Parser, PathStyle};
+use super::{FollowedByType, ForceCollect, Parser, PathStyle};
 
-use crate::maybe_whole;
+use crate::{maybe_collect_tokens, maybe_whole};
 
 use rustc_ast::ptr::P;
 use rustc_ast::token::{self, TokenKind};
@@ -69,7 +69,7 @@ impl<'a> Parser<'a> {
         unsafety: Unsafe,
     ) -> PResult<'a, Mod> {
         let mut items = vec![];
-        while let Some(item) = self.parse_item()? {
+        while let Some(item) = self.parse_item(ForceCollect::No)? {
             items.push(item);
             self.maybe_consume_incorrect_semicolon(&items);
         }
@@ -93,13 +93,17 @@ impl<'a> Parser<'a> {
 pub(super) type ItemInfo = (Ident, ItemKind);
 
 impl<'a> Parser<'a> {
-    pub fn parse_item(&mut self) -> PResult<'a, Option<P<Item>>> {
-        self.parse_item_(|_| true).map(|i| i.map(P))
+    pub fn parse_item(&mut self, force_collect: ForceCollect) -> PResult<'a, Option<P<Item>>> {
+        self.parse_item_(|_| true, force_collect).map(|i| i.map(P))
     }
 
-    fn parse_item_(&mut self, req_name: ReqName) -> PResult<'a, Option<Item>> {
+    fn parse_item_(
+        &mut self,
+        req_name: ReqName,
+        force_collect: ForceCollect,
+    ) -> PResult<'a, Option<Item>> {
         let attrs = self.parse_outer_attributes()?;
-        self.parse_item_common(attrs, true, false, req_name)
+        self.parse_item_common(attrs, true, false, req_name, force_collect)
     }
 
     pub(super) fn parse_item_common(
@@ -108,6 +112,7 @@ impl<'a> Parser<'a> {
         mac_allowed: bool,
         attrs_allowed: bool,
         req_name: ReqName,
+        force_collect: ForceCollect,
     ) -> PResult<'a, Option<Item>> {
         maybe_whole!(self, NtItem, |item| {
             let mut item = item;
@@ -116,16 +121,12 @@ impl<'a> Parser<'a> {
             Some(item.into_inner())
         });
 
-        let needs_tokens = super::attr::maybe_needs_tokens(&attrs);
-
         let mut unclosed_delims = vec![];
-        let parse_item = |this: &mut Self| {
+        let item = maybe_collect_tokens!(self, force_collect, &attrs, |this: &mut Self| {
             let item = this.parse_item_common_(attrs, mac_allowed, attrs_allowed, req_name);
             unclosed_delims.append(&mut this.unclosed_delims);
             item
-        };
-
-        let item = if needs_tokens { self.collect_tokens(parse_item) } else { parse_item(self) }?;
+        })?;
 
         self.unclosed_delims.append(&mut unclosed_delims);
         Ok(item)
@@ -731,20 +732,22 @@ impl<'a> Parser<'a> {
 
     /// Parses associated items.
     fn parse_assoc_item(&mut self, req_name: ReqName) -> PResult<'a, Option<Option<P<AssocItem>>>> {
-        Ok(self.parse_item_(req_name)?.map(|Item { attrs, id, span, vis, ident, kind, tokens }| {
-            let kind = match AssocItemKind::try_from(kind) {
-                Ok(kind) => kind,
-                Err(kind) => match kind {
-                    ItemKind::Static(a, _, b) => {
-                        self.struct_span_err(span, "associated `static` items are not allowed")
-                            .emit();
-                        AssocItemKind::Const(Defaultness::Final, a, b)
-                    }
-                    _ => return self.error_bad_item_kind(span, &kind, "`trait`s or `impl`s"),
-                },
-            };
-            Some(P(Item { attrs, id, span, vis, ident, kind, tokens }))
-        }))
+        Ok(self.parse_item_(req_name, ForceCollect::No)?.map(
+            |Item { attrs, id, span, vis, ident, kind, tokens }| {
+                let kind = match AssocItemKind::try_from(kind) {
+                    Ok(kind) => kind,
+                    Err(kind) => match kind {
+                        ItemKind::Static(a, _, b) => {
+                            self.struct_span_err(span, "associated `static` items are not allowed")
+                                .emit();
+                            AssocItemKind::Const(Defaultness::Final, a, b)
+                        }
+                        _ => return self.error_bad_item_kind(span, &kind, "`trait`s or `impl`s"),
+                    },
+                };
+                Some(P(Item { attrs, id, span, vis, ident, kind, tokens }))
+            },
+        ))
     }
 
     /// Parses a `type` alias with the following grammar:
@@ -921,19 +924,21 @@ impl<'a> Parser<'a> {
 
     /// Parses a foreign item (one in an `extern { ... }` block).
     pub fn parse_foreign_item(&mut self) -> PResult<'a, Option<Option<P<ForeignItem>>>> {
-        Ok(self.parse_item_(|_| true)?.map(|Item { attrs, id, span, vis, ident, kind, tokens }| {
-            let kind = match ForeignItemKind::try_from(kind) {
-                Ok(kind) => kind,
-                Err(kind) => match kind {
-                    ItemKind::Const(_, a, b) => {
-                        self.error_on_foreign_const(span, ident);
-                        ForeignItemKind::Static(a, Mutability::Not, b)
-                    }
-                    _ => return self.error_bad_item_kind(span, &kind, "`extern` blocks"),
-                },
-            };
-            Some(P(Item { attrs, id, span, vis, ident, kind, tokens }))
-        }))
+        Ok(self.parse_item_(|_| true, ForceCollect::No)?.map(
+            |Item { attrs, id, span, vis, ident, kind, tokens }| {
+                let kind = match ForeignItemKind::try_from(kind) {
+                    Ok(kind) => kind,
+                    Err(kind) => match kind {
+                        ItemKind::Const(_, a, b) => {
+                            self.error_on_foreign_const(span, ident);
+                            ForeignItemKind::Static(a, Mutability::Not, b)
+                        }
+                        _ => return self.error_bad_item_kind(span, &kind, "`extern` blocks"),
+                    },
+                };
+                Some(P(Item { attrs, id, span, vis, ident, kind, tokens }))
+            },
+        ))
     }
 
     fn error_bad_item_kind<T>(&self, span: Span, kind: &ItemKind, ctx: &str) -> Option<T> {
@@ -1515,7 +1520,7 @@ impl<'a> Parser<'a> {
         {
             let kw_token = self.token.clone();
             let kw_str = pprust::token_to_string(&kw_token);
-            let item = self.parse_item()?;
+            let item = self.parse_item(ForceCollect::No)?;
 
             self.struct_span_err(
                 kw_token.span,
diff --git a/compiler/rustc_parse/src/parser/mod.rs b/compiler/rustc_parse/src/parser/mod.rs
index 5d7ea5b8d57..c85b7a00732 100644
--- a/compiler/rustc_parse/src/parser/mod.rs
+++ b/compiler/rustc_parse/src/parser/mod.rs
@@ -54,6 +54,13 @@ enum BlockMode {
     Ignore,
 }
 
+/// Whether or not we should force collection of tokens for an AST node,
+/// regardless of whether or not it has attributes
+pub enum ForceCollect {
+    Yes,
+    No,
+}
+
 /// Like `maybe_whole_expr`, but for things other than expressions.
 #[macro_export]
 macro_rules! maybe_whole {
@@ -1413,3 +1420,16 @@ fn make_token_stream(
     assert!(stack.is_empty(), "Stack should be empty: final_buf={:?} stack={:?}", final_buf, stack);
     TokenStream::new(final_buf.inner)
 }
+
+#[macro_export]
+macro_rules! maybe_collect_tokens {
+    ($self:ident, $force_collect:expr, $attrs:expr, $f:expr) => {
+        if matches!($force_collect, ForceCollect::Yes)
+            || $crate::parser::attr::maybe_needs_tokens($attrs)
+        {
+            $self.collect_tokens($f)
+        } else {
+            $f($self)
+        }
+    };
+}
diff --git a/compiler/rustc_parse/src/parser/nonterminal.rs b/compiler/rustc_parse/src/parser/nonterminal.rs
index 97d0c0d8745..012b76d3d18 100644
--- a/compiler/rustc_parse/src/parser/nonterminal.rs
+++ b/compiler/rustc_parse/src/parser/nonterminal.rs
@@ -5,7 +5,7 @@ use rustc_errors::PResult;
 use rustc_span::symbol::{kw, Ident};
 
 use crate::parser::pat::{GateOr, RecoverComma};
-use crate::parser::{FollowedByType, Parser, PathStyle};
+use crate::parser::{FollowedByType, ForceCollect, Parser, PathStyle};
 
 impl<'a> Parser<'a> {
     /// Checks whether a non-terminal may begin with a particular token.
@@ -98,7 +98,7 @@ impl<'a> Parser<'a> {
         // in advance whether or not a proc-macro will be (transitively) invoked,
         // we always capture tokens for any `Nonterminal` which needs them.
         Ok(match kind {
-            NonterminalKind::Item => match self.collect_tokens(|this| this.parse_item())? {
+            NonterminalKind::Item => match self.parse_item(ForceCollect::Yes)? {
                 Some(item) => token::NtItem(item),
                 None => {
                     return Err(self.struct_span_err(self.token.span, "expected an item keyword"));
@@ -107,7 +107,7 @@ impl<'a> Parser<'a> {
             NonterminalKind::Block => {
                 token::NtBlock(self.collect_tokens(|this| this.parse_block())?)
             }
-            NonterminalKind::Stmt => match self.collect_tokens(|this| this.parse_stmt())? {
+            NonterminalKind::Stmt => match self.parse_stmt(ForceCollect::Yes)? {
                 Some(s) => token::NtStmt(s),
                 None => {
                     return Err(self.struct_span_err(self.token.span, "expected a statement"));
diff --git a/compiler/rustc_parse/src/parser/stmt.rs b/compiler/rustc_parse/src/parser/stmt.rs
index 641b29227db..da60ba8472b 100644
--- a/compiler/rustc_parse/src/parser/stmt.rs
+++ b/compiler/rustc_parse/src/parser/stmt.rs
@@ -3,8 +3,8 @@ use super::diagnostics::{AttemptLocalParseRecovery, Error};
 use super::expr::LhsExpr;
 use super::pat::{GateOr, RecoverComma};
 use super::path::PathStyle;
-use super::{BlockMode, Parser, Restrictions, SemiColonMode};
-use crate::maybe_whole;
+use super::{BlockMode, ForceCollect, Parser, Restrictions, SemiColonMode};
+use crate::{maybe_collect_tokens, maybe_whole};
 
 use rustc_ast as ast;
 use rustc_ast::attr::HasAttrs;
@@ -24,17 +24,21 @@ impl<'a> Parser<'a> {
     /// Parses a statement. This stops just before trailing semicolons on everything but items.
     /// e.g., a `StmtKind::Semi` parses to a `StmtKind::Expr`, leaving the trailing `;` unconsumed.
     // Public for rustfmt usage.
-    pub fn parse_stmt(&mut self) -> PResult<'a, Option<Stmt>> {
-        Ok(self.parse_stmt_without_recovery().unwrap_or_else(|mut e| {
+    pub fn parse_stmt(&mut self, force_collect: ForceCollect) -> PResult<'a, Option<Stmt>> {
+        Ok(self.parse_stmt_without_recovery(force_collect).unwrap_or_else(|mut e| {
             e.emit();
             self.recover_stmt_(SemiColonMode::Break, BlockMode::Ignore);
             None
         }))
     }
 
-    fn parse_stmt_without_recovery(&mut self) -> PResult<'a, Option<Stmt>> {
+    /// If `force_capture` is true, forces collection of tokens regardless of whether
+    /// or not we have attributes
+    fn parse_stmt_without_recovery(
+        &mut self,
+        force_collect: ForceCollect,
+    ) -> PResult<'a, Option<Stmt>> {
         let mut attrs = self.parse_outer_attributes()?;
-        let has_attrs = !attrs.is_empty();
         let lo = self.token.span;
 
         maybe_whole!(self, NtStmt, |stmt| {
@@ -46,7 +50,7 @@ impl<'a> Parser<'a> {
             Some(stmt)
         });
 
-        let parse_stmt_inner = |this: &mut Self| {
+        maybe_collect_tokens!(self, force_collect, &attrs, |this: &mut Self| {
             let stmt = if this.eat_keyword(kw::Let) {
                 this.parse_local_mk(lo, attrs.into())?
             } else if this.is_kw_followed_by_ident(kw::Mut) {
@@ -69,7 +73,7 @@ impl<'a> Parser<'a> {
                 // Also, we avoid stealing syntax from `parse_item_`.
                 this.parse_stmt_path_start(lo, attrs)?
             } else if let Some(item) =
-                this.parse_item_common(attrs.clone(), false, true, |_| true)?
+                this.parse_item_common(attrs.clone(), false, true, |_| true, force_collect)?
             {
                 // FIXME: Bad copy of attrs
                 this.mk_stmt(lo.to(item.span), StmtKind::Item(P(item)))
@@ -86,14 +90,7 @@ impl<'a> Parser<'a> {
                 return Ok(None);
             };
             Ok(Some(stmt))
-        };
-
-        let stmt = if has_attrs {
-            self.collect_tokens(parse_stmt_inner)?
-        } else {
-            parse_stmt_inner(self)?
-        };
-        Ok(stmt)
+        })
     }
 
     fn parse_stmt_path_start(&mut self, lo: Span, attrs: Vec<Attribute>) -> PResult<'a, Stmt> {
@@ -292,7 +289,7 @@ impl<'a> Parser<'a> {
         //      bar;
         //
         // which is valid in other languages, but not Rust.
-        match self.parse_stmt_without_recovery() {
+        match self.parse_stmt_without_recovery(ForceCollect::No) {
             // If the next token is an open brace (e.g., `if a b {`), the place-
             // inside-a-block suggestion would be more likely wrong than right.
             Ok(Some(_))
@@ -395,7 +392,7 @@ impl<'a> Parser<'a> {
         // Skip looking for a trailing semicolon when we have an interpolated statement.
         maybe_whole!(self, NtStmt, |x| Some(x));
 
-        let mut stmt = match self.parse_stmt_without_recovery()? {
+        let mut stmt = match self.parse_stmt_without_recovery(ForceCollect::No)? {
             Some(stmt) => stmt,
             None => return Ok(None),
         };