parser: Ensure that all nonterminals have tokens after parsing

5 files changed, 33 insertions, 30 deletions

diff --git a/compiler/rustc_parse/src/parser/attr_wrapper.rs b/compiler/rustc_parse/src/parser/attr_wrapper.rs
index 35759a396e8..e1d0b84f419 100644
--- a/compiler/rustc_parse/src/parser/attr_wrapper.rs
+++ b/compiler/rustc_parse/src/parser/attr_wrapper.rs
@@ -342,16 +342,10 @@ impl<'a> Parser<'a> {
 
         // If we support tokens at all
         if let Some(target_tokens) = ret.tokens_mut() {
-            if let Some(target_tokens) = target_tokens {
-                assert!(
-                    !self.capture_cfg,
-                    "Encountered existing tokens with capture_cfg set: {:?}",
-                    target_tokens
-                );
-            } else {
+            if target_tokens.is_none() {
                 // Store se our newly captured tokens into the AST node
                 *target_tokens = Some(tokens.clone());
-            };
+            }
         }
 
         let final_attrs = ret.attrs();
diff --git a/compiler/rustc_parse/src/parser/expr.rs b/compiler/rustc_parse/src/parser/expr.rs
index 56c97b59476..a764cf6bdb0 100644
--- a/compiler/rustc_parse/src/parser/expr.rs
+++ b/compiler/rustc_parse/src/parser/expr.rs
@@ -94,17 +94,7 @@ impl<'a> Parser<'a> {
 
     /// Parses an expression, forcing tokens to be collected
     pub fn parse_expr_force_collect(&mut self) -> PResult<'a, P<Expr>> {
-        // If we have outer attributes, then the call to `collect_tokens_trailing_token`
-        // will be made for us.
-        if matches!(self.token.kind, TokenKind::Pound | TokenKind::DocComment(..)) {
-            self.parse_expr()
-        } else {
-            // If we don't have outer attributes, then we need to ensure
-            // that collection happens by using `collect_tokens_no_attrs`.
-            // Expression don't support custom inner attributes, so `parse_expr`
-            // will never try to collect tokens if we don't have outer attributes.
-            self.collect_tokens_no_attrs(|this| this.parse_expr())
-        }
+        self.collect_tokens_no_attrs(|this| this.parse_expr())
     }
 
     pub fn parse_anon_const_expr(&mut self) -> PResult<'a, AnonConst> {
diff --git a/compiler/rustc_parse/src/parser/mod.rs b/compiler/rustc_parse/src/parser/mod.rs
index 4c2bc6ebf31..cd9f84db5e5 100644
--- a/compiler/rustc_parse/src/parser/mod.rs
+++ b/compiler/rustc_parse/src/parser/mod.rs
@@ -63,6 +63,7 @@ enum BlockMode {
 
 /// Whether or not we should force collection of tokens for an AST node,
 /// regardless of whether or not it has attributes
+#[derive(Clone, Copy, PartialEq)]
 pub enum ForceCollect {
     Yes,
     No,
diff --git a/compiler/rustc_parse/src/parser/nonterminal.rs b/compiler/rustc_parse/src/parser/nonterminal.rs
index 0c43e304f1e..30a6b61407f 100644
--- a/compiler/rustc_parse/src/parser/nonterminal.rs
+++ b/compiler/rustc_parse/src/parser/nonterminal.rs
@@ -1,5 +1,6 @@
 use rustc_ast::ptr::P;
 use rustc_ast::token::{self, Nonterminal, NonterminalKind, Token};
+use rustc_ast::AstLike;
 use rustc_ast_pretty::pprust;
 use rustc_errors::PResult;
 use rustc_span::symbol::{kw, Ident};
@@ -102,7 +103,7 @@ impl<'a> Parser<'a> {
         // which requires having captured tokens available. Since we cannot determine
         // 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 {
+        let mut nt = match kind {
             NonterminalKind::Item => match self.parse_item(ForceCollect::Yes)? {
                 Some(item) => token::NtItem(item),
                 None => {
@@ -169,7 +170,19 @@ impl<'a> Parser<'a> {
                     return Err(self.struct_span_err(self.token.span, msg));
                 }
             }
-        })
+        };
+
+        // If tokens are supported at all, they should be collected.
+        if matches!(nt.tokens_mut(), Some(None)) {
+            panic!(
+                "Missing tokens for nt {:?} at {:?}: {:?}",
+                nt,
+                nt.span(),
+                pprust::nonterminal_to_string(&nt)
+            );
+        }
+
+        Ok(nt)
     }
 }
 
diff --git a/compiler/rustc_parse/src/parser/stmt.rs b/compiler/rustc_parse/src/parser/stmt.rs
index b40eed8c5d1..4f0dcfeb5da 100644
--- a/compiler/rustc_parse/src/parser/stmt.rs
+++ b/compiler/rustc_parse/src/parser/stmt.rs
@@ -73,7 +73,11 @@ impl<'a> Parser<'a> {
             // or `auto trait` items. We aim to parse an arbitrary path `a::b` but not something
             // that starts like a path (1 token), but it fact not a path.
             // Also, we avoid stealing syntax from `parse_item_`.
-            self.parse_stmt_path_start(lo, attrs, force_collect)?
+            if force_collect == ForceCollect::Yes {
+                self.collect_tokens_no_attrs(|this| this.parse_stmt_path_start(lo, attrs))
+            } else {
+                self.parse_stmt_path_start(lo, attrs)
+            }?
         } else if let Some(item) =
             self.parse_item_common(attrs.clone(), false, true, |_| true, force_collect)?
         {
@@ -85,7 +89,13 @@ impl<'a> Parser<'a> {
             self.mk_stmt(lo, StmtKind::Empty)
         } else if self.token != token::CloseDelim(token::Brace) {
             // Remainder are line-expr stmts.
-            let e = self.parse_expr_res(Restrictions::STMT_EXPR, Some(attrs))?;
+            let e = if force_collect == ForceCollect::Yes {
+                self.collect_tokens_no_attrs(|this| {
+                    this.parse_expr_res(Restrictions::STMT_EXPR, Some(attrs))
+                })
+            } else {
+                self.parse_expr_res(Restrictions::STMT_EXPR, Some(attrs))
+            }?;
             self.mk_stmt(lo.to(e.span), StmtKind::Expr(e))
         } else {
             self.error_outer_attrs(&attrs.take_for_recovery());
@@ -93,13 +103,8 @@ impl<'a> Parser<'a> {
         }))
     }
 
-    fn parse_stmt_path_start(
-        &mut self,
-        lo: Span,
-        attrs: AttrWrapper,
-        force_collect: ForceCollect,
-    ) -> PResult<'a, Stmt> {
-        let stmt = self.collect_tokens_trailing_token(attrs, force_collect, |this, attrs| {
+    fn parse_stmt_path_start(&mut self, lo: Span, attrs: AttrWrapper) -> PResult<'a, Stmt> {
+        let stmt = self.collect_tokens_trailing_token(attrs, ForceCollect::No, |this, attrs| {
             let path = this.parse_path(PathStyle::Expr)?;
 
             if this.eat(&token::Not) {