parser: Ensure that all nonterminals have tokens after parsing

compiler/rustc_ast/src/ast_like.rs

@@ -82,7 +82,8 @@ impl AstLike for crate::token::Nonterminal {
             Nonterminal::NtMeta(attr_item) => attr_item.tokens_mut(),
             Nonterminal::NtPath(path) => path.tokens_mut(),
             Nonterminal::NtVis(vis) => vis.tokens_mut(),
-            _ => panic!("Called tokens_mut on {:?}", self),
+            Nonterminal::NtBlock(block) => block.tokens_mut(),
+            Nonterminal::NtIdent(..) | Nonterminal::NtLifetime(..) | Nonterminal::NtTT(..) => None,
         }
     }
 }

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();

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

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,

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

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,21 +89,22 @@ 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());
             return Ok(None);
         }))
     }
 
-    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) {

src/test/ui/proc-macro/expr-stmt-nonterminal-tokens.rs

@@ -0,0 +1,37 @@
+// check-pass
+// aux-build:test-macros.rs
+
+#![feature(decl_macro)]
+#![feature(stmt_expr_attributes)]
+
+#![no_std] // Don't load unnecessary hygiene information from std
+extern crate std;
+
+#[macro_use]
+extern crate test_macros;
+
+macro mac {
+    (expr $expr:expr) => {
+        #[derive(Print)]
+        enum E {
+            V = { let _ = $expr; 0 },
+        }
+    },
+    (stmt $stmt:stmt) => {
+        #[derive(Print)]
+        enum E {
+            V = { let _ = { $stmt }; 0 },
+        }
+    },
+}
+
+const PATH: u8 = 2;
+
+fn main() {
+    mac!(expr #[allow(warnings)] 0);
+    mac!(stmt 0);
+    mac!(stmt {});
+    mac!(stmt PATH);
+    mac!(stmt 0 + 1);
+    mac!(stmt PATH + 1);
+}