From 07b9123f3bc797c05e3cedf63b119743c2521631 Mon Sep 17 00:00:00 2001 From: xmakro Date: Tue, 8 Sep 2026 22:46:02 -0700 Subject: [PATCH] Index macro arms lazily by their first literal token [skip ci] --- compiler/rustc_expand/src/mbe/diagnostics.rs | 2 +- compiler/rustc_expand/src/mbe/macro_parser.rs | 2 +- compiler/rustc_expand/src/mbe/macro_rules.rs | 147 +++++++++++++++++- 3 files changed, 145 insertions(+), 6 deletions(-) diff --git a/compiler/rustc_expand/src/mbe/diagnostics.rs b/compiler/rustc_expand/src/mbe/diagnostics.rs index 0e79dadb3503e..9243cb8fe13ee 100644 --- a/compiler/rustc_expand/src/mbe/diagnostics.rs +++ b/compiler/rustc_expand/src/mbe/diagnostics.rs @@ -48,7 +48,7 @@ pub(super) fn failed_to_match_macro( let mut tracker = CollectTrackerAndEmitter::new(name, psess.dcx(), sp); let try_success_result = match args { - FailedMacro::Func => try_match_macro(psess, body, rules, &mut tracker), + FailedMacro::Func => try_match_macro(psess, body, rules, None, &mut tracker), FailedMacro::Attr(attr_args) => { try_match_macro_attr(psess, attr_args, body, rules, &mut tracker) } diff --git a/compiler/rustc_expand/src/mbe/macro_parser.rs b/compiler/rustc_expand/src/mbe/macro_parser.rs index 3e94e0ca34773..a6ecc12ceb73e 100644 --- a/compiler/rustc_expand/src/mbe/macro_parser.rs +++ b/compiler/rustc_expand/src/mbe/macro_parser.rs @@ -393,7 +393,7 @@ impl NamedMatch { } /// Performs a token equality check, ignoring syntax context (that is, an unhygienic comparison) -fn token_name_eq(t1: &Token, t2: &Token) -> bool { +pub(super) fn token_name_eq(t1: &Token, t2: &Token) -> bool { if let (Some((ident1, is_raw1)), Some((ident2, is_raw2))) = (t1.ident(), t2.ident()) { ident1.name == ident2.name && is_raw1 == is_raw2 } else if let (Some((ident1, is_raw1)), Some((ident2, is_raw2))) = diff --git a/compiler/rustc_expand/src/mbe/macro_rules.rs b/compiler/rustc_expand/src/mbe/macro_rules.rs index 2212724c68bc1..7379fe991b1d9 100644 --- a/compiler/rustc_expand/src/mbe/macro_rules.rs +++ b/compiler/rustc_expand/src/mbe/macro_rules.rs @@ -1,6 +1,6 @@ use std::borrow::Cow; use std::collections::hash_map::Entry; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use std::{mem, slice}; use ast::token::IdentIsRaw; @@ -199,9 +199,68 @@ pub struct MacroRulesMacroExpander { transparency: Transparency, kinds: MacroKinds, rules: Vec, + rule_index: OnceLock>, macro_rules: bool, } +/// First-token pruning for larger function-style macros. Lists retain original arm order. +#[derive(Default, Debug)] +pub(super) struct RuleIndex { + by_first: FxHashMap>, + fallback: Vec, +} + +impl RuleIndex { + fn token_key(token: &Token) -> Option { + if let Some((ident, raw)) = token.ident() { + Some(TokenKind::Ident(ident.name, raw)) + } else if let Some((ident, raw)) = token.lifetime() { + Some(TokenKind::Lifetime(ident.name, raw)) + } else if matches!(token.kind, OpenInvisible(_) | CloseInvisible(_)) { + None + } else { + Some(token.kind) + } + } + + fn new(rules: &[MacroRule]) -> Self { + let mut index = Self::default(); + for (i, rule) in rules.iter().enumerate() { + let MacroRule::Func { lhs, .. } = rule else { continue }; + if let Some(MatcherLoc::Token { token }) = lhs.first() + && !matches!(token.kind, DocComment(..)) + && let Some(key) = Self::token_key(token) + { + index.by_first.entry(key).or_default().push(i); + } else { + // Metavariables, repetitions, ignored doc comments and delimiters retain + // the original matcher, including ambiguity and feature-gating behavior. + index.fallback.push(i); + } + } + index + } + + fn candidates(&self, token: &Token) -> impl Iterator { + let mut matching = Self::token_key(token) + .and_then(|key| self.by_first.get(&key)) + .map_or(&[][..], Vec::as_slice); + let mut fallback = self.fallback.as_slice(); + std::iter::from_fn(move || { + let next = match (matching.first(), fallback.first()) { + (Some(a), Some(b)) if a < b => &mut matching, + (Some(_), Some(_)) => &mut fallback, + (Some(_), None) => &mut matching, + (None, Some(_)) => &mut fallback, + (None, None) => return None, + }; + let i = next[0]; + *next = &next[1..]; + Some(i) + }) + } +} + impl MacroRulesMacroExpander { pub fn get_unused_rule(&self, rule_i: usize) -> Option<(&Ident, MultiSpan)> { // If the rhs contains an invocation like `compile_error!`, don't report it as unused. @@ -303,6 +362,9 @@ impl TTMacroExpander for MacroRulesMacroExpander { self.transparency, input, &self.rules, + (self.rules.len() >= 8).then(|| { + self.rule_index.get_or_init(|| Box::new(RuleIndex::new(&self.rules))).as_ref() + }), self.on_unmatched_args.as_ref(), )) } @@ -362,6 +424,9 @@ fn trace_macros_note(cx_expansions: &mut FxIndexMap>, sp: Span } pub(super) trait Tracker<'matcher> { + /// Diagnostic trackers need every failure callback, even for an immediate mismatch. + const TRACK_FAILURES: bool = true; + /// Provide context on the arm that's about to be matched. fn prepare(&mut self, which_matcher: WhichMatcher, matcher: &'matcher [MatcherLoc]); @@ -405,6 +470,8 @@ pub(super) trait Tracker<'matcher> { pub(super) struct NoopTracker; impl<'matcher> Tracker<'matcher> for NoopTracker { + const TRACK_FAILURES: bool = false; + fn prepare(&mut self, _which_matcher: WhichMatcher, _matcher: &'matcher [MatcherLoc]) {} fn before_match_loc(&mut self, _parser: &TtParser, _matcher: &'matcher MatcherLoc) {} @@ -427,7 +494,7 @@ impl<'matcher> Tracker<'matcher> for NoopTracker { } /// Expands the rules based macro defined by `rules` for a given input `arg`. -#[instrument(skip(cx, transparency, arg, rules, on_unmatched_args))] +#[instrument(skip(cx, transparency, arg, rules, rule_index, on_unmatched_args))] fn expand_macro<'cx, 'a: 'cx>( cx: &'cx mut ExtCtxt<'_>, sp: Span, @@ -437,6 +504,7 @@ fn expand_macro<'cx, 'a: 'cx>( transparency: Transparency, arg: TokenStream, rules: &'a [MacroRule], + rule_index: Option<&RuleIndex>, on_unmatched_args: Option<&Directive>, ) -> Box { let psess = &cx.sess.psess; @@ -447,7 +515,7 @@ fn expand_macro<'cx, 'a: 'cx>( } // Track nothing for the best performance. - let try_success_result = try_match_macro(psess, &arg, rules, &mut NoopTracker); + let try_success_result = try_match_macro(psess, &arg, rules, rule_index, &mut NoopTracker); match try_success_result { Ok((rule_index, rule, named_matches)) => { @@ -608,6 +676,7 @@ pub(super) fn try_match_macro<'matcher, T: Tracker<'matcher>>( psess: &ParseSess, arg: &TokenStream, rules: &'matcher [MacroRule], + rule_index: Option<&RuleIndex>, track: &mut T, ) -> Result<(usize, &'matcher MacroRule, NamedMatches), CanRetry> { // We create a base parser that can be used for the "black box" parts. @@ -632,7 +701,15 @@ pub(super) fn try_match_macro<'matcher, T: Tracker<'matcher>>( let parser = parser_from_cx(psess, arg.clone(), T::recovery()); // Try each arm's matchers. let mut tt_parser = TtParser::new(); - for (i, rule) in rules.iter().enumerate() { + let mut all_rules = 0..rules.len(); + let mut indexed_rules = + rule_index.filter(|_| !T::TRACK_FAILURES).map(|index| index.candidates(&parser.token)); + let rule_indices = std::iter::from_fn(|| match &mut indexed_rules { + Some(indices) => indices.next(), + None => all_rules.next(), + }); + for i in rule_indices { + let rule = &rules[i]; let MacroRule::Func { lhs, .. } = rule else { continue }; let _tracing_span = trace_span!("Matching arm", %i); @@ -940,6 +1017,7 @@ pub fn compile_declarative_macro( on_unmatched_args, transparency, rules, + rule_index: OnceLock::new(), macro_rules, }; mk_syn_ext(SyntaxExtensionKind::MacroRules(Arc::new(exp))) @@ -1868,3 +1946,64 @@ pub(super) fn parser_from_cx( tts.desugar_doc_comments(); Parser::new(psess, tts, rustc_parse::MACRO_ARGUMENTS).recovery(recovery) } + +#[cfg(test)] +mod rule_index_tests { + use super::*; + use crate::mbe::macro_parser::token_name_eq; + + #[test] + fn first_token_index_preserves_scalar_filter_order() { + rustc_span::create_default_session_globals_then(|| { + let span = rustc_span::DUMMY_SP; + let x = sym::cfg; + let ident = Ident::with_dummy_span(x); + let kinds = [ + Eq, + Comma, + TokenKind::Ident(x, IdentIsRaw::No), + TokenKind::Ident(x, IdentIsRaw::Yes), + NtIdent(ident, IdentIsRaw::No), + Lifetime(x, IdentIsRaw::No), + NtLifetime(ident, IdentIsRaw::No), + OpenInvisible(token::InvisibleOrigin::ProcMacro), + CloseInvisible(token::InvisibleOrigin::ProcMacro), + DocComment(token::CommentKind::Line, ast::AttrStyle::Outer, x), + Eof, + ]; + let rules: Vec<_> = (0..55) + .map(|i| { + let token = Token::new(kinds[i % kinds.len()], span); + let first = if i % 7 == 0 { + MatcherLoc::Delimited + } else { + MatcherLoc::Token { token } + }; + MacroRule::Func { + lhs: vec![first, MatcherLoc::Eof], + lhs_span: span, + rhs: mbe::TokenTree::Token(token), + } + }) + .collect(); + let index = RuleIndex::new(&rules); + for kind in kinds { + let token = Token::new(kind, span); + let eligible = |i: &usize| match &rules[*i] { + MacroRule::Func { lhs, .. } => match lhs.first() { + Some(MatcherLoc::Token { token: first }) + if !matches!(first.kind, DocComment(..)) => + { + token_name_eq(first, &token) + } + _ => true, + }, + _ => false, + }; + let expected: Vec<_> = (0..rules.len()).filter(eligible).collect(); + let actual: Vec<_> = index.candidates(&token).filter(eligible).collect(); + assert_eq!(actual, expected, "token={kind:?}"); + } + }); + } +}