diff --git a/README.md b/README.md index 3cbf7aa6..304466b6 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ For Rust projects, add the runtime crate: ```toml [dependencies] -antlr-rust-runtime = "0.5" +antlr-rust-runtime = "0.6" ``` The library crate is imported as `antlr4_runtime`: @@ -38,7 +38,9 @@ cargo install antlr-rust-runtime ``` This installs `antlr4-rust-gen`, which turns ANTLR `.interp` metadata into Rust -lexer and parser modules. +lexer and parser modules. During generation it also compiles the lexer's DFA +ahead of time and embeds the tables in the generated lexer, so tokenization +runs at full speed from the first character with no per-process warmup. ### 3. Generate your parser @@ -212,6 +214,10 @@ The runtime contains: - ANTLR v4 serialized ATN deserialization - lexer ATN recognition with longest-match/rule-priority behavior and lexer actions +- ahead-of-time compiled lexer DFA tables, built by `antlr4-rust-gen` and + embedded in generated lexers, with per-token escape to ATN interpretation + for constructs a finite DFA cannot represent (semantic predicates, + recursive lexer rules) - parser ATN rule recognition with backtracking over token stream indices - `antlr4-rust-gen`, a Rust generator that consumes ANTLR `.interp` metadata and emits Rust modules @@ -280,11 +286,13 @@ group (**> 1.0** means Rust is faster than Go; **< 1.0** means slower): Rust is faster than Go on every fixture in all four language groups, with Kotlin leading dramatically (expression-ladder memoization in the generated -walker). Learned lexer and parser DFAs are shared across recognizer -instances, so repeated parses of the same grammar — the common case for a -CLI tool or language server — skip relearning entirely. Numbers are -warm-parse minimums on an Apple M3 Pro and are indicative — re-run the -benchmark on your own hardware for authoritative figures. +walker). Lexer DFAs are compiled at generation time and embedded in the +generated lexer, so tokenization needs no warmup at all; learned parser +decision DFAs are shared across parser instances, so repeated parses of the +same grammar — the common case for a CLI tool or language server — skip +relearning entirely. Numbers are warm-parse minimums on an Apple M3 Pro and +are indicative — re-run the benchmark on your own hardware for authoritative +figures. ## Useful Information diff --git a/src/atn/lexer.rs b/src/atn/lexer.rs index 21a04c55..16e29791 100644 --- a/src/atn/lexer.rs +++ b/src/atn/lexer.rs @@ -1,6 +1,7 @@ use std::collections::{BTreeSet, HashSet}; use std::hash::BuildHasherDefault; +use crate::atn::lexer_dfa::{CompiledLexerAccept, CompiledLexerDfa, DEAD_STATE, ESCAPE_STATE}; use crate::atn::{Atn, AtnStateKind, LexerAction, Transition}; use crate::char_stream::{CharStream, TextInterval}; use crate::int_stream::EOF; @@ -18,33 +19,33 @@ const MIN_CHAR_VALUE: i32 = 0; const MAX_CHAR_VALUE: i32 = 0x0010_FFFF; #[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] -struct LexerConfig { - state: usize, - position: usize, - consumed_eof: bool, - alt_rule_index: Option, - passed_non_greedy: bool, - stack: Vec, - actions: Vec, +pub(super) struct LexerConfig { + pub(super) state: usize, + pub(super) position: usize, + pub(super) consumed_eof: bool, + pub(super) alt_rule_index: Option, + pub(super) passed_non_greedy: bool, + pub(super) stack: Vec, + pub(super) actions: Vec, } #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] -struct LexerActionTrace { - action_index: usize, - position: usize, +pub(super) struct LexerActionTrace { + pub(super) action_index: usize, + pub(super) position: usize, /// Lexer rule that the action transition belonged to. ANTLR suppresses /// commands of nested non-fragment rule references, so the dispatcher /// must compare this against the accepted rule before applying side /// effects like `pushMode` / `popMode`. - rule_index: usize, + pub(super) rule_index: usize, } #[derive(Clone, Debug)] -struct AcceptState { - position: usize, - rule_index: usize, - consumed_eof: bool, - actions: Vec, +pub(super) struct AcceptState { + pub(super) position: usize, + pub(super) rule_index: usize, + pub(super) consumed_eof: bool, + pub(super) actions: Vec, } #[derive(Clone, Debug)] @@ -54,9 +55,9 @@ enum MatchResult { } #[derive(Clone, Debug)] -struct ClosureResult { - configs: Vec, - has_semantic_context: bool, +pub(super) struct ClosureResult { + pub(super) configs: Vec, + pub(super) has_semantic_context: bool, } /// Mutable emission state produced by executing lexer actions for one token. @@ -215,7 +216,70 @@ where &mut custom_action, &mut semantic_predicate, &mut accept_adjuster, - false, + LexerMatchStrategy { + compiled: None, + use_cache: false, + }, + ) +} + +/// Runs one lexer-token match against an ahead-of-time compiled lexer DFA. +/// +/// Tokens starting in a compiled mode are matched by walking static tables; +/// modes the compiler left dynamic fall back to cached ATN interpretation per +/// token, so behavior always matches [`next_token`]. +pub fn next_token_compiled( + lexer: &mut BaseLexer, + atn: &Atn, + dfa: &CompiledLexerDfa, +) -> CommonToken +where + I: CharStream, + F: TokenFactory, +{ + next_token_with_hooks_impl( + lexer, + atn, + &mut |_, _| {}, + &mut |_, _| true, + &mut |_, _, _| {}, + LexerMatchStrategy { + compiled: Some(dfa), + use_cache: true, + }, + ) +} + +/// Runs one compiled-DFA lexer-token match with all generated extension hooks. +/// +/// Compiled modes never contain semantic predicates, so hook grammars still +/// take the table walk for their static modes; predicate-bearing modes re-run +/// the ATN interpreter exactly like [`next_token_with_hooks`]. +pub fn next_token_compiled_with_hooks( + lexer: &mut BaseLexer, + atn: &Atn, + dfa: &CompiledLexerDfa, + mut custom_action: A, + mut semantic_predicate: P, + mut accept_adjuster: E, +) -> CommonToken +where + I: CharStream, + F: TokenFactory, + A: FnMut(&mut BaseLexer, LexerCustomAction), + P: FnMut(&BaseLexer, LexerPredicate) -> bool, + E: FnMut(&mut BaseLexer, i32, usize), +{ + next_token_with_hooks_impl( + lexer, + atn, + &mut custom_action, + &mut semantic_predicate, + &mut accept_adjuster, + LexerMatchStrategy { + compiled: Some(dfa), + use_cache: false, + }, ) } @@ -239,17 +303,59 @@ where &mut custom_action, &mut semantic_predicate, &mut accept_adjuster, - true, + LexerMatchStrategy { + compiled: None, + use_cache: true, + }, ) } +/// Token-matching backend chosen by a lexer entry point: an optional +/// ahead-of-time compiled DFA, and whether ATN interpretation (used directly +/// or as the compiled path's per-mode fallback) may replay the learned-DFA +/// cache. Hook entry points disable the cache, matching their interpreted +/// counterparts. +#[derive(Clone, Copy)] +struct LexerMatchStrategy<'a> { + compiled: Option<&'a CompiledLexerDfa>, + use_cache: bool, +} + +/// Dispatches one token match to the strategy's backend. +fn match_token_with_strategy( + lexer: &mut BaseLexer, + atn: &Atn, + mode: i32, + start: usize, + semantic_predicate: &mut P, + strategy: LexerMatchStrategy<'_>, +) -> MatchResult +where + I: CharStream, + F: TokenFactory, + P: FnMut(&BaseLexer, LexerPredicate) -> bool, +{ + if let Some(dfa) = strategy.compiled + && !lexer.force_interpreted() + && let Some(start_state) = dfa.mode_start(mode) + && let Some(result) = match_token_compiled(lexer, dfa, start_state, start) + { + return result; + } + if strategy.use_cache { + match_token_cached(lexer, atn, mode, start, semantic_predicate) + } else { + match_token(lexer, atn, mode, start, semantic_predicate) + } +} + fn next_token_with_hooks_impl( lexer: &mut BaseLexer, atn: &Atn, custom_action: &mut A, semantic_predicate: &mut P, accept_adjuster: &mut E, - use_cache: bool, + strategy: LexerMatchStrategy<'_>, ) -> CommonToken where I: CharStream, @@ -269,11 +375,8 @@ where } let mode = lexer.mode(); let start = lexer.input().index(); - let token_match = if use_cache { - match_token_cached(lexer, atn, mode, start, semantic_predicate) - } else { - match_token(lexer, atn, mode, start, semantic_predicate) - }; + let token_match = + match_token_with_strategy(lexer, atn, mode, start, semantic_predicate, strategy); let accept = match token_match { MatchResult::Accept(accept) => accept, MatchResult::NoViableAlt { stop } => { @@ -353,7 +456,11 @@ where /// text, but its embedded action does not run unless that rule itself accepts /// the token. Fragment-rule actions remain eligible because fragments have no /// token type of their own. -fn lexer_action_belongs_to_accept(atn: &Atn, accept_rule: usize, action_rule: usize) -> bool { +pub(super) fn lexer_action_belongs_to_accept( + atn: &Atn, + accept_rule: usize, + action_rule: usize, +) -> bool { action_rule == accept_rule || atn .rule_to_token_type() @@ -389,7 +496,6 @@ where return MatchResult::NoViableAlt { stop: start }; }; let start_closure = epsilon_closure( - lexer, atn, [LexerConfig { state: start_state, @@ -400,7 +506,7 @@ where stack: Vec::new(), actions: Vec::new(), }], - semantic_predicate, + &mut |predicate| semantic_predicate(lexer, predicate), ); let mut active = prune_after_accepts(atn, start_closure.configs); let mut dfa_state = lexer.lexer_dfa_state( @@ -440,7 +546,9 @@ where } } - let closure = epsilon_closure(lexer, atn, next, semantic_predicate); + let closure = epsilon_closure(atn, next, &mut |predicate| { + semantic_predicate(lexer, predicate) + }); let target_has_semantic_context = closure.has_semantic_context; let suppress_edge = source_has_semantic_context || target_has_semantic_context; active = prune_after_accepts(atn, closure.configs); @@ -556,7 +664,9 @@ where } } - let closure = epsilon_closure(lexer, atn, next, semantic_predicate); + let closure = epsilon_closure(atn, next, &mut |predicate| { + semantic_predicate(lexer, predicate) + }); let target_has_semantic_context = closure.has_semantic_context; if target_has_semantic_context { return match_token(lexer, atn, mode, start, semantic_predicate); @@ -597,6 +707,98 @@ where ) } +/// Bounds EOF-edge traversals in the compiled walk. EOF transitions do not +/// advance the cursor, so past this bound the walk stops guessing and escapes +/// to the ATN interpreter, which owns the semantics of longer EOF chains +/// (including grammars whose chains never terminate). +const MAX_COMPILED_EOF_EDGES: u32 = 8; + +/// Matches one token by walking the ahead-of-time compiled lexer DFA. +/// +/// The walk reproduces the interpreter's longest-match selection: remember +/// the best accept seen so far, advance until the table has no transition, +/// then return the remembered accept — or a recognition error spanning every +/// character the walk looked at, exactly like `match_token`. Reaching an +/// escape edge (semantic predicate, recursive lexer rule, state budget) +/// returns `None`, and the caller re-matches the token with the interpreter. +fn match_token_compiled( + lexer: &mut BaseLexer, + dfa: &CompiledLexerDfa, + start_state: u16, + start: usize, +) -> Option +where + I: CharStream, + F: TokenFactory, +{ + let mut state = start_state; + let mut position = start; + let mut best: Option = None; + let mut error_stop = start; + let mut eof_edges = 0_u32; + loop { + if let Some(accept) = dfa.accept(state) { + record_compiled_accept(accept, position, &mut best); + } + let symbol = symbol_at(lexer, position); + let target = if symbol == EOF { + eof_edges += 1; + if eof_edges > MAX_COMPILED_EOF_EDGES { + return None; + } + dfa.eof_target(state) + } else { + error_stop = error_stop.max(position.saturating_add(1)); + dfa.char_target(state, symbol) + }; + if target == DEAD_STATE { + break; + } + if target == ESCAPE_STATE { + return None; + } + if symbol != EOF { + position += 1; + } + state = target; + } + Some(best.map_or( + MatchResult::NoViableAlt { stop: error_stop }, + MatchResult::Accept, + )) +} + +/// Applies the interpreter's longest-match / lowest-rule preference to one +/// compiled accept state, materializing its action traces at absolute input +/// positions. +fn record_compiled_accept( + accept: &CompiledLexerAccept, + position: usize, + best: &mut Option, +) { + let replaces = best.as_ref().is_none_or(|current| { + position > current.position + || (position == current.position && accept.rule_index < current.rule_index) + }); + if !replaces { + return; + } + *best = Some(AcceptState { + position, + rule_index: accept.rule_index, + consumed_eof: accept.consumed_eof, + actions: accept + .actions + .iter() + .map(|trace| LexerActionTrace { + action_index: trace.action_index, + position: position.saturating_sub(trace.behind), + rule_index: trace.rule_index, + }) + .collect(), + }); +} + fn cached_mode_start_state( lexer: &BaseLexer, atn: &Atn, @@ -616,7 +818,6 @@ where let mode_index = usize::try_from(mode).ok()?; let start_state = atn.mode_to_start_state().get(mode_index).copied()?; let start_closure = epsilon_closure( - lexer, atn, [LexerConfig { state: start_state, @@ -627,7 +828,7 @@ where stack: Vec::new(), actions: Vec::new(), }], - semantic_predicate, + &mut |predicate| semantic_predicate(lexer, predicate), ); let active = prune_after_accepts(atn, start_closure.configs); let state = cache_dfa_state( @@ -717,16 +918,13 @@ fn cached_accept_state( /// Lexer rule calls use an explicit return-state stack in `LexerConfig` because /// fragment rules and nested lexer constructs compile to rule transitions in the /// serialized ATN. -fn epsilon_closure( - lexer: &BaseLexer, +pub(super) fn epsilon_closure

( atn: &Atn, configs: impl IntoIterator, semantic_predicate: &mut P, ) -> ClosureResult where - I: CharStream, - F: TokenFactory, - P: FnMut(&BaseLexer, LexerPredicate) -> bool, + P: FnMut(LexerPredicate) -> bool, { let mut state = ClosureState { seen: FxHashSet::default(), @@ -735,7 +933,7 @@ where }; for config in configs { - close_config(lexer, atn, config, &mut state, semantic_predicate); + close_config(atn, config, &mut state, semantic_predicate); } ClosureResult { @@ -750,16 +948,13 @@ where /// Ordered DFS matters for lexer greediness: greedy loop entries serialize the /// loop path before the exit path, while non-greedy entries serialize the exit /// path first. The later accept-pruning step relies on this order. -fn close_config( - lexer: &BaseLexer, +fn close_config

( atn: &Atn, config: LexerConfig, closure: &mut ClosureState, semantic_predicate: &mut P, ) where - I: CharStream, - F: TokenFactory, - P: FnMut(&BaseLexer, LexerPredicate) -> bool, + P: FnMut(LexerPredicate) -> bool, { if !closure.seen.insert(config.clone()) { return; @@ -774,7 +969,7 @@ fn close_config( let mut returned = config.clone(); set_config_state(atn, &mut returned, follow_state); returned.stack = rest.to_vec(); - close_config(lexer, atn, returned, closure, semantic_predicate); + close_config(atn, returned, closure, semantic_predicate); } closure.closed.push(config); return; @@ -786,7 +981,7 @@ fn close_config( let mut next = config.clone(); set_config_state(atn, &mut next, *target); next.passed_non_greedy |= state.non_greedy; - close_config(lexer, atn, next, closure, semantic_predicate); + close_config(atn, next, closure, semantic_predicate); } Transition::Rule { target, @@ -797,7 +992,7 @@ fn close_config( set_config_state(atn, &mut next, *target); next.passed_non_greedy |= state.non_greedy; next.stack.push(*follow_state); - close_config(lexer, atn, next, closure, semantic_predicate); + close_config(atn, next, closure, semantic_predicate); } Transition::Predicate { target, @@ -806,21 +1001,22 @@ fn close_config( .. } => { closure.has_semantic_context = true; - if semantic_predicate( - lexer, - LexerPredicate::new(*rule_index, *pred_index, config.position), - ) { + if semantic_predicate(LexerPredicate::new( + *rule_index, + *pred_index, + config.position, + )) { let mut next = config.clone(); set_config_state(atn, &mut next, *target); next.passed_non_greedy |= state.non_greedy; - close_config(lexer, atn, next, closure, semantic_predicate); + close_config(atn, next, closure, semantic_predicate); } } Transition::Precedence { target, .. } => { let mut next = config.clone(); set_config_state(atn, &mut next, *target); next.passed_non_greedy |= state.non_greedy; - close_config(lexer, atn, next, closure, semantic_predicate); + close_config(atn, next, closure, semantic_predicate); } Transition::Action { target, @@ -838,7 +1034,7 @@ fn close_config( rule_index: *rule_index, }); } - close_config(lexer, atn, next, closure, semantic_predicate); + close_config(atn, next, closure, semantic_predicate); } Transition::Atom { .. } | Transition::Range { .. } @@ -864,7 +1060,7 @@ fn close_config( /// Once such a path reaches the rule stop state, later same-rule configs should /// not continue to grow into a longer token. Greedy decisions still need all /// paths to remain available so longest-match selection can win. -fn prune_after_accepts(atn: &Atn, configs: Vec) -> Vec { +pub(super) fn prune_after_accepts(atn: &Atn, configs: Vec) -> Vec { let mut accepted_rules = BTreeSet::new(); let mut pruned = Vec::with_capacity(configs.len()); for config in configs { @@ -892,7 +1088,7 @@ fn prune_after_accepts(atn: &Atn, configs: Vec) -> Vec /// ANTLR lexer priority is encoded by rule order. `match_token` already handles /// longest-match selection across input positions; within a single position the /// lower rule index wins. -fn best_accept(atn: &Atn, configs: &[LexerConfig]) -> Option { +pub(super) fn best_accept(atn: &Atn, configs: &[LexerConfig]) -> Option { configs .iter() .filter_map(|config| { @@ -992,7 +1188,7 @@ fn cached_configs_to_configs( /// Moves a lexer config to `state_number` and records the top-level lexer rule /// once the config leaves a mode start state. -fn set_config_state(atn: &Atn, config: &mut LexerConfig, state_number: usize) { +pub(super) fn set_config_state(atn: &Atn, config: &mut LexerConfig, state_number: usize) { config.state = state_number; if config.alt_rule_index.is_none() { config.alt_rule_index = atn.state(state_number).and_then(|state| state.rule_index); diff --git a/src/atn/lexer_dfa.rs b/src/atn/lexer_dfa.rs new file mode 100644 index 00000000..2c687405 --- /dev/null +++ b/src/atn/lexer_dfa.rs @@ -0,0 +1,1124 @@ +//! Ahead-of-time lexer DFA compilation. +//! +//! ANTLR runtimes normally discover the lexer DFA lazily: the ATN simulator +//! computes epsilon closures per input character and caches the resulting +//! config sets. This module runs the same subset construction eagerly over +//! the entire character alphabet, once per grammar, so token matching becomes +//! one table lookup per character with no closure computation, hashing, or +//! config allocation on the hot path. +//! +//! Compilation is conservative at the edge level: a transition whose target +//! closure crosses a semantic predicate (whose outcome exists only at parse +//! time), grows an unbounded rule-call stack (recursive lexer rules such as +//! nested comments), or exceeds the state budget is compiled as an *escape* +//! edge. A token walk that reaches an escape edge is re-matched from the +//! token start by the ATN interpreter, so rare dynamic constructs never +//! poison the rest of the mode. Because the construction reuses the +//! interpreter's own closure, pruning, and accept-selection code, a compiled +//! walk that does not escape reproduces interpreter behavior exactly. + +use std::collections::HashMap; +use std::hash::BuildHasherDefault; + +use crate::atn::lexer::{ + LexerConfig, best_accept, epsilon_closure, lexer_action_belongs_to_accept, + prune_after_accepts, set_config_state, +}; +use crate::atn::{Atn, Transition}; +use crate::int_stream::EOF; +use crate::lexer::{LexerDfaActionKey, LexerDfaConfigKey, LexerDfaKey}; +use crate::prediction::PredictionFxHasher; + +#[allow(clippy::disallowed_types)] +type FxHashMap = HashMap>; + +const MIN_CHAR_VALUE: i32 = 0; +const MAX_CHAR_VALUE: i32 = 0x0010_FFFF; + +/// Sentinel state id meaning "no transition". +pub(super) const DEAD_STATE: u16 = u16::MAX; + +/// Sentinel state id meaning "re-match this token with the ATN interpreter". +pub(super) const ESCAPE_STATE: u16 = u16::MAX - 1; + +/// Per-mode state budget; targets past it compile as escape edges. The cap +/// also bounds compile time for pathological grammars. +const MAX_MODE_STATES: usize = 4096; + +/// Rule-call stacks deeper than this escape to the interpreter, as a backstop +/// for grammars with extraordinarily long non-recursive fragment chains. +const MAX_STACK_DEPTH: usize = 32; + +/// Configs whose surviving action trace grows past this escape to the +/// interpreter: a custom action crossed inside a loop is genuinely +/// position-dependent and cannot compile to finitely many DFA states. +const MAX_ACTION_TRACES: usize = 16; + +/// Dense per-state edge row width, matching the interpreter's DFA cache rows. +const ASCII_EDGE_SYMBOLS: usize = 128; +/// [`ASCII_EDGE_SYMBOLS`] as a code point for segment arithmetic. +const ASCII_EDGE_LIMIT: i32 = 128; + +/// A lexer DFA compiled ahead of time from a lexer ATN. +/// +/// Build one per grammar with [`CompiledLexerDfa::compile`] (generated lexers +/// cache it in a `OnceLock` beside the deserialized ATN) and match tokens +/// through [`crate::atn::lexer::next_token_compiled`] or +/// [`crate::atn::lexer::next_token_compiled_with_hooks`]. +#[derive(Clone, Debug)] +pub struct CompiledLexerDfa { + mode_starts: Vec>, + states: Vec, + ascii_rows: Vec<[u16; ASCII_EDGE_SYMBOLS]>, + wide_rows: Vec>, + accepts: Vec, +} + +/// One compiled DFA state; rows are pooled indices because many states share +/// identical edge rows (identifier continuations, string bodies, …). +#[derive(Clone, Copy, Debug)] +struct CompiledLexerState { + ascii_row: u32, + wide_row: u32, + eof_target: u16, + /// Index into `accepts`, or `u32::MAX` when the state does not accept. + accept: u32, +} + +/// Inclusive code-point range above the ASCII row mapping to one target. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +struct WideRange { + low: u32, + high: u32, + target: u16, +} + +/// Accept metadata for one DFA state: the winning lexer rule plus the action +/// transitions collected on the accepted ATN path. +#[derive(Clone, Debug)] +pub(super) struct CompiledLexerAccept { + pub(super) rule_index: usize, + pub(super) consumed_eof: bool, + pub(super) actions: Vec, +} + +/// One recorded lexer action, positioned relative to the accept boundary so +/// the same DFA state serves every input offset. +#[derive(Clone, Copy, Debug)] +pub(super) struct CompiledLexerActionTrace { + pub(super) action_index: usize, + pub(super) rule_index: usize, + /// Characters consumed between the action transition and the accept. + pub(super) behind: usize, +} + +impl CompiledLexerDfa { + /// Compiles every lexer mode of `atn` that has no semantic predicates and + /// fits the state budget; the rest stay interpreter-matched. + pub fn compile(atn: &Atn) -> Self { + let mut dfa = Self { + mode_starts: Vec::new(), + states: Vec::new(), + ascii_rows: Vec::new(), + wide_rows: Vec::new(), + accepts: Vec::new(), + }; + let mut pools = RowPools::default(); + for mode in 0..atn.mode_to_start_state().len() { + let start = build_mode(atn, mode, &mut dfa, &mut pools); + dfa.mode_starts.push(start); + } + dfa + } + + /// True when at least one lexer mode compiled to static tables. + pub fn has_compiled_modes(&self) -> bool { + self.mode_starts.iter().any(Option::is_some) + } + + /// Number of compiled DFA states across all modes (diagnostics). + pub const fn state_count(&self) -> usize { + self.states.len() + } + + /// Per-mode compilation outcome (diagnostics): `true` = static tables. + pub fn compiled_mode_flags(&self) -> Vec { + self.mode_starts.iter().map(Option::is_some).collect() + } + + /// Per-mode state counts (diagnostics), derived from start offsets. + pub fn mode_state_counts(&self) -> Vec { + let mut starts: Vec = self + .mode_starts + .iter() + .flatten() + .map(|&start| usize::from(start)) + .collect(); + starts.push(self.states.len()); + starts.windows(2).map(|pair| pair[1] - pair[0]).collect() + } + + /// Compiled start state for `mode`, or `None` when the mode is + /// interpreter-matched. + pub(super) fn mode_start(&self, mode: i32) -> Option { + let mode = usize::try_from(mode).ok()?; + self.mode_starts.get(mode).copied().flatten() + } + + pub(super) fn accept(&self, state: u16) -> Option<&CompiledLexerAccept> { + self.accepts + .get(self.states[usize::from(state)].accept as usize) + } + + /// Transition target for a non-EOF symbol, or [`DEAD_STATE`]. + pub(super) fn char_target(&self, state: u16, symbol: i32) -> u16 { + let compiled = &self.states[usize::from(state)]; + let code_point = symbol.cast_unsigned(); + if let Ok(ascii) = usize::try_from(symbol) + && ascii < ASCII_EDGE_SYMBOLS + { + return self.ascii_rows[compiled.ascii_row as usize][ascii]; + } + let row = &self.wide_rows[compiled.wide_row as usize]; + match row.binary_search_by(|range| range.low.cmp(&code_point)) { + Ok(found) => row[found].target, + Err(insert) => { + if insert > 0 && row[insert - 1].high >= code_point { + row[insert - 1].target + } else { + DEAD_STATE + } + } + } + } + + /// Transition target for the EOF symbol, or [`DEAD_STATE`]. + pub(super) fn eof_target(&self, state: u16) -> u16 { + self.states[usize::from(state)].eof_target + } + + /// Flattens the compiled DFA into a `u32` stream for embedding in + /// generated code. + /// + /// The format is internal to this runtime version; [`Self::from_serialized`] + /// rejects streams from other versions so generated lexers can fall back + /// to [`Self::compile`]. + pub fn serialize(&self) -> Vec { + // Exact word count: the tag, five section-length words, and each + // section's payload (states are 4 words; ASCII rows pack 2 targets + // per word; wide ranges and action traces are 3 words each behind + // their per-row/per-accept length words). + let wide_words: usize = self.wide_rows.iter().map(|row| 1 + row.len() * 3).sum(); + let accept_words: usize = self + .accepts + .iter() + .map(|accept| 3 + accept.actions.len() * 3) + .sum(); + let capacity = 6 + + self.mode_starts.len() + + self.states.len() * 4 + + self.ascii_rows.len() * (ASCII_EDGE_SYMBOLS / 2) + + wide_words + + accept_words; + let mut out = Vec::with_capacity(capacity); + out.push(SERIALIZED_TAG); + out.push(self.mode_starts.len() as u32); + for start in &self.mode_starts { + out.push(start.map_or(u32::MAX, u32::from)); + } + out.push(self.states.len() as u32); + for state in &self.states { + out.push(state.ascii_row); + out.push(state.wide_row); + out.push(u32::from(state.eof_target)); + out.push(state.accept); + } + out.push(self.ascii_rows.len() as u32); + for row in &self.ascii_rows { + for pair in row.chunks(2) { + out.push(u32::from(pair[0]) | (u32::from(pair[1]) << 16)); + } + } + out.push(self.wide_rows.len() as u32); + for row in &self.wide_rows { + out.push(row.len() as u32); + for range in &**row { + out.push(range.low); + out.push(range.high); + out.push(u32::from(range.target)); + } + } + out.push(self.accepts.len() as u32); + for accept in &self.accepts { + out.push(accept.rule_index as u32); + out.push(u32::from(accept.consumed_eof)); + out.push(accept.actions.len() as u32); + for action in &accept.actions { + out.push(action.action_index as u32); + out.push(action.rule_index as u32); + out.push(action.behind as u32); + } + } + debug_assert_eq!(out.len(), capacity, "serialized stream fills its capacity exactly"); + out + } + + /// Rebuilds a compiled DFA from [`Self::serialize`] output; `None` when + /// the stream comes from a different runtime version or is malformed. + pub fn from_serialized(data: &[u32]) -> Option { + let mut reader = SerializedReader { data, position: 0 }; + if reader.next()? != SERIALIZED_TAG { + return None; + } + let mode_count = reader.next_len()?; + let mut mode_starts = Vec::with_capacity(mode_count); + for _ in 0..mode_count { + let word = reader.next()?; + let start = if word == u32::MAX { + None + } else { + Some(u16::try_from(word).ok()?) + }; + mode_starts.push(start); + } + let states = reader.read_states()?; + let ascii_rows = reader.read_ascii_rows()?; + let wide_rows = reader.read_wide_rows()?; + let accepts = reader.read_accepts()?; + if reader.position != data.len() { + return None; + } + let dfa = Self { + mode_starts, + states, + ascii_rows, + wide_rows, + accepts, + }; + dfa.table_indexes_are_valid().then_some(dfa) + } + + /// Cheap structural validation so a corrupted embedded stream degrades to + /// runtime compilation instead of an out-of-bounds panic mid-parse. + fn table_indexes_are_valid(&self) -> bool { + let state_ok = |target: u16| { + usize::from(target) < self.states.len() || target >= ESCAPE_STATE + }; + self.mode_starts + .iter() + .flatten() + .all(|&start| usize::from(start) < self.states.len()) + && self.states.iter().all(|state| { + (state.ascii_row as usize) < self.ascii_rows.len() + && (state.wide_row as usize) < self.wide_rows.len() + && state_ok(state.eof_target) + && (state.accept == u32::MAX || (state.accept as usize) < self.accepts.len()) + }) + && self.ascii_rows.iter().all(|row| row.iter().all(|&target| state_ok(target))) + && self.wide_rows.iter().all(|row| { + wide_row_is_searchable(row) && row.iter().all(|range| state_ok(range.target)) + }) + } +} + +/// Wide rows must hold well-formed, sorted, disjoint ranges for +/// [`CompiledLexerDfa::char_target`]'s binary search; anything else would +/// silently misroute transitions instead of degrading to recompilation. +fn wide_row_is_searchable(row: &[WideRange]) -> bool { + row.iter().all(|range| range.low <= range.high) + && row.windows(2).all(|pair| pair[0].high < pair[1].low) +} + +/// Version tag guarding embedded tables against serialization format drift. +const SERIALIZED_TAG: u32 = 0x4C58_4401; + +/// Cursor over a serialized DFA stream. +struct SerializedReader<'a> { + data: &'a [u32], + position: usize, +} + +impl SerializedReader<'_> { + fn next(&mut self) -> Option { + let value = self.data.get(self.position).copied(); + self.position += 1; + value + } + + fn next_u16(&mut self) -> Option { + u16::try_from(self.next()?).ok() + } + + fn next_len(&mut self) -> Option { + usize::try_from(self.next()?).ok() + } + + fn read_states(&mut self) -> Option> { + let count = self.next_len()?; + let mut states = Vec::with_capacity(count.min(self.data.len())); + for _ in 0..count { + states.push(CompiledLexerState { + ascii_row: self.next()?, + wide_row: self.next()?, + eof_target: self.next_u16()?, + accept: self.next()?, + }); + } + Some(states) + } + + fn read_ascii_rows(&mut self) -> Option> { + let count = self.next_len()?; + let mut rows = Vec::with_capacity(count.min(self.data.len())); + for _ in 0..count { + let mut row = [DEAD_STATE; ASCII_EDGE_SYMBOLS]; + for pair in 0..ASCII_EDGE_SYMBOLS / 2 { + let word = self.next()?; + row[pair * 2] = (word & 0xFFFF) as u16; + row[pair * 2 + 1] = (word >> 16) as u16; + } + rows.push(row); + } + Some(rows) + } + + fn read_wide_rows(&mut self) -> Option>> { + let count = self.next_len()?; + let mut rows = Vec::with_capacity(count.min(self.data.len())); + for _ in 0..count { + let len = self.next_len()?; + let mut row = Vec::with_capacity(len.min(self.data.len())); + for _ in 0..len { + row.push(WideRange { + low: self.next()?, + high: self.next()?, + target: self.next_u16()?, + }); + } + rows.push(row.into()); + } + Some(rows) + } + + fn read_accepts(&mut self) -> Option> { + let count = self.next_len()?; + let mut accepts = Vec::with_capacity(count.min(self.data.len())); + for _ in 0..count { + let rule_index = self.next_len()?; + let consumed_eof = self.next()? != 0; + let action_count = self.next_len()?; + let mut actions = Vec::with_capacity(action_count.min(self.data.len())); + for _ in 0..action_count { + actions.push(CompiledLexerActionTrace { + action_index: self.next_len()?, + rule_index: self.next_len()?, + behind: self.next_len()?, + }); + } + accepts.push(CompiledLexerAccept { + rule_index, + consumed_eof, + actions, + }); + } + Some(accepts) + } +} + +/// Deduplicating pools for edge rows shared by many DFA states. +#[derive(Debug, Default)] +struct RowPools { + ascii_ids: FxHashMap<[u16; ASCII_EDGE_SYMBOLS], u32>, + wide_ids: FxHashMap, u32>, +} + +impl RowPools { + fn intern_ascii(&mut self, rows: &mut Vec<[u16; ASCII_EDGE_SYMBOLS]>, row: [u16; ASCII_EDGE_SYMBOLS]) -> u32 { + *self.ascii_ids.entry(row).or_insert_with(|| { + rows.push(row); + (rows.len() - 1) as u32 + }) + } + + fn intern_wide(&mut self, rows: &mut Vec>, row: Vec) -> u32 { + let row: Box<[WideRange]> = row.into(); + if let Some(&id) = self.wide_ids.get(&row) { + return id; + } + rows.push(row.clone()); + let id = (rows.len() - 1) as u32; + self.wide_ids.insert(row, id); + id + } +} + +/// In-progress subset construction for one lexer mode. +/// +/// States are numbered globally (`base` + discovery order) so edges can be +/// written directly into the final table, but nothing is committed to the +/// shared [`CompiledLexerDfa`] until the whole mode succeeds. +struct ModeBuild { + base: usize, + ids: FxHashMap, + configs: Vec>, + steps: Vec, + accepts: Vec>, +} + +/// Edge rows produced by expanding one DFA state. +struct StateRows { + /// Sorted, disjoint code-point segments with live targets. + segments: Vec<(i32, i32, u16)>, + eof_target: u16, +} + +impl ModeBuild { + fn new(base: usize) -> Self { + Self { + base, + ids: FxHashMap::default(), + configs: Vec::new(), + steps: Vec::new(), + accepts: Vec::new(), + } + } + + const fn len(&self) -> usize { + self.configs.len() + } + + /// Returns the state id for a closed, pruned config set, creating the + /// state when the (input-offset-normalized) identity is new. + /// [`ESCAPE_STATE`] means the state budget is exhausted and the edge must + /// hand the token to the interpreter. + fn intern(&mut self, atn: &Atn, configs: Vec, step: usize) -> u16 { + let key = LexerDfaKey::new( + configs + .iter() + .map(|config| relative_config_key(config, step)) + .collect(), + ); + if let Some(&id) = self.ids.get(&key) { + return id; + } + let local = self.configs.len(); + let global = self.base + local; + if local >= MAX_MODE_STATES || global >= usize::from(ESCAPE_STATE) { + return ESCAPE_STATE; + } + let Ok(id) = u16::try_from(global) else { + return ESCAPE_STATE; + }; + self.ids.insert(key, id); + self.accepts.push(compiled_accept(atn, &configs, step)); + self.configs.push(configs); + self.steps.push(step); + id + } +} + +/// Normalizes one config for DFA-state identity, measuring action positions +/// backwards from the current input offset (`step`). +/// +/// This differs from the interpreter cache's token-start-relative deltas on +/// purpose: rule-final lexer commands (`skip`, `pushMode`, …) fire a fixed +/// distance before the accept, so anchoring at the read position keeps the +/// state space finite regardless of token length. +fn relative_config_key(config: &LexerConfig, step: usize) -> LexerDfaConfigKey { + LexerDfaConfigKey::new( + config.state, + config.alt_rule_index, + config.consumed_eof, + config.passed_non_greedy, + config.stack.clone(), + config + .actions + .iter() + .map(|action| LexerDfaActionKey { + action_index: action.action_index, + position_delta: step.saturating_sub(action.position), + rule_index: action.rule_index, + }) + .collect(), + ) +} + +/// Computes the accept metadata for a DFA state from its config set, using +/// the interpreter's own rule-priority selection. +fn compiled_accept(atn: &Atn, configs: &[LexerConfig], step: usize) -> Option { + let accept = best_accept(atn, configs)?; + debug_assert!( + accept.position == step, + "every config in a lexer DFA state shares the state's input offset" + ); + Some(CompiledLexerAccept { + rule_index: accept.rule_index, + consumed_eof: accept.consumed_eof, + actions: accept + .actions + .iter() + .map(|trace| CompiledLexerActionTrace { + action_index: trace.action_index, + rule_index: trace.rule_index, + behind: accept.position.saturating_sub(trace.position), + }) + .collect(), + }) +} + +/// Runs subset construction for one mode; `None` leaves the whole mode to the +/// interpreter (only when its very first closure already escapes). +fn build_mode( + atn: &Atn, + mode: usize, + dfa: &mut CompiledLexerDfa, + pools: &mut RowPools, +) -> Option { + let start_state = atn.mode_to_start_state().get(mode).copied()?; + let mut build = ModeBuild::new(dfa.states.len()); + let start_configs = closed_configs( + atn, + vec![LexerConfig { + state: start_state, + position: 0, + consumed_eof: false, + alt_rule_index: None, + passed_non_greedy: false, + stack: Vec::new(), + actions: Vec::new(), + }], + )?; + let start_id = build.intern(atn, start_configs, 0); + if start_id == ESCAPE_STATE { + return None; + } + + let mut rows = Vec::new(); + let mut cursor = 0; + while cursor < build.len() { + rows.push(expand_state(atn, &mut build, cursor)); + cursor += 1; + } + + commit_mode(dfa, pools, build, rows); + Some(start_id) +} + +/// Closes and prunes a moved config set exactly like the interpreter does. +/// `None` means the closure crossed a semantic predicate (which only the +/// interpreter can evaluate) or entered a recursive lexer rule (nested +/// comments never determinize), so the edge must escape. +fn closed_configs(atn: &Atn, moved: Vec) -> Option> { + let closure = epsilon_closure(atn, moved, &mut |_| true); + if closure.has_semantic_context { + return None; + } + if closure.configs.iter().any(has_recursive_stack) { + return None; + } + let mut configs = closure.configs; + for config in &mut configs { + prune_dead_action_traces(atn, config); + if config.actions.len() > MAX_ACTION_TRACES { + return None; + } + } + Some(prune_after_accepts(atn, configs)) +} + +/// Drops action traces the accept-time dispatcher would suppress anyway. +/// +/// The interpreter keeps traces of every action transition it crosses and +/// filters them per accept with `lexer_action_belongs_to_accept`; a token +/// rule referenced from another rule leaves traces that can never fire (its +/// commands belong to itself, not the enclosing rule). Carrying them into +/// DFA-state identity would mint a fresh state per input offset — rules that +/// loop over comment/whitespace references would never determinize. +fn prune_dead_action_traces(atn: &Atn, config: &mut LexerConfig) { + let Some(accept_rule) = config.alt_rule_index else { + return; + }; + config + .actions + .retain(|trace| lexer_action_belongs_to_accept(atn, accept_rule, trace.rule_index)); +} + +/// Detects lexer-rule recursion: re-entering a rule from the same call site +/// pushes the same follow state again, so a duplicated stack entry (or an +/// implausibly deep stack) marks a config a finite DFA cannot represent. +fn has_recursive_stack(config: &LexerConfig) -> bool { + let stack = &config.stack; + if stack.len() > MAX_STACK_DEPTH { + return true; + } + stack + .iter() + .enumerate() + .any(|(index, follow)| stack[..index].contains(follow)) +} + +/// Computes every outgoing edge of one interned DFA state. +fn expand_state(atn: &Atn, build: &mut ModeBuild, local: usize) -> StateRows { + let configs = build.configs[local].clone(); + let step = build.steps[local]; + let entries = consuming_entries(atn, &configs); + let eof_target = eof_move(atn, build, &configs, step, &entries); + + let entry_intervals: Vec> = entries + .iter() + .map(|(_, transition)| transition_char_intervals(transition)) + .collect(); + let segments = char_segments(&entry_intervals); + let matrix = segment_mask_matrix(&segments, &entry_intervals, entries.len()); + let words = entries.len().div_ceil(64); + + let mut rows = StateRows { + segments: Vec::new(), + eof_target, + }; + // Distinct transition sets are few even when segments are many (large + // Unicode classes fragment the alphabet), so closures run once per + // matching-transition mask, not once per segment. + let mut mask_targets: FxHashMap, u16> = FxHashMap::default(); + for (index, &(low, high)) in segments.iter().enumerate() { + let mask = &matrix[index * words..(index + 1) * words]; + if mask.iter().all(|&word| word == 0) { + continue; + } + let target = match mask_targets.get(mask) { + Some(&target) => target, + None => { + let target = move_target(atn, build, &configs, step, &entries, mask); + mask_targets.insert(mask.to_vec(), target); + target + } + }; + if target != DEAD_STATE { + rows.segments.push((low, high, target)); + } + } + rows +} + +/// Lists each config's consuming transitions in the interpreter's move order. +fn consuming_entries<'a>(atn: &'a Atn, configs: &[LexerConfig]) -> Vec<(usize, &'a Transition)> { + let mut entries = Vec::new(); + for (config_index, config) in configs.iter().enumerate() { + let Some(state) = atn.state(config.state) else { + continue; + }; + for transition in &state.transitions { + if !transition.is_epsilon() { + entries.push((config_index, transition)); + } + } + } + entries +} + +/// Splits the code-point alphabet at every interval boundary, so each segment +/// is matched uniformly by every transition. +fn char_segments(entry_intervals: &[Vec<(i32, i32)>]) -> Vec<(i32, i32)> { + let mut cuts = Vec::new(); + for intervals in entry_intervals { + for &(low, high) in intervals { + cuts.push(low); + cuts.push(high + 1); + } + } + cuts.sort_unstable(); + cuts.dedup(); + cuts.windows(2).map(|pair| (pair[0], pair[1] - 1)).collect() +} + +/// Marks, for every segment, which entries match it — one bit row per +/// segment. Sweeping each entry's intervals over the sorted segment starts +/// keeps the work proportional to interval count, not `segments × entries`. +fn segment_mask_matrix( + segments: &[(i32, i32)], + entry_intervals: &[Vec<(i32, i32)>], + entry_count: usize, +) -> Vec { + let words = entry_count.div_ceil(64); + let mut matrix = vec![0_u64; segments.len() * words]; + for (bit, intervals) in entry_intervals.iter().enumerate() { + for &(low, high) in intervals { + // Interval boundaries are cut points, so the covered segments are + // exactly those whose start lies inside the interval. + let from = segments.partition_point(|&(start, _)| start < low); + let to = segments.partition_point(|&(start, _)| start <= high); + for segment in from..to { + matrix[segment * words + bit / 64] |= 1 << (bit % 64); + } + } + } + matrix +} + +/// Materializes the code-point intervals a transition consumes, clamped to +/// the valid character range (EOF is handled separately). +fn transition_char_intervals(transition: &Transition) -> Vec<(i32, i32)> { + let mut intervals = Vec::new(); + let mut push_clamped = |low: i32, high: i32| { + let low = low.max(MIN_CHAR_VALUE); + let high = high.min(MAX_CHAR_VALUE); + if low <= high { + intervals.push((low, high)); + } + }; + match transition { + Transition::Atom { label, .. } => push_clamped(*label, *label), + Transition::Range { start, stop, .. } => push_clamped(*start, *stop), + Transition::Set { set, .. } => { + for &(low, high) in set.ranges() { + push_clamped(low, high); + } + } + Transition::NotSet { set, .. } => { + // `NotSet` matches the complement within the character range; + // `IntervalSet` ranges are sorted and coalesced. + let mut next = MIN_CHAR_VALUE; + for &(low, high) in set.ranges() { + if low > next { + push_clamped(next, low - 1); + } + next = next.max(high.saturating_add(1)); + } + push_clamped(next, MAX_CHAR_VALUE); + } + Transition::Wildcard { .. } => push_clamped(MIN_CHAR_VALUE, MAX_CHAR_VALUE), + _ => {} + } + intervals +} + +/// Advances the masked entries by one character and interns the result; +/// closures that escape compile as [`ESCAPE_STATE`] edges. +fn move_target( + atn: &Atn, + build: &mut ModeBuild, + configs: &[LexerConfig], + step: usize, + entries: &[(usize, &Transition)], + mask: &[u64], +) -> u16 { + let mut moved = Vec::new(); + for (bit, (config_index, transition)) in entries.iter().enumerate() { + if mask[bit / 64] & (1 << (bit % 64)) == 0 { + continue; + } + let mut advanced = configs[*config_index].clone(); + set_config_state(atn, &mut advanced, transition.target()); + advanced.position += 1; + moved.push(advanced); + } + let Some(active) = closed_configs(atn, moved) else { + return ESCAPE_STATE; + }; + if active.is_empty() { + return DEAD_STATE; + } + build.intern(atn, active, step + 1) +} + +/// Advances the EOF-matching entries; EOF consumes no character, so the input +/// offset stays put and the moved configs record `consumed_eof`. +fn eof_move( + atn: &Atn, + build: &mut ModeBuild, + configs: &[LexerConfig], + step: usize, + entries: &[(usize, &Transition)], +) -> u16 { + let mut moved = Vec::new(); + for (config_index, transition) in entries { + if !transition.matches(EOF, MIN_CHAR_VALUE, MAX_CHAR_VALUE) { + continue; + } + let mut advanced = configs[*config_index].clone(); + set_config_state(atn, &mut advanced, transition.target()); + advanced.consumed_eof = true; + moved.push(advanced); + } + if moved.is_empty() { + return DEAD_STATE; + } + let Some(active) = closed_configs(atn, moved) else { + return ESCAPE_STATE; + }; + if active.is_empty() { + return DEAD_STATE; + } + build.intern(atn, active, step) +} + +/// Converts a finished mode's edge rows into pooled table entries. +fn commit_mode(dfa: &mut CompiledLexerDfa, pools: &mut RowPools, build: ModeBuild, rows: Vec) { + for (accept, state_rows) in build.accepts.into_iter().zip(rows) { + let accept_id = accept.map_or(u32::MAX, |accept| { + dfa.accepts.push(accept); + (dfa.accepts.len() - 1) as u32 + }); + let (ascii_row, wide_row) = split_rows(&state_rows.segments); + dfa.states.push(CompiledLexerState { + ascii_row: pools.intern_ascii(&mut dfa.ascii_rows, ascii_row), + wide_row: pools.intern_wide(&mut dfa.wide_rows, wide_row), + eof_target: state_rows.eof_target, + accept: accept_id, + }); + } +} + +/// Splits sorted segments into the dense ASCII row and merged wide ranges. +fn split_rows(segments: &[(i32, i32, u16)]) -> ([u16; ASCII_EDGE_SYMBOLS], Vec) { + let mut ascii = [DEAD_STATE; ASCII_EDGE_SYMBOLS]; + let mut wide: Vec = Vec::new(); + for &(low, high, target) in segments { + let ascii_high = high.min(ASCII_EDGE_LIMIT - 1); + for code_point in low..=ascii_high { + ascii[code_point.cast_unsigned() as usize] = target; + } + if high >= ASCII_EDGE_LIMIT { + let low = low.max(ASCII_EDGE_LIMIT).cast_unsigned(); + let high = high.cast_unsigned(); + if let Some(last) = wide.last_mut() + && last.target == target + && last.high + 1 == low + { + last.high = high; + continue; + } + wide.push(WideRange { low, high, target }); + } + } + (ascii, wide) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::atn::lexer::{next_token, next_token_compiled, next_token_compiled_with_hooks}; + use crate::atn::serialized::{AtnDeserializer, SerializedAtn}; + use crate::char_stream::InputStream; + use crate::lexer::BaseLexer; + use crate::recognizer::RecognizerData; + use crate::token::{TOKEN_EOF, Token}; + use crate::vocabulary::Vocabulary; + + fn recognizer_data() -> RecognizerData { + RecognizerData::new( + "T", + Vocabulary::new( + [None, Some("'ab'"), Some("' '")], + [None, Some("AB"), Some("WS")], + [None::<&str>, None, None], + ), + ) + } + + /// Two-rule lexer (`AB: 'ab';` and `WS: ' ' -> skip;`), with rule 0's + /// final epsilon optionally replaced by a semantic predicate transition. + fn two_rule_atn(with_predicate: bool) -> Atn { + let epsilon_or_predicate = if with_predicate { 4 } else { 1 }; + AtnDeserializer::new(&SerializedAtn::from_i32(&[ + 4, 0, 2, // version, lexer, max token type + 9, // states + 6, -1, // 0 token start + 2, 0, // 1 rule 0 start + 1, 0, // 2 + 1, 0, // 3 + 7, 0, // 4 rule 0 stop + 2, 1, // 5 rule 1 start + 1, 1, // 6 + 1, 1, // 7 + 7, 1, // 8 rule 1 stop + 0, // non-greedy + 0, // precedence + 2, // rules + 1, 1, // rule 0 starts at 1, token type 1 + 5, 2, // rule 1 starts at 5, token type 2 + 1, // modes + 0, // default mode starts at 0 + 0, // sets + 8, // edges + 0, 1, 1, 0, 0, 0, // start -> rule 0 + 0, 5, 1, 0, 0, 0, // start -> rule 1 + 1, 2, 5, 'a' as i32, 0, 0, // 'a' + 2, 3, 5, 'b' as i32, 0, 0, // 'b' + 3, 4, epsilon_or_predicate, 0, 0, 0, // epsilon or predicate to stop + 5, 6, 5, ' ' as i32, 0, 0, // ' ' + 6, 7, 1, 0, 0, 0, // + 7, 8, 6, 1, 0, 0, // action 0, then stop + 1, // decisions + 0, 1, // lexer actions + 6, 0, 0, // skip + ])) + .deserialize() + .expect("artificial lexer ATN should deserialize") + } + + /// One token rule matching `[\u{100}-\u{200}]+`, exercising wide rows. + fn wide_range_atn() -> Atn { + AtnDeserializer::new(&SerializedAtn::from_i32(&[ + 4, 0, 1, // version, lexer, max token type + 5, // states + 6, -1, // 0 token start + 2, 0, // 1 rule 0 start + 1, 0, // 2 + 1, 0, // 3 + 7, 0, // 4 rule 0 stop + 0, // non-greedy + 0, // precedence + 1, // rules + 1, 1, // rule 0 starts at 1, token type 1 + 1, // modes + 0, // default mode starts at 0 + 0, // sets + 5, // edges + 0, 1, 1, 0, 0, 0, // start -> rule 0 + 1, 2, 1, 0, 0, 0, // + 2, 3, 2, 0x100, 0x200, 0, // range + 3, 2, 1, 0, 0, 0, // greedy loop continues first + 3, 4, 1, 0, 0, 0, // then exits to stop + 0, // decisions + 0, // lexer actions + ])) + .deserialize() + .expect("artificial wide-range lexer ATN should deserialize") + } + + #[test] + fn compiled_dfa_matches_longest_token_and_skips() { + let atn = two_rule_atn(false); + let dfa = CompiledLexerDfa::compile(&atn); + assert!(dfa.has_compiled_modes()); + assert!(dfa.mode_start(0).is_some()); + + let mut lexer = BaseLexer::new(InputStream::new(" ab"), recognizer_data()); + let token = next_token_compiled(&mut lexer, &atn, &dfa); + assert_eq!(token.token_type(), 1); + assert_eq!(token.text(), Some("ab")); + assert_eq!( + next_token_compiled(&mut lexer, &atn, &dfa).token_type(), + TOKEN_EOF + ); + } + + #[test] + fn predicate_edge_escapes_to_the_interpreter() { + let atn = two_rule_atn(true); + let dfa = CompiledLexerDfa::compile(&atn); + // The predicate sits mid-rule, so the mode still compiles; only the + // edge that would cross it escapes to the interpreter. + assert!(dfa.mode_start(0).is_some()); + + let mut lexer = BaseLexer::new(InputStream::new(" ab"), recognizer_data()); + let token = next_token_compiled_with_hooks( + &mut lexer, + &atn, + &dfa, + |_, _| {}, + |_, _| true, + |_, _, _| {}, + ); + assert_eq!(token.token_type(), 1); + assert_eq!(token.text(), Some("ab")); + } + + #[test] + fn compiled_dfa_walks_wide_ranges() { + let atn = wide_range_atn(); + let dfa = CompiledLexerDfa::compile(&atn); + assert!(dfa.mode_start(0).is_some()); + + let mut lexer = BaseLexer::new(InputStream::new("ĀĂ"), recognizer_data()); + let token = next_token_compiled(&mut lexer, &atn, &dfa); + assert_eq!(token.token_type(), 1); + assert_eq!(token.text(), Some("ĀĂ")); + assert_eq!( + next_token_compiled(&mut lexer, &atn, &dfa).token_type(), + TOKEN_EOF + ); + } + + #[test] + fn compiled_dfa_reports_recognition_errors_like_the_interpreter() { + let atn = wide_range_atn(); + let dfa = CompiledLexerDfa::compile(&atn); + + let mut compiled = BaseLexer::new(InputStream::new("zĀ"), recognizer_data()); + let mut interpreted = BaseLexer::new(InputStream::new("zĀ"), recognizer_data()); + loop { + let compiled_token = next_token_compiled(&mut compiled, &atn, &dfa); + let interpreted_token = next_token(&mut interpreted, &atn); + assert_eq!(compiled_token.token_type(), interpreted_token.token_type()); + assert_eq!(compiled_token.text(), interpreted_token.text()); + if compiled_token.token_type() == TOKEN_EOF { + break; + } + } + let compiled_errors: Vec = compiled + .drain_errors() + .into_iter() + .map(|error| error.message) + .collect(); + let interpreted_errors: Vec = interpreted + .drain_errors() + .into_iter() + .map(|error| error.message) + .collect(); + assert_eq!(compiled_errors, vec!["token recognition error at: 'z'"]); + assert_eq!(compiled_errors, interpreted_errors); + } + + #[test] + fn serialization_round_trips() { + let atn = two_rule_atn(false); + let dfa = CompiledLexerDfa::compile(&atn); + let stream = dfa.serialize(); + + let restored = + CompiledLexerDfa::from_serialized(&stream).expect("stream should deserialize"); + assert_eq!(restored.serialize(), stream); + + let mut lexer = BaseLexer::new(InputStream::new(" ab"), recognizer_data()); + let token = next_token_compiled(&mut lexer, &atn, &restored); + assert_eq!(token.token_type(), 1); + assert_eq!(token.text(), Some("ab")); + + // A stream from a different runtime version is rejected, not trusted. + let mut wrong_tag = stream; + wrong_tag[0] ^= 1; + assert!(CompiledLexerDfa::from_serialized(&wrong_tag).is_none()); + } + + #[test] + fn malformed_wide_rows_are_rejected() { + let atn = wide_range_atn(); + let stream = CompiledLexerDfa::compile(&atn).serialize(); + + // Invert the [0x100, 0x200] range's bounds in place; a broken wide + // row must fail validation, not silently misroute binary searches. + let position = stream + .windows(2) + .position(|pair| pair == [0x100, 0x200]) + .expect("wide-range test grammar serializes its range bounds"); + let mut inverted = stream; + inverted.swap(position, position + 1); + assert!(CompiledLexerDfa::from_serialized(&inverted).is_none()); + } + + #[test] + fn force_interpreted_bypasses_compiled_tables() { + let atn = two_rule_atn(false); + let dfa = CompiledLexerDfa::compile(&atn); + + let mut lexer = BaseLexer::new(InputStream::new("ab"), recognizer_data()); + lexer.set_force_interpreted(true); + let token = next_token_compiled(&mut lexer, &atn, &dfa); + assert_eq!(token.token_type(), 1); + // The interpreter path records the learned-DFA trace; the compiled + // walk does not. + assert!(!lexer.lexer_dfa_string().is_empty()); + } +} diff --git a/src/atn/mod.rs b/src/atn/mod.rs index 330646e7..adb3fea2 100644 --- a/src/atn/mod.rs +++ b/src/atn/mod.rs @@ -6,6 +6,7 @@ //! these compact Rust structures for simulation. pub mod lexer; +pub mod lexer_dfa; pub mod parser; pub mod serialized; diff --git a/src/bin/antlr4-runtime-testsuite.rs b/src/bin/antlr4-runtime-testsuite.rs index 4f329712..863129a4 100644 --- a/src/bin/antlr4-runtime-testsuite.rs +++ b/src/bin/antlr4-runtime-testsuite.rs @@ -1501,18 +1501,25 @@ fn smoke_main(descriptor: &Descriptor) -> String { } let module_name = module_name(&descriptor.grammar_name); let type_name = rust_type_name(&descriptor.grammar_name); - let dfa_dump = if descriptor.flags.trim() == "showDFA" { + let show_dfa = descriptor.flags.trim() == "showDFA"; + let dfa_dump = if show_dfa { " print!(\"{}\", tokens.token_source().lexer_dfa_string());\n" } else { "" }; - let token_source_import = if descriptor.flags.trim() == "showDFA" { - ", TokenSource" + let token_source_import = if show_dfa { ", TokenSource" } else { "" }; + // The learned-DFA trace only exists when tokens go through ATN + // interpretation, so showDFA cases opt out of the compiled lexer DFA. + let (lexer_binding, force_interpreted) = if show_dfa { + ( + "let mut lexer", + " lexer.set_force_interpreted(true);\n", + ) } else { - "" + ("let lexer", "") }; format!( - "pub mod generated {{\n pub mod {module_name};\n}}\n\nuse antlr4_runtime::{{CommonTokenStream, InputStream{token_source_import}}};\nuse generated::{module_name}::{type_name};\n\nfn main() {{\n let lexer = {type_name}::new(InputStream::new(\"{}\"));\n let mut tokens = CommonTokenStream::new(lexer);\n tokens.fill();\n for error in tokens.drain_source_errors() {{\n eprintln!(\"line {{}}:{{}} {{}}\", error.line, error.column, error.message);\n }}\n for token in tokens.tokens() {{\n println!(\"{{token}}\");\n }}\n{dfa_dump}}}\n", + "pub mod generated {{\n pub mod {module_name};\n}}\n\nuse antlr4_runtime::{{CommonTokenStream, InputStream{token_source_import}}};\nuse generated::{module_name}::{type_name};\n\nfn main() {{\n {lexer_binding} = {type_name}::new(InputStream::new(\"{}\"));\n{force_interpreted} let mut tokens = CommonTokenStream::new(lexer);\n tokens.fill();\n for error in tokens.drain_source_errors() {{\n eprintln!(\"line {{}}:{{}} {{}}\", error.line, error.column, error.message);\n }}\n for token in tokens.tokens() {{\n println!(\"{{token}}\");\n }}\n{dfa_dump}}}\n", rust_string(&descriptor.input) ) } diff --git a/src/bin/antlr4-rust-gen.rs b/src/bin/antlr4-rust-gen.rs index d693979d..41a408e9 100644 --- a/src/bin/antlr4-rust-gen.rs +++ b/src/bin/antlr4-rust-gen.rs @@ -6,6 +6,7 @@ use std::io; use std::ops::AddAssign; use std::path::{Path, PathBuf}; +use antlr4_runtime::atn::lexer_dfa::CompiledLexerDfa; use antlr4_runtime::atn::serialized::{AtnDeserializer, SerializedAtn}; use antlr4_runtime::atn::{Atn, AtnStateKind, LexerAction, Transition}; @@ -304,6 +305,7 @@ fn render_lexer( |source| lexer_predicate_templates(data, source), )?; let adjusts_accept_position = grammar_source.is_some_and(uses_position_adjusting_lexer); + let lexer_dfa_data = compiled_lexer_dfa_words(data); let has_action_dispatch = lexer_actions_need_dispatch(&actions); let action_method = render_lexer_action_method(&actions); let predicate_method = render_lexer_predicate_method(&predicates); @@ -312,42 +314,29 @@ fn render_lexer( } else { String::new() }; - let next_token_call = match ( - !has_action_dispatch, - predicates.is_empty(), - adjusts_accept_position, - ) { - (true, true, false) => { - "antlr4_runtime::atn::lexer::next_token(&mut self.base, atn())".to_owned() - } - (false, true, false) => { - "antlr4_runtime::atn::lexer::next_token_with_actions(&mut self.base, atn(), Self::run_action)" - .to_owned() - } - (true, false, false) => { - "antlr4_runtime::atn::lexer::next_token_with_actions_and_predicates(&mut self.base, atn(), |_, _| {}, Self::run_predicate)" - .to_owned() - } - (false, false, false) => { - "antlr4_runtime::atn::lexer::next_token_with_actions_and_predicates(&mut self.base, atn(), Self::run_action, Self::run_predicate)" - .to_owned() - } - (true, true, true) => { - "antlr4_runtime::atn::lexer::next_token_with_accept_adjuster(&mut self.base, atn(), Self::adjust_accept_position)" - .to_owned() - } - (false, true, true) => { - "antlr4_runtime::atn::lexer::next_token_with_hooks(&mut self.base, atn(), Self::run_action, |_, _| true, Self::adjust_accept_position)" - .to_owned() - } - (true, false, true) => { - "antlr4_runtime::atn::lexer::next_token_with_hooks(&mut self.base, atn(), |_, _| {}, Self::run_predicate, Self::adjust_accept_position)" - .to_owned() - } - (false, false, true) => { - "antlr4_runtime::atn::lexer::next_token_with_hooks(&mut self.base, atn(), Self::run_action, Self::run_predicate, Self::adjust_accept_position)" - .to_owned() - } + let next_token_call = if !has_action_dispatch && predicates.is_empty() && !adjusts_accept_position + { + "antlr4_runtime::atn::lexer::next_token_compiled(&mut self.base, atn(), lexer_dfa())" + .to_owned() + } else { + let action = if has_action_dispatch { + "Self::run_action" + } else { + "|_, _| {}" + }; + let predicate = if predicates.is_empty() { + "|_, _| true" + } else { + "Self::run_predicate" + }; + let adjuster = if adjusts_accept_position { + "Self::adjust_accept_position" + } else { + "|_, _, _| {}" + }; + format!( + "antlr4_runtime::atn::lexer::next_token_compiled_with_hooks(&mut self.base, atn(), lexer_dfa(), {action}, {predicate}, {adjuster})" + ) }; let generated_header = GENERATED_MODULE_HEADER; let generated_footer = GENERATED_MODULE_FOOTER; @@ -357,6 +346,7 @@ fn render_lexer( use antlr4_runtime::recognizer::RecognizerData; use antlr4_runtime::token::{{CommonToken, TokenSource}}; use antlr4_runtime::atn::Atn; +use antlr4_runtime::atn::lexer_dfa::CompiledLexerDfa; use antlr4_runtime::atn::serialized::AtnDeserializer; use antlr4_runtime::{{BaseLexer, GeneratedLexer, GrammarMetadata, Lexer, Recognizer}}; use std::sync::OnceLock; @@ -376,6 +366,20 @@ fn atn() -> &'static Atn {{ }}) }} +static LEXER_DFA_DATA: &[u32] = &[{lexer_dfa_data}]; + +static LEXER_DFA_CELL: OnceLock = OnceLock::new(); + +/// Ahead-of-time lexer DFA tables compiled by antlr4-rust-gen, embedded so +/// runtime startup only deserializes them. Rebuilt from the ATN instead when +/// the embedded stream comes from a different runtime version. +fn lexer_dfa() -> &'static CompiledLexerDfa {{ + LEXER_DFA_CELL.get_or_init(|| {{ + CompiledLexerDfa::from_serialized(LEXER_DFA_DATA) + .unwrap_or_else(|| CompiledLexerDfa::compile(atn())) + }}) +}} + #[derive(Clone, Debug)] pub struct {type_name} where @@ -404,6 +408,13 @@ where metadata() }} + /// Routes every token through ATN interpretation instead of the compiled + /// lexer DFA, so the learned-DFA trace (`lexer_dfa_string`) observes each + /// match. + pub fn set_force_interpreted(&mut self, force_interpreted: bool) {{ + self.base.set_force_interpreted(force_interpreted); + }} + {action_method} {predicate_method} {accept_adjust_method} @@ -463,6 +474,23 @@ where )) } +/// Compiles the lexer DFA at generation time and flattens it for embedding. +/// +/// An empty stream makes the generated lexer fall back to compiling the DFA +/// from its ATN at first use, so generation never fails on this step. +fn compiled_lexer_dfa_words(data: &InterpData) -> String { + if data.atn.is_empty() { + return String::new(); + } + let serialized = SerializedAtn::from_i32(&data.atn); + let Ok(atn) = AtnDeserializer::new(&serialized).deserialize() else { + return String::new(); + }; + let words = CompiledLexerDfa::compile(&atn).serialize(); + let rendered: Vec = words.iter().map(u32::to_string).collect(); + rendered.join(",") +} + #[derive(Clone, Debug, Eq, PartialEq)] struct GeneratedParserRule { rule_index: usize, @@ -4874,88 +4902,143 @@ fn next_action_block(source: &str, offset: usize) -> Option Option { - let mut index = offset; - let mut single_quoted = false; - let mut double_quoted = false; - let mut escaped = false; - let mut line_comment = false; - let mut block_comment = false; - let mut char_set = false; - while let Some(ch) = source[index..].chars().next() { - let size = ch.len_utf8(); - if line_comment { - line_comment = ch != '\n'; - index += size; - continue; + let mut cursor = GrammarSourceCursor::new(source, offset); + while let Some((index, ch)) = cursor.next_significant() { + if ch == '{' { + return Some(index); + } + } + None +} + +/// Lexical cursor over ANTLR grammar source that skips line and block +/// comments, string literals, and `[...]` character sets, yielding only +/// characters that are significant to grammar structure. +/// +/// Action extraction and rule-header scanning need the same skip rules; +/// sharing one state machine keeps them from drifting apart. +struct GrammarSourceCursor<'a> { + source: &'a str, + index: usize, + single_quoted: bool, + double_quoted: bool, + escaped: bool, + line_comment: bool, + block_comment: bool, + char_set: bool, +} + +impl<'a> GrammarSourceCursor<'a> { + const fn new(source: &'a str, offset: usize) -> Self { + Self { + source, + index: offset, + single_quoted: false, + double_quoted: false, + escaped: false, + line_comment: false, + block_comment: false, + char_set: false, + } + } + + /// Moves the cursor to `index`, which must be a char boundary outside any + /// comment, string literal, or character set. + const fn seek(&mut self, index: usize) { + self.index = index; + } + + /// Returns the next structurally significant character with its byte + /// offset, consuming it. + fn next_significant(&mut self) -> Option<(usize, char)> { + while let Some(ch) = self.source[self.index..].chars().next() { + let index = self.index; + let size = ch.len_utf8(); + if self.consume_skipped(ch, size) { + continue; + } + match ch { + '/' if self.source.as_bytes().get(index..index + 2) == Some(b"//") => { + self.line_comment = true; + self.index += 2; + } + '/' if self.source.as_bytes().get(index..index + 2) == Some(b"/*") => { + self.block_comment = true; + self.index += 2; + } + '\'' => { + self.single_quoted = true; + self.index += size; + } + '"' => { + self.double_quoted = true; + self.index += size; + } + '[' => { + self.char_set = true; + self.index += size; + } + _ => { + self.index += size; + return Some((index, ch)); + } + } } - if block_comment { - if source.as_bytes().get(index..index + 2) == Some(b"*/") { - block_comment = false; - index += 2; + None + } + + /// Consumes one character belonging to an active comment, string, or + /// character-set region; false when the cursor is at top level. + fn consume_skipped(&mut self, ch: char, size: usize) -> bool { + if self.line_comment { + self.line_comment = ch != '\n'; + self.index += size; + return true; + } + if self.block_comment { + if self.source.as_bytes().get(self.index..self.index + 2) == Some(b"*/") { + self.block_comment = false; + self.index += 2; } else { - index += size; + self.index += size; } - continue; + return true; } - if char_set { + if self.char_set { match ch { - _ if escaped => escaped = false, - '\\' => escaped = true, - ']' => char_set = false, + _ if self.escaped => self.escaped = false, + '\\' => self.escaped = true, + ']' => self.char_set = false, _ => {} } - index += size; - continue; + self.index += size; + return true; } - if escaped { - escaped = false; - index += size; - continue; + if self.escaped { + self.escaped = false; + self.index += size; + return true; } - if single_quoted { + if self.single_quoted { match ch { - '\\' => escaped = true, - '\'' => single_quoted = false, + '\\' => self.escaped = true, + '\'' => self.single_quoted = false, _ => {} } - index += size; - continue; + self.index += size; + return true; } - if double_quoted { + if self.double_quoted { match ch { - '\\' => escaped = true, - '"' => double_quoted = false, + '\\' => self.escaped = true, + '"' => self.double_quoted = false, _ => {} } - index += size; - continue; - } - match ch { - '/' if source.as_bytes().get(index..index + 2) == Some(b"//") => { - line_comment = true; - index += 2; - } - '/' if source.as_bytes().get(index..index + 2) == Some(b"/*") => { - block_comment = true; - index += 2; - } - '\'' => { - single_quoted = true; - index += size; - } - '"' => { - double_quoted = true; - index += size; - } - '[' => { - char_set = true; - index += size; - } - '{' => return Some(index), - _ => index += size, + self.index += size; + return true; } + false } - None } /// Finds grammar predicate templates in the same order as ANTLR serializes @@ -5049,93 +5132,22 @@ fn statement_rule_header(source: &str, position: usize) -> Option fn last_rule_header_colon(source: &str, position: usize) -> Option { let mut last = None; - let mut index = 0; - let mut single_quoted = false; - let mut double_quoted = false; - let mut escaped = false; - let mut line_comment = false; - let mut block_comment = false; - let mut char_set = false; - while index < position { - let ch = source[index..].chars().next()?; - let size = ch.len_utf8(); - if line_comment { - line_comment = ch != '\n'; - index += size; - continue; - } - if block_comment { - if source.as_bytes().get(index..index + 2) == Some(b"*/") { - block_comment = false; - index += 2; - } else { - index += size; - } - continue; - } - if char_set { - match ch { - _ if escaped => escaped = false, - '\\' => escaped = true, - ']' => char_set = false, - _ => {} - } - index += size; - continue; - } - if escaped { - escaped = false; - index += size; - continue; - } - if single_quoted { - match ch { - '\\' => escaped = true, - '\'' => single_quoted = false, - _ => {} - } - index += size; - continue; - } - if double_quoted { - match ch { - '\\' => escaped = true, - '"' => double_quoted = false, - _ => {} - } - index += size; - continue; + let mut cursor = GrammarSourceCursor::new(source, 0); + while let Some((index, ch)) = cursor.next_significant() { + if index >= position { + break; } match ch { - '/' if source.as_bytes().get(index..index + 2) == Some(b"//") => { - line_comment = true; - index += 2; - } - '/' if source.as_bytes().get(index..index + 2) == Some(b"/*") => { - block_comment = true; - index += 2; - } - '\'' => { - single_quoted = true; - index += size; - } - '"' => { - double_quoted = true; - index += size; - } - '[' => { - char_set = true; - index += size; - } - '{' => { - index = matching_action_brace(source, index + 1) - .map_or(index + size, |close| close.saturating_add(1).min(position)); - } - ':' => { - last = Some(index); - index += size; - } - _ => index += size, + // Embedded action bodies may contain colons (Rust paths, ternary + // templates); skip the balanced block instead of scanning it. + '{' => cursor.seek( + matching_action_brace(source, index + 1) + .map_or_else(|| index + ch.len_utf8(), |close| { + close.saturating_add(1).min(position) + }), + ), + ':' => last = Some(index), + _ => {} } } last diff --git a/src/lexer.rs b/src/lexer.rs index 2b6234c5..ad87084e 100644 --- a/src/lexer.rs +++ b/src/lexer.rs @@ -113,6 +113,7 @@ pub struct BaseLexer { line: usize, column: usize, hit_eof: bool, + force_interpreted: bool, errors: Vec, dfa_cache: Rc>, } @@ -273,6 +274,7 @@ where line: 1, column: 0, hit_eof: false, + force_interpreted: false, errors: Vec::new(), dfa_cache: Rc::new(RefCell::new(LexerDfaCache::default())), } @@ -586,6 +588,22 @@ where self.hit_eof = hit_eof; } + /// Routes every token through ATN interpretation even when the generated + /// lexer carries an ahead-of-time compiled DFA. + /// + /// Interpretation is what learns the replayable DFA that + /// [`Self::lexer_dfa_string`] reports, so harnesses asserting on the + /// observed-DFA trace (ANTLR's `showDFA` descriptors) enable this before + /// lexing. + pub const fn set_force_interpreted(&mut self, force_interpreted: bool) { + self.force_interpreted = force_interpreted; + } + + /// Whether compiled-DFA entry points must fall back to interpretation. + pub const fn force_interpreted(&self) -> bool { + self.force_interpreted + } + /// Buffers a lexer diagnostic until the token stream consumer is ready to /// emit errors in parser-compatible order. pub fn record_error(&mut self, line: usize, column: usize, message: impl Into) { diff --git a/tools/parse-bench/run.py b/tools/parse-bench/run.py index 1655b2ca..b93ed2bd 100755 --- a/tools/parse-bench/run.py +++ b/tools/parse-bench/run.py @@ -600,14 +600,30 @@ def write_tree_sitter_runner(work_dir: Path, specs: list[LanguageSpec]) -> Path: SINK = None -def parse_once(parser, src: bytes): +def parse_once(parser, src): global SINK tree = parser.parse(src) - if tree.root_node.has_error: + # tree-sitter-language-pack ships py-tree-sitter bindings on some + # platforms (attributes) and its own Rust binding on others (methods); + # support both so the runner works on every wheel of the pinned version. + root = tree.root_node() if callable(tree.root_node) else tree.root_node + has_error = root.has_error() if callable(root.has_error) else root.has_error + if has_error: raise RuntimeError("tree-sitter parse produced ERROR nodes") SINK = tree +def probe_source(parser, raw: bytes): + try: + parse_once(parser, raw) + return raw + except TypeError: + # The Rust-binding wheels accept str input, not bytes. + src = raw.decode("utf-8") + parse_once(parser, src) + return src + + def main() -> int: arg_parser = argparse.ArgumentParser() arg_parser.add_argument("--language", required=True, choices=sorted(SPECS)) @@ -618,7 +634,7 @@ def main() -> int: if args.iters <= 0: raise SystemExit("--iters must be greater than 0") parser = get_parser(SPECS[args.language]) - src = Path(args.input).read_bytes() + src = probe_source(parser, Path(args.input).read_bytes()) for _ in range(args.warmups): parse_once(parser, src) elapsed = []