Use depth-first traversal for macro parsing

compiler/rustc_expand/src/mbe/diagnostics.rs

@@ -2,7 +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_data_structures::fx::FxHashMap;
 use rustc_errors::{Applicability, Diag, DiagCtxtHandle, DiagMessage, pluralize};
 use rustc_hir::attrs::diagnostic::{CustomDiagnostic, Directive, FormatArgs};
 use rustc_macros::Subdiagnostic;
@@ -155,21 +155,26 @@ struct CollectTrackerAndEmitter<'dcx, 'matcher> {
     // FIXME: Factor out a per-arm `Tracker` so that the `Option` is unnecessary.
     current: Option<(WhichMatcher, &'matcher [MatcherLoc])>,
 
-    /// Matches of [`MatcherLoc`]s that successfully consumed input from the parser.
+    /// Matches of [`MatcherLoc`]s.
     ///
-    /// This accumulates all calls to [`Tracker::matched_one()`]. It is used to identify all
-    /// competing matches for ambiguity errors.
-    matches: FxHashSet<SuccessfulMatch>,
+    /// This accumulates all calls to [`Tracker::trying_match()`] and [`Tracker::matched_one()`]. It
+    /// is used to identify all competing matches for ambiguity errors, and relevant match attempts
+    /// for failures.
+    matches: FxHashMap<Match, MatchResult>,
+
+    /// Tokens seen during parsing.
+    tokens: FxHashMap<u32, Token>,
+
+    input_pos_offset: u32,
 
-    remaining_matcher: Option<&'matcher MatcherLoc>,
     /// Which arm's failure should we report? (the one furthest along)
     best_failure: Option<BestFailure>,
     root_span: Span,
     result: Option<(Span, ErrorGuaranteed)>,
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
-struct SuccessfulMatch {
+struct Match {
     /// The position in the parser.
     ///
     /// As per [`Parser::approx_token_stream_pos()`].
@@ -179,6 +184,12 @@ struct SuccessfulMatch {
     loc_index: u32,
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+enum MatchResult {
+    Success,
+    Failure,
+}
+
 struct BestFailure {
     token: Token,
 
@@ -209,19 +220,27 @@ impl<'dcx, 'matcher> Tracker<'matcher> for CollectTrackerAndEmitter<'dcx, 'match
         self.current = Some((which_matcher, matcher));
     }
 
-    fn before_match_loc(&mut self, parser: &TtParser, matcher: &'matcher MatcherLoc) {
-        if self.remaining_matcher.is_none()
-            || (parser.has_no_remaining_items_for_step() && *matcher != MatcherLoc::Eof)
-        {
-            self.remaining_matcher = Some(matcher);
-        }
-    }
+    fn trying_match(&mut self, input_pos: u32, token: &Token, loc_index: u32) {
+        let input_pos = self.input_pos_offset + input_pos;
+        let old_token = self.tokens.insert(input_pos, *token);
+        debug_assert!(old_token.is_none_or(|t| t == *token));
 
-    fn matched_one(&mut self, parser: &Parser<'_>, loc_index: usize) {
-        let input_pos = parser.approx_token_stream_pos();
+        // Insert failure for now, will be updated in `matched_one()`.
         let loc_index: u32 = loc_index.try_into().unwrap();
-        let m = SuccessfulMatch { input_pos, loc_index };
-        self.matches.insert(m);
+        let m = Match { input_pos, loc_index };
+        self.matches.entry(m).or_insert(MatchResult::Failure);
+    }
+
+    fn matched_one(&mut self, input_pos: u32, loc_index: u32) {
+        let input_pos = self.input_pos_offset + input_pos;
+        let m = Match { input_pos, loc_index };
+        let match_result =
+            self.matches.get_mut(&m).unwrap_or_else(|| bug!("no corresponding `trying_match()`"));
+        *match_result = MatchResult::Success;
+    }
+
+    fn reset_input_pos(&mut self, parser: &Parser<'_>) {
+        self.input_pos_offset = parser.approx_token_stream_pos();
     }
 
     fn after_arm(&mut self, result: &NamedParseResult) {
@@ -249,10 +268,12 @@ impl<'dcx, 'matcher> Tracker<'matcher> for CollectTrackerAndEmitter<'dcx, 'match
 
         self.current = None;
         self.matches.clear();
+        self.tokens.clear();
+        self.input_pos_offset = 0;
     }
 
     fn failure(&mut self, parser: &Parser<'_>) {
-        let Some((which_matcher, _)) = self.current else {
+        let Some((which_matcher, matcher)) = self.current else {
             bug!("`Self::prepare()` was not called to initialize context");
         };
 
@@ -275,20 +296,37 @@ impl<'dcx, 'matcher> Tracker<'matcher> for CollectTrackerAndEmitter<'dcx, 'match
             .as_ref()
             .is_none_or(|failure| failure.is_better_position(which_matcher, approx_position))
         {
+            // Use the furthest-along non-EOF matcher. If none exists, use an EOF matcher.
+            #[expect(rustc::potential_query_instability, reason = "finding a unique maximum value")]
+            let (&Match { input_pos, loc_index }, &result) = self
+                .matches
+                .iter()
+                .max_by_key(|&(m, _res)| {
+                    // sort EOFs before others to deprioritize them
+                    let eof = match matcher[m.loc_index as usize] {
+                        MatcherLoc::Eof => 0,
+                        _ => 1,
+                    };
+                    (eof, m)
+                })
+                .unwrap_or_else(|| bug!("failure without a `trying_match()` call"));
+            // NOTE: `input_pos` might differ from `approx_position` (it might be a little older,
+            // if there were no non-EOF candidates at the right position), and it might have been
+            // a successful match for the same reason.
+            let _ = (input_pos, result);
+            let matcher = matcher[loc_index as usize].clone();
+
             self.best_failure = Some(BestFailure {
                 token,
                 matcher: which_matcher,
                 position: approx_position,
                 msg,
-                remaining_matcher: self
-                    .remaining_matcher
-                    .expect("must have collected matcher already")
-                    .clone(),
+                remaining_matcher: matcher,
             })
         }
     }
 
-    fn ambiguity(&mut self, parser: &Parser<'_>) {
+    fn ambiguity(&mut self) {
         let Some((_, matcher)) = self.current else {
             bug!("`Self::prepare()` was not called to initialize context");
         };
@@ -297,22 +335,40 @@ impl<'dcx, 'matcher> Tracker<'matcher> for CollectTrackerAndEmitter<'dcx, 'match
             rustc::potential_query_instability,
             reason = "sorting the results deterministically afterwards"
         )]
-        let (mut bb_locs, mut next_locs) = self
+        let mut matches = self
             .matches
             .iter()
-            .filter(|m| m.input_pos == parser.approx_token_stream_pos())
-            .partition::<Vec<&SuccessfulMatch>, _>(|m| {
+            .filter_map(|(m, result)| matches!(result, MatchResult::Success).then_some(m))
+            .collect::<Vec<_>>();
+        // Sort by input position, then `MatcherLoc` index.
+        matches.sort_unstable();
+
+        // Identify the earliest position where ambiguity occurred.
+        let input_pos = matches
+            .array_windows::<2>()
+            .find(|ms @ [a, b]| {
+                let mut locs = ms.iter().map(|x| &matcher[x.loc_index as usize]);
+                a.input_pos == b.input_pos
+                    && locs
+                        .any(|loc| matches!(loc, MatcherLoc::MetaVarDecl { .. } | MatcherLoc::Eof))
+            })
+            .map(|[a, _]| a.input_pos)
+            .unwrap_or_else(|| bug!("no ambiguity detected"));
+
+        let (bb_locs, next_locs) =
+            matches.iter().filter(|m| m.input_pos == input_pos).partition::<Vec<&Match>, _>(|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);
+        debug_assert!(bb_locs.iter().is_sorted());
+        debug_assert!(next_locs.iter().is_sorted());
+
+        let token = *self.tokens.get(&input_pos).unwrap();
 
-        let span = parser.token.span.substitute_dummy(self.root_span);
+        let span = token.span.substitute_dummy(self.root_span);
 
-        if parser.token == token::Eof {
+        if token == token::Eof {
             let msg = "ambiguity: multiple successful parses".to_string();
             let guar = self.dcx.span_err(span, msg);
             self.result = Some((span, guar));
@@ -357,8 +413,9 @@ impl<'dcx> CollectTrackerAndEmitter<'dcx, '_> {
             macro_name,
             dcx,
             current: None,
-            matches: FxHashSet::default(),
-            remaining_matcher: None,
+            matches: FxHashMap::default(),
+            tokens: FxHashMap::default(),
+            input_pos_offset: 0,
             best_failure: None,
             root_span,
             result: None,