Eagerly check for ambiguity in macro parsing

compiler/rustc_expand/src/mbe/diagnostics.rs

@@ -2,6 +2,7 @@ use std::borrow::Cow;
 
 use rustc_ast::token::{self, Token};
 use rustc_ast::tokenstream::TokenStream;
+use rustc_data_structures::fx::FxHashSet;
 use rustc_errors::{Applicability, Diag, DiagCtxtHandle, DiagMessage, pluralize};
 use rustc_hir::attrs::diagnostic::{CustomDiagnostic, Directive, FormatArgs};
 use rustc_macros::Subdiagnostic;
@@ -152,7 +153,13 @@ struct CollectTrackerAndEmitter<'dcx, 'matcher> {
     /// The matcher currently being parsed.
     //
     // FIXME: Factor out a per-arm `Tracker` so that the `Option` is unnecessary.
-    current: Option<WhichMatcher>,
+    current: Option<(WhichMatcher, &'matcher [MatcherLoc])>,
+
+    /// Matches of [`MatcherLoc`]s that successfully consumed input from the parser.
+    ///
+    /// This accumulates all calls to [`Tracker::matched_one()`]. It is used to identify all
+    /// competing matches for ambiguity errors.
+    matches: FxHashSet<SuccessfulMatch>,
 
     remaining_matcher: Option<&'matcher MatcherLoc>,
     /// Which arm's failure should we report? (the one furthest along)
@@ -161,6 +168,17 @@ struct CollectTrackerAndEmitter<'dcx, 'matcher> {
     result: Option<(Span, ErrorGuaranteed)>,
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+struct SuccessfulMatch {
+    /// The position in the parser.
+    ///
+    /// As per [`Parser::approx_token_stream_pos()`].
+    input_pos: u32,
+
+    /// The index of the [`MatcherLoc`].
+    loc_index: u32,
+}
+
 struct BestFailure {
     token: Token,
 
@@ -183,12 +201,12 @@ impl BestFailure {
 }
 
 impl<'dcx, 'matcher> Tracker<'matcher> for CollectTrackerAndEmitter<'dcx, 'matcher> {
-    fn prepare(&mut self, which_matcher: WhichMatcher) {
+    fn prepare(&mut self, which_matcher: WhichMatcher, matcher: &'matcher [MatcherLoc]) {
         if self.current.is_some() {
             bug!("`Self::after_arm()` was not called to clean up context");
         }
 
-        self.current = Some(which_matcher);
+        self.current = Some((which_matcher, matcher));
     }
 
     fn before_match_loc(&mut self, parser: &TtParser, matcher: &'matcher MatcherLoc) {
@@ -199,6 +217,13 @@ impl<'dcx, 'matcher> Tracker<'matcher> for CollectTrackerAndEmitter<'dcx, 'match
         }
     }
 
+    fn matched_one(&mut self, parser: &Parser<'_>, loc_index: usize) {
+        let input_pos = parser.approx_token_stream_pos();
+        let loc_index: u32 = loc_index.try_into().unwrap();
+        let m = SuccessfulMatch { input_pos, loc_index };
+        self.matches.insert(m);
+    }
+
     fn after_arm(&mut self, result: &NamedParseResult) {
         match *result {
             Success(_) => {
@@ -223,10 +248,11 @@ impl<'dcx, 'matcher> Tracker<'matcher> for CollectTrackerAndEmitter<'dcx, 'match
         }
 
         self.current = None;
+        self.matches.clear();
     }
 
     fn failure(&mut self, parser: &Parser<'_>) {
-        let Some(which_matcher) = self.current else {
+        let Some((which_matcher, _)) = self.current else {
             bug!("`Self::prepare()` was not called to initialize context");
         };
 
@@ -262,12 +288,28 @@ impl<'dcx, 'matcher> Tracker<'matcher> for CollectTrackerAndEmitter<'dcx, 'match
         }
     }
 
-    fn ambiguity(
-        &mut self,
-        parser: &Parser<'_>,
-        bb_locs: impl IntoIterator<Item = &'matcher MatcherLoc>,
-        next_locs: impl IntoIterator<Item = &'matcher MatcherLoc>,
-    ) {
+    fn ambiguity(&mut self, parser: &Parser<'_>) {
+        let Some((_, matcher)) = self.current else {
+            bug!("`Self::prepare()` was not called to initialize context");
+        };
+
+        #[expect(
+            rustc::potential_query_instability,
+            reason = "sorting the results deterministically afterwards"
+        )]
+        let (mut bb_locs, mut next_locs) = self
+            .matches
+            .iter()
+            .filter(|m| m.input_pos == parser.approx_token_stream_pos())
+            .partition::<Vec<&SuccessfulMatch>, _>(|m| {
+                let loc = &matcher[m.loc_index as usize];
+                matches!(loc, MatcherLoc::MetaVarDecl { .. })
+            });
+
+        // Use a reasonable and deterministic ordering for data in the error message.
+        bb_locs.sort_unstable_by_key(|m| m.loc_index);
+        next_locs.sort_unstable_by_key(|m| m.loc_index);
+
         let span = parser.token.span.substitute_dummy(self.root_span);
 
         if parser.token == token::Eof {
@@ -279,19 +321,18 @@ impl<'dcx, 'matcher> Tracker<'matcher> for CollectTrackerAndEmitter<'dcx, 'match
 
         let nts = bb_locs
             .into_iter()
-            .map(|loc| match loc {
-                MatcherLoc::MetaVarDecl { bind, kind, .. } => {
-                    format!("{kind} ('{bind}')")
-                }
-                _ => unreachable!(),
+            .map(|m| {
+                let loc = &matcher[m.loc_index as usize];
+                let MatcherLoc::MetaVarDecl { bind, kind, .. } = loc else { unreachable!() };
+                format!("{kind} ('{bind}')")
             })
             .collect::<Vec<String>>()
             .join(" or ");
 
         let msg = format!(
             "local ambiguity when calling macro `{}`: multiple parsing options: {}",
             self.macro_name,
-            match next_locs.into_iter().count() {
+            match next_locs.len() {
                 0 => format!("built-in NTs {nts}."),
                 n => format!("built-in NTs {nts} or {n} other option{s}.", s = pluralize!(n)),
             }
@@ -316,6 +357,7 @@ impl<'dcx> CollectTrackerAndEmitter<'dcx, '_> {
             macro_name,
             dcx,
             current: None,
+            matches: FxHashSet::default(),
             remaining_matcher: None,
             best_failure: None,
             root_span,