rustfmt: Check that a token can begin a nonterminal kind before parsing it as a macro arg

src/tools/rustfmt/src/parse/macros/mod.rs

@@ -1,7 +1,7 @@
-use rustc_ast::token::{Delimiter, TokenKind};
+use rustc_ast::token::{Delimiter, NonterminalKind, TokenKind};
 use rustc_ast::tokenstream::TokenStream;
 use rustc_ast::{ast, ptr};
-use rustc_parse::parser::{ForceCollect, Parser};
+use rustc_parse::parser::{ForceCollect, Parser, Recovery};
 use rustc_parse::{stream_to_parser, MACRO_ARGUMENTS};
 use rustc_session::parse::ParseSess;
 use rustc_span::symbol::{self, kw};
@@ -15,7 +15,7 @@ pub(crate) mod cfg_if;
 pub(crate) mod lazy_static;
 
 fn build_stream_parser<'a>(sess: &'a ParseSess, tokens: TokenStream) -> Parser<'a> {
-    stream_to_parser(sess, tokens, MACRO_ARGUMENTS)
+    stream_to_parser(sess, tokens, MACRO_ARGUMENTS).recovery(Recovery::Forbidden)
 }
 
 fn build_parser<'a>(context: &RewriteContext<'a>, tokens: TokenStream) -> Parser<'a> {
@@ -24,45 +24,51 @@ fn build_parser<'a>(context: &RewriteContext<'a>, tokens: TokenStream) -> Parser
 
 fn parse_macro_arg<'a, 'b: 'a>(parser: &'a mut Parser<'b>) -> Option<MacroArg> {
     macro_rules! parse_macro_arg {
-        ($macro_arg:ident, $parser:expr, $f:expr) => {
+        ($macro_arg:ident, $nt_kind:expr, $try_parse:expr, $then:expr) => {
             let mut cloned_parser = (*parser).clone();
-            match $parser(&mut cloned_parser) {
-                Ok(x) => {
-                    if parser.sess.dcx.has_errors().is_some() {
+            if Parser::nonterminal_may_begin_with($nt_kind, &cloned_parser.token) {
+                match $try_parse(&mut cloned_parser) {
+                    Ok(x) => {
+                        if parser.sess.dcx.has_errors().is_some() {
+                            parser.sess.dcx.reset_err_count();
+                        } else {
+                            // Parsing succeeded.
+                            *parser = cloned_parser;
+                            return Some(MacroArg::$macro_arg($then(x)?));
+                        }
+                    }
+                    Err(e) => {
+                        e.cancel();
                         parser.sess.dcx.reset_err_count();
-                    } else {
-                        // Parsing succeeded.
-                        *parser = cloned_parser;
-                        return Some(MacroArg::$macro_arg($f(x)?));
                     }
                 }
-                Err(e) => {
-                    e.cancel();
-                    parser.sess.dcx.reset_err_count();
-                }
             }
         };
     }
 
     parse_macro_arg!(
         Expr,
-        |parser: &mut rustc_parse::parser::Parser<'b>| parser.parse_expr(),
+        NonterminalKind::Expr,
+        |parser: &mut Parser<'b>| parser.parse_expr(),
         |x: ptr::P<ast::Expr>| Some(x)
     );
     parse_macro_arg!(
         Ty,
-        |parser: &mut rustc_parse::parser::Parser<'b>| parser.parse_ty(),
+        NonterminalKind::Ty,
+        |parser: &mut Parser<'b>| parser.parse_ty(),
         |x: ptr::P<ast::Ty>| Some(x)
     );
     parse_macro_arg!(
         Pat,
-        |parser: &mut rustc_parse::parser::Parser<'b>| parser.parse_pat_no_top_alt(None, None),
+        NonterminalKind::PatParam { inferred: false },
+        |parser: &mut Parser<'b>| parser.parse_pat_no_top_alt(None, None),
         |x: ptr::P<ast::Pat>| Some(x)
     );
     // `parse_item` returns `Option<ptr::P<ast::Item>>`.
     parse_macro_arg!(
         Item,
-        |parser: &mut rustc_parse::parser::Parser<'b>| parser.parse_item(ForceCollect::No),
+        NonterminalKind::Item,
+        |parser: &mut Parser<'b>| parser.parse_item(ForceCollect::No),
         |x: Option<ptr::P<ast::Item>>| x
     );
 

src/tools/rustfmt/tests/source/macros/rewrite-const-item.rs

@@ -0,0 +1 @@
+m!(const N: usize = 0;);

src/tools/rustfmt/tests/target/macros/rewrite-const-item.rs

@@ -0,0 +1,3 @@
+m!(
+    const N: usize = 0;
+);