diff --git a/README.md b/README.md index 4be9a27e..03f23b2c 100644 --- a/README.md +++ b/README.md @@ -510,10 +510,10 @@ group (**> 1.0** means Rust is faster than Go; **< 1.0** means slower): | Language | Fixtures | Rust vs Go (parse time) | |----------|---------:|-------------------------| -| Kotlin | 4 | 24.924x | -| Java | 4 | 3.088x | -| C# | 4 | 2.199x | -| Trino SQL | 5 | 3.390x | +| Kotlin | 4 | 30.336x | +| Java | 4 | 3.166x | +| C# | 4 | 2.177x | +| Trino SQL | 5 | 3.311x | ## Useful Information diff --git a/src/atn/lexer.rs b/src/atn/lexer.rs index bccb6803..cc28ed65 100644 --- a/src/atn/lexer.rs +++ b/src/atn/lexer.rs @@ -1,22 +1,27 @@ use std::cell::RefCell; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::hash::BuildHasherDefault; use crate::atn::lexer_dfa::{ - CompiledLexerAccept, CompiledLexerContinuation, CompiledLexerDfa, DEAD_STATE, ESCAPE_STATE, + CompiledLexerAccept, CompiledLexerContext, CompiledLexerContinuation, CompiledLexerDfa, + DEAD_STATE, ESCAPE_STATE, }; use crate::atn::{AtnStateKind, LexerAction, LexerAtn, LexerTransition}; use crate::char_stream::{CharStream, TextInterval}; use crate::int_stream::EOF; use crate::lexer::{ - BaseLexer, Lexer, LexerCustomAction, LexerDfaActionKey, LexerDfaCachedAccept, - LexerDfaCachedState, LexerDfaCachedTransition, LexerDfaConfigKey, LexerDfaKey, - LexerLifecycleCtx, LexerPredicate, LexerSemCtx, + BaseLexer, EMPTY_LEXER_CONTEXT, Lexer, LexerContextArena, LexerContextId, LexerContextNode, + LexerCustomAction, LexerDfaActionKey, LexerDfaCachedAccept, LexerDfaCachedState, + LexerDfaCachedTransition, LexerDfaConfigKey, LexerDfaKey, LexerLifecycleCtx, LexerPredicate, + LexerSemCtx, }; use crate::parser::{SemanticHooks, UnknownSemanticPolicy}; -use crate::prediction::PredictionFxHasher; +use crate::prediction::{PredictionFxHasher, PredictionWorkspace}; use crate::token::{INVALID_TOKEN_TYPE, TokenId, TokenSink, TokenStoreError}; +#[allow(clippy::disallowed_types)] +type FxHashMap = HashMap>; + #[allow(clippy::disallowed_types)] type FxHashSet = HashSet>; @@ -30,10 +35,117 @@ pub(super) struct LexerConfig { pub(super) consumed_eof: bool, pub(super) alt_rule_index: Option, pub(super) passed_non_greedy: bool, - pub(super) stack: Vec, + pub(super) context: LexerContextId, pub(super) actions: Vec, } +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct LexerConfigKey { + state: usize, + position: usize, + consumed_eof: bool, + alt_rule_index: Option, + passed_non_greedy: bool, + actions: Vec, +} + +impl From<&LexerConfig> for LexerConfigKey { + fn from(config: &LexerConfig) -> Self { + Self { + state: config.state, + position: config.position, + consumed_eof: config.consumed_eof, + alt_rule_index: config.alt_rule_index, + passed_non_greedy: config.passed_non_greedy, + actions: config.actions.clone(), + } + } +} + +trait LexerContextOps { + fn singleton(&mut self, parent: LexerContextId, return_state: usize) -> LexerContextId; + fn merge(&mut self, left: LexerContextId, right: LexerContextId) -> LexerContextId; + fn node(&self, context: LexerContextId) -> LexerContextNode; +} + +struct BorrowedLexerContexts<'a> { + contexts: &'a mut LexerContextArena, + workspace: &'a mut PredictionWorkspace, +} + +impl LexerContextOps for BorrowedLexerContexts<'_> { + fn singleton(&mut self, parent: LexerContextId, return_state: usize) -> LexerContextId { + self.contexts.singleton(parent, return_state) + } + + fn merge(&mut self, left: LexerContextId, right: LexerContextId) -> LexerContextId { + self.contexts.merge(left, right, &mut *self.workspace) + } + + fn node(&self, context: LexerContextId) -> LexerContextNode { + self.contexts.node(context) + } +} + +struct SharedLexerContexts<'a, I> { + lexer: &'a BaseLexer, +} + +impl LexerContextOps for SharedLexerContexts<'_, I> +where + I: CharStream, +{ + fn singleton(&mut self, parent: LexerContextId, return_state: usize) -> LexerContextId { + self.lexer + .lexer_prediction_store() + .contexts + .singleton(parent, return_state) + } + + fn merge(&mut self, left: LexerContextId, right: LexerContextId) -> LexerContextId { + let mut prediction = self.lexer.lexer_prediction_store(); + let prediction = &mut *prediction; + prediction + .contexts + .merge(left, right, &mut prediction.workspace) + } + + fn node(&self, context: LexerContextId) -> LexerContextNode { + self.lexer.lexer_prediction_store().contexts.node(context) + } +} + +/// Ordered lexer configurations with graph-structured caller contexts. +/// +/// Configurations which differ only by their caller path are behaviorally +/// equivalent until a rule stop pops that path, so retaining one config with +/// the union of those paths avoids materializing every concrete call stack. +#[derive(Debug, Default)] +struct LexerConfigSet { + configs: Vec, + config_index: FxHashMap, +} + +impl LexerConfigSet { + fn add(&mut self, config: LexerConfig, contexts: &mut C) + where + C: LexerContextOps, + { + let key = LexerConfigKey::from(&config); + if let Some(&index) = self.config_index.get(&key) { + let existing = self.configs[index].context; + self.configs[index].context = contexts.merge(existing, config.context); + return; + } + self.config_index.insert(key, self.configs.len()); + self.configs.push(config); + } + + fn into_configs(self) -> Vec { + self.configs + } +} + #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] pub(super) struct LexerActionTrace { pub(super) action_index: usize, @@ -59,6 +171,14 @@ enum MatchResult { NoViableAlt { stop: usize }, } +struct InterpretedMatchState { + active: Vec, + dfa_state: usize, + dfa_state_has_semantic_context: bool, + best: Option, + error_stop: usize, +} + #[derive(Clone, Debug)] pub(super) struct ClosureResult { pub(super) configs: Vec, @@ -100,8 +220,8 @@ where /// Accumulates one epsilon-closure expansion, including whether predicate /// evaluation made the closure input-position-sensitive. struct ClosureState { - seen: FxHashSet, - closed: Vec, + expanded: FxHashMap, + closed: LexerConfigSet, has_semantic_context: bool, } @@ -673,6 +793,7 @@ where I: CharStream, P: FnMut(&BaseLexer, LexerPredicate) -> bool, { + lexer.reset_lexer_prediction_workspace(); if let Some(dfa) = strategy.compiled && !lexer.force_interpreted() && let Some(start_state) = dfa.mode_start(mode) @@ -681,7 +802,14 @@ where CompiledMatch::Complete(result) => return result, CompiledMatch::Resume(resume) => { if compiled_resume_matches_atn(atn, start, &resume) { - return match_token_from_continuation(lexer, atn, resume, semantic_predicate); + return match_token_from_continuation( + lexer, + atn, + mode, + start, + resume, + semantic_predicate, + ); } } CompiledMatch::Restart => {} @@ -860,7 +988,8 @@ where let Some(start_state) = atn.mode_to_start_state().get(mode_index).copied() else { return MatchResult::NoViableAlt { stop: start }; }; - let start_closure = epsilon_closure( + let start_closure = epsilon_closure_with_lexer( + lexer, atn, [LexerConfig { state: start_state, @@ -868,20 +997,50 @@ where consumed_eof: false, alt_rule_index: None, passed_non_greedy: false, - stack: Vec::new(), + context: EMPTY_LEXER_CONTEXT, actions: Vec::new(), }], - &mut |predicate| semantic_predicate(lexer, predicate), + semantic_predicate, ); - let mut active = prune_after_accepts(atn, start_closure.configs); - let mut dfa_state = lexer.lexer_dfa_state( + let active = prune_after_accepts(atn, start_closure.configs); + let dfa_state = lexer.lexer_dfa_state( lexer_dfa_key(&active, start), accept_prediction(atn, &active), ); - let mut dfa_state_has_semantic_context = start_closure.has_semantic_context; + let best = best_accept(atn, &active); + continue_interpreted_match( + lexer, + atn, + start, + InterpretedMatchState { + active, + dfa_state, + dfa_state_has_semantic_context: start_closure.has_semantic_context, + best, + error_stop: start, + }, + semantic_predicate, + ) +} - let mut best = best_accept(atn, &active); - let mut error_stop = start; +fn continue_interpreted_match( + lexer: &mut BaseLexer, + atn: &LexerAtn, + start: usize, + state: InterpretedMatchState, + semantic_predicate: &mut P, +) -> MatchResult +where + I: CharStream, + P: FnMut(&BaseLexer, LexerPredicate) -> bool, +{ + let InterpretedMatchState { + mut active, + mut dfa_state, + mut dfa_state_has_semantic_context, + mut best, + mut error_stop, + } = state; while !active.is_empty() { let position = active[0].position; debug_assert!( @@ -914,9 +1073,7 @@ where } } - let closure = epsilon_closure(atn, next, &mut |predicate| { - semantic_predicate(lexer, predicate) - }); + let closure = epsilon_closure_with_lexer(lexer, atn, next, semantic_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); @@ -932,15 +1089,7 @@ where } } } - if let Some(accept) = best_accept(atn, &active) { - if best.as_ref().is_none_or(|current| { - accept.position > current.position - || (accept.position == current.position - && accept.rule_index < current.rule_index) - }) { - best = Some(accept); - } - } + update_best_accept(atn, &active, &mut best); } best.map_or( @@ -949,6 +1098,14 @@ where ) } +enum CachedModeStart { + Cached(usize), + Evaluated { + dfa_state: usize, + active: Vec, + }, +} + fn match_token_cached( lexer: &mut BaseLexer, atn: &LexerAtn, @@ -960,18 +1117,59 @@ where I: CharStream, P: FnMut(&BaseLexer, LexerPredicate) -> bool, { - let Some((mut dfa_state, mode_start_has_semantic_context)) = - cached_mode_start_state(lexer, atn, mode, start, semantic_predicate) + let Some(mode_start) = cached_mode_start_state(lexer, atn, mode, start, semantic_predicate) else { return MatchResult::NoViableAlt { stop: start }; }; - if mode_start_has_semantic_context { - return match_token(lexer, atn, mode, start, semantic_predicate); - } + let dfa_state = match mode_start { + CachedModeStart::Cached(dfa_state) => dfa_state, + CachedModeStart::Evaluated { dfa_state, active } => { + let best = best_accept(atn, &active); + return continue_interpreted_match( + lexer, + atn, + start, + InterpretedMatchState { + active, + dfa_state, + dfa_state_has_semantic_context: true, + best, + error_stop: start, + }, + semantic_predicate, + ); + } + }; - let mut position = start; - let mut best = None; - let mut error_stop = start; + match_token_cached_from_state( + lexer, + atn, + mode, + start, + dfa_state, + start, + None, + start, + semantic_predicate, + ) +} + +#[allow(clippy::too_many_arguments)] +fn match_token_cached_from_state( + lexer: &mut BaseLexer, + atn: &LexerAtn, + mode: i32, + start: usize, + mut dfa_state: usize, + mut position: usize, + mut best: Option, + mut error_stop: usize, + semantic_predicate: &mut P, +) -> MatchResult +where + I: CharStream, + P: FnMut(&BaseLexer, LexerPredicate) -> bool, +{ loop { let Some(cached_state) = lexer.cached_lexer_dfa_state(dfa_state) else { return match_token(lexer, atn, mode, start, semantic_predicate); @@ -1031,21 +1229,16 @@ where } } - let closure = epsilon_closure(atn, next, &mut |predicate| { - semantic_predicate(lexer, predicate) - }); + let closure = epsilon_closure_with_lexer(lexer, atn, next, semantic_predicate); let target_has_semantic_context = closure.has_semantic_context; - if target_has_semantic_context { - return match_token(lexer, atn, mode, start, semantic_predicate); - } let suppress_edge = source_has_semantic_context || target_has_semantic_context; let active = prune_after_accepts(atn, closure.configs); + update_best_accept(atn, &active, &mut best); if active.is_empty() { break; } - let Some(target_position) = shared_config_position(&active) else { - return match_token(lexer, atn, mode, start, semantic_predicate); - }; + let target_position = + shared_config_position(&active).expect("lexer ATN configs advance in lockstep"); dfa_state = cache_dfa_state( lexer, atn, @@ -1054,6 +1247,21 @@ where start, target_position, ); + if target_has_semantic_context { + return continue_interpreted_match( + lexer, + atn, + start, + InterpretedMatchState { + active, + dfa_state, + dfa_state_has_semantic_context: true, + best, + error_stop, + }, + semantic_predicate, + ); + } if !suppress_edge && symbol != EOF { lexer.record_lexer_dfa_edge(source_dfa_state, symbol, dfa_state); lexer.cache_lexer_dfa_transition( @@ -1098,16 +1306,32 @@ fn compiled_resume_matches_atn(atn: &LexerAtn, start: usize, resume: &CompiledRe return false; }; let rule_count = atn.rule_to_token_type().len(); - !resume.continuation.configs.is_empty() + let contexts_valid = resume + .continuation + .contexts + .iter() + .enumerate() + .all(|(index, context)| { + let local_id = u32::try_from(index + 1).ok(); + local_id.is_some_and(|local_id| match context { + CompiledLexerContext::Singleton { + parent, + return_state, + } => *parent < local_id && atn.state(*return_state).is_some(), + CompiledLexerContext::Union { left, right } => { + *left < local_id && *right < local_id + } + }) + }); + contexts_valid + && !resume.continuation.configs.is_empty() && resume.continuation.configs.iter().all(|config| { atn.state(config.state).is_some() && config .alt_rule_index .is_some_and(|rule_index| rule_index < rule_count) - && config - .stack - .iter() - .all(|&state_number| atn.state(state_number).is_some()) + && usize::try_from(config.context) + .is_ok_and(|context| context <= resume.continuation.contexts.len()) && config.actions.iter().all(|action| { action.action_index < atn.lexer_actions().len() && action.rule_index < rule_count @@ -1302,6 +1526,8 @@ where fn match_token_from_continuation( lexer: &mut BaseLexer, atn: &LexerAtn, + mode: i32, + start: usize, resume: CompiledResume<'_>, semantic_predicate: &mut P, ) -> MatchResult @@ -1313,73 +1539,88 @@ where continuation, position, mut best, - mut error_stop, + error_stop, } = resume; - let moved = continuation.configs.iter().map(|config| LexerConfig { - state: config.state, - position, - consumed_eof: config.consumed_eof, - alt_rule_index: config.alt_rule_index, - passed_non_greedy: config.passed_non_greedy, - stack: config.stack.clone(), - actions: config - .actions + let moved = { + let mut prediction = lexer.lexer_prediction_store(); + let prediction = &mut *prediction; + let mut contexts = Vec::with_capacity(continuation.contexts.len() + 1); + contexts.push(EMPTY_LEXER_CONTEXT); + for compiled in &continuation.contexts { + let imported = match *compiled { + CompiledLexerContext::Singleton { + parent, + return_state, + } => prediction + .contexts + .singleton(contexts[parent as usize], return_state), + CompiledLexerContext::Union { left, right } => prediction.contexts.merge( + contexts[left as usize], + contexts[right as usize], + &mut prediction.workspace, + ), + }; + contexts.push(imported); + } + continuation + .configs .iter() - .map(|action| LexerActionTrace { - action_index: action.action_index, - position: position.saturating_sub(action.behind), - rule_index: action.rule_index, + .map(|config| LexerConfig { + state: config.state, + position, + consumed_eof: config.consumed_eof, + alt_rule_index: config.alt_rule_index, + passed_non_greedy: config.passed_non_greedy, + context: contexts[config.context as usize], + actions: config + .actions + .iter() + .map(|action| LexerActionTrace { + action_index: action.action_index, + position: position.saturating_sub(action.behind), + rule_index: action.rule_index, + }) + .collect(), }) - .collect(), - }); - let closure = epsilon_closure(atn, moved, &mut |predicate| { - semantic_predicate(lexer, predicate) - }); - let mut active = prune_after_accepts(atn, closure.configs); + .collect::>() + }; + let closure = epsilon_closure_with_lexer(lexer, atn, moved, semantic_predicate); + let has_semantic_context = closure.has_semantic_context; + let active = prune_after_accepts(atn, closure.configs); update_best_accept(atn, &active, &mut best); - - while let Some(config) = active.first() { - let active_position = config.position; - debug_assert!( - active - .iter() - .all(|config| config.position == active_position), - "lexer ATN configs must advance through the input in lockstep" + if active.is_empty() { + return best.map_or( + MatchResult::NoViableAlt { stop: error_stop }, + MatchResult::Accept, ); - let symbol = symbol_at(lexer, active_position); - if symbol != EOF { - error_stop = error_stop.max(active_position.saturating_add(1)); - } - let mut next = Vec::new(); - for config in active { - let Some(state) = atn.state(config.state) else { - continue; - }; - for transition in &state.transitions { - if !transition.matches(symbol, MIN_CHAR_VALUE, MAX_CHAR_VALUE) { - continue; - } - let mut advanced = config.clone(); - set_config_state(atn, &mut advanced, transition.target()); - if symbol == EOF { - advanced.consumed_eof = true; - } else { - advanced.position += 1; - } - next.push(advanced); - } - } - - let closure = epsilon_closure(atn, next, &mut |predicate| { - semantic_predicate(lexer, predicate) - }); - active = prune_after_accepts(atn, closure.configs); - update_best_accept(atn, &active, &mut best); } - - best.map_or( - MatchResult::NoViableAlt { stop: error_stop }, - MatchResult::Accept, + let position = shared_config_position(&active).expect("lexer ATN configs advance in lockstep"); + let dfa_state = cache_dfa_state(lexer, atn, &active, has_semantic_context, start, position); + if has_semantic_context { + return continue_interpreted_match( + lexer, + atn, + start, + InterpretedMatchState { + active, + dfa_state, + dfa_state_has_semantic_context: true, + best, + error_stop, + }, + semantic_predicate, + ); + } + match_token_cached_from_state( + lexer, + atn, + mode, + start, + dfa_state, + position, + best, + error_stop, + semantic_predicate, ) } @@ -1432,18 +1673,19 @@ fn cached_mode_start_state( mode: i32, start: usize, semantic_predicate: &mut P, -) -> Option<(usize, bool)> +) -> Option where I: CharStream, P: FnMut(&BaseLexer, LexerPredicate) -> bool, { if let Some(state) = lexer.cached_lexer_mode_start(mode) { - return Some((state, false)); + return Some(CachedModeStart::Cached(state)); } 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( + let start_closure = epsilon_closure_with_lexer( + lexer, atn, [LexerConfig { state: start_state, @@ -1451,10 +1693,10 @@ where consumed_eof: false, alt_rule_index: None, passed_non_greedy: false, - stack: Vec::new(), + context: EMPTY_LEXER_CONTEXT, actions: Vec::new(), }], - &mut |predicate| semantic_predicate(lexer, predicate), + semantic_predicate, ); let active = prune_after_accepts(atn, start_closure.configs); let state = cache_dfa_state( @@ -1465,10 +1707,14 @@ where start, start, ); - if !start_closure.has_semantic_context { - lexer.cache_lexer_mode_start(mode, state); + if start_closure.has_semantic_context { + return Some(CachedModeStart::Evaluated { + dfa_state: state, + active, + }); } - Some((state, start_closure.has_semantic_context)) + lexer.cache_lexer_mode_start(mode, state); + Some(CachedModeStart::Cached(state)) } fn cache_dfa_state( @@ -1537,32 +1783,67 @@ fn cached_accept_state( } } +fn epsilon_closure_with_lexer( + lexer: &BaseLexer, + atn: &LexerAtn, + configs: impl IntoIterator, + semantic_predicate: &mut P, +) -> ClosureResult +where + I: CharStream, + P: FnMut(&BaseLexer, LexerPredicate) -> bool, +{ + let mut contexts = SharedLexerContexts { lexer }; + epsilon_closure_with_contexts(atn, configs, &mut contexts, &mut |predicate| { + semantic_predicate(lexer, predicate) + }) +} + /// Expands epsilon, rule-call, predicate, precedence, and action transitions /// without consuming input. /// -/// 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. +/// Lexer rule calls use graph-structured caller contexts. Equivalent configs +/// merge their contexts in the returned ordered set, retaining all return +/// paths without cloning every concrete stack. pub(super) fn epsilon_closure

( atn: &LexerAtn, configs: impl IntoIterator, + contexts: &mut LexerContextArena, + workspace: &mut PredictionWorkspace, semantic_predicate: &mut P, ) -> ClosureResult where P: FnMut(LexerPredicate) -> bool, +{ + let mut contexts = BorrowedLexerContexts { + contexts, + workspace, + }; + epsilon_closure_with_contexts(atn, configs, &mut contexts, semantic_predicate) +} + +fn epsilon_closure_with_contexts( + atn: &LexerAtn, + configs: impl IntoIterator, + contexts: &mut C, + semantic_predicate: &mut P, +) -> ClosureResult +where + C: LexerContextOps, + P: FnMut(LexerPredicate) -> bool, { let mut state = ClosureState { - seen: FxHashSet::default(), - closed: Vec::new(), + expanded: FxHashMap::default(), + closed: LexerConfigSet::default(), has_semantic_context: false, }; for config in configs { - close_config(atn, config, &mut state, semantic_predicate); + close_config(atn, config, contexts, &mut state, semantic_predicate); } ClosureResult { - configs: state.closed, + configs: state.closed.into_configs(), has_semantic_context: state.has_semantic_context, } } @@ -1573,16 +1854,27 @@ 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

( +fn close_config( atn: &LexerAtn, config: LexerConfig, + contexts: &mut C, closure: &mut ClosureState, semantic_predicate: &mut P, ) where + C: LexerContextOps, P: FnMut(LexerPredicate) -> bool, { - if !closure.seen.insert(config.clone()) { - return; + let key = LexerConfigKey::from(&config); + if let Some(existing) = closure.expanded.get(&key).copied() { + let merged = contexts.merge(existing, config.context); + if merged == existing { + return; + } + closure.expanded.insert(key, merged); + // `config` is the newly discovered ordered delta. Re-expanding the + // merged prefix would duplicate earlier paths ahead of this one. + } else { + closure.expanded.insert(key, config.context); } let Some(state) = atn.state(config.state) else { @@ -1590,13 +1882,33 @@ fn close_config

( }; if state.kind == AtnStateKind::RuleStop { - if let Some((&follow_state, rest)) = config.stack.split_last() { - let mut returned = config.clone(); - set_config_state(atn, &mut returned, follow_state); - returned.stack = rest.to_vec(); - close_config(atn, returned, closure, semantic_predicate); + let mut pending = vec![config.context]; + let mut visited = FxHashSet::default(); + while let Some(context) = pending.pop() { + if !visited.insert(context) { + continue; + } + match contexts.node(context) { + LexerContextNode::Empty => { + let mut accepted = config.clone(); + accepted.context = EMPTY_LEXER_CONTEXT; + closure.closed.add(accepted, contexts); + } + LexerContextNode::Singleton { + parent, + return_state, + } => { + let mut returned = config.clone(); + set_config_state(atn, &mut returned, return_state); + returned.context = parent; + close_config(atn, returned, contexts, closure, semantic_predicate); + } + LexerContextNode::Union { left, right } => { + pending.push(right); + pending.push(left); + } + } } - closure.closed.push(config); return; } @@ -1606,7 +1918,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(atn, next, closure, semantic_predicate); + close_config(atn, next, contexts, closure, semantic_predicate); } LexerTransition::Rule { target, @@ -1616,8 +1928,8 @@ fn close_config

( let mut next = config.clone(); set_config_state(atn, &mut next, *target); next.passed_non_greedy |= state.non_greedy; - next.stack.push(*follow_state); - close_config(atn, next, closure, semantic_predicate); + next.context = contexts.singleton(config.context, *follow_state); + close_config(atn, next, contexts, closure, semantic_predicate); } LexerTransition::Predicate { target, @@ -1634,14 +1946,14 @@ fn close_config

( let mut next = config.clone(); set_config_state(atn, &mut next, *target); next.passed_non_greedy |= state.non_greedy; - close_config(atn, next, closure, semantic_predicate); + close_config(atn, next, contexts, closure, semantic_predicate); } } LexerTransition::Precedence { target, .. } => { let mut next = config.clone(); set_config_state(atn, &mut next, *target); next.passed_non_greedy |= state.non_greedy; - close_config(atn, next, closure, semantic_predicate); + close_config(atn, next, contexts, closure, semantic_predicate); } LexerTransition::Action { target, @@ -1653,13 +1965,19 @@ fn close_config

( set_config_state(atn, &mut next, *target); next.passed_non_greedy |= state.non_greedy; if let Some(action_index) = action_index { - next.actions.push(LexerActionTrace { + let trace = LexerActionTrace { action_index: *action_index, position: config.position, rule_index: *rule_index, + }; + let keep = next.alt_rule_index.is_none_or(|accept_rule| { + lexer_action_belongs_to_accept(atn, accept_rule, *rule_index) }); + if keep { + append_lexer_action_trace(atn, &mut next.actions, trace); + } } - close_config(atn, next, closure, semantic_predicate); + close_config(atn, next, contexts, closure, semantic_predicate); } LexerTransition::Atom { .. } | LexerTransition::Range { .. } @@ -1674,10 +1992,58 @@ fn close_config

( .iter() .any(|transition| !transition.is_epsilon()) { - closure.closed.push(config); + closure.closed.add(config, contexts); } } +/// Appends one executable action while removing an earlier setter whose +/// result cannot be observed before this one overwrites it. +fn append_lexer_action_trace( + atn: &LexerAtn, + actions: &mut Vec, + trace: LexerActionTrace, +) { + #[derive(Clone, Copy, Eq, PartialEq)] + enum Setter { + Channel, + Mode, + TokenType, + } + + const fn setter(action: &LexerAction) -> Option { + match action { + LexerAction::Channel(_) => Some(Setter::Channel), + LexerAction::Mode(_) => Some(Setter::Mode), + LexerAction::More | LexerAction::Skip | LexerAction::Type(_) => Some(Setter::TokenType), + LexerAction::Custom { .. } | LexerAction::PopMode | LexerAction::PushMode(_) => None, + } + } + + let Some(action) = atn.lexer_actions().get(trace.action_index) else { + actions.push(trace); + return; + }; + let Some(slot) = setter(action) else { + actions.push(trace); + return; + }; + for index in (0..actions.len()).rev() { + let Some(previous) = atn.lexer_actions().get(actions[index].action_index) else { + break; + }; + if matches!(previous, LexerAction::Custom { .. }) + || (slot == Setter::Mode && matches!(previous, LexerAction::PushMode(_))) + { + break; + } + if setter(previous) == Some(slot) { + actions.remove(index); + break; + } + } + actions.push(trace); +} + /// Removes lower-priority non-greedy configs ordered after a top-level accept /// for the same lexer rule. /// @@ -1697,7 +2063,7 @@ pub(super) fn prune_after_accepts(atn: &LexerAtn, configs: Vec) -> if config.passed_non_greedy && accepted_rules.contains(&rule_index) { continue; } - let is_top_level_accept = config.stack.is_empty() + let is_top_level_accept = config.context == EMPTY_LEXER_CONTEXT && atn .state(config.state) .is_some_and(crate::atn::LexerAtnState::is_rule_stop); @@ -1719,7 +2085,7 @@ pub(super) fn best_accept(atn: &LexerAtn, configs: &[LexerConfig]) -> Option LexerDfaCo config.alt_rule_index, config.consumed_eof, config.passed_non_greedy, - config.stack.clone(), + config.context, config .actions .iter() @@ -1798,7 +2164,7 @@ fn cached_configs_to_configs( consumed_eof: config.consumed_eof, alt_rule_index: config.alt_rule_index, passed_non_greedy: config.passed_non_greedy, - stack: config.stack.clone(), + context: config.context, actions: config .actions .iter() @@ -1870,7 +2236,9 @@ where #[cfg(test)] mod tests { use super::*; - use crate::atn::lexer_dfa::{CompiledLexerActionTrace, CompiledLexerConfig}; + use crate::atn::lexer_dfa::{ + CompiledLexerActionTrace, CompiledLexerConfig, CompiledLexerContext, + }; use crate::atn::serialized::{AtnDeserializer, SerializedAtn}; use crate::atn::{LexerAtnState, LexerTransition}; use crate::char_stream::InputStream; @@ -1893,6 +2261,47 @@ mod tests { ) } + fn predicate_atn() -> LexerAtn { + let mut atn = LexerAtn::new(1); + + let mut mode_start = LexerAtnState::new(0, AtnStateKind::TokenStart); + mode_start.add_transition(LexerTransition::Epsilon { target: 1 }); + atn.add_state(mode_start); + + let mut rule_start = LexerAtnState::new(1, AtnStateKind::RuleStart).with_rule_index(0); + rule_start.add_transition(LexerTransition::Atom { + target: 2, + label: 'a' as i32, + }); + atn.add_state(rule_start); + + let mut call_fragment = LexerAtnState::new(2, AtnStateKind::Basic).with_rule_index(0); + call_fragment.add_transition(LexerTransition::Rule { + target: 4, + rule_index: 1, + follow_state: 3, + precedence: 0, + }); + atn.add_state(call_fragment); + atn.add_state(LexerAtnState::new(3, AtnStateKind::RuleStop).with_rule_index(0)); + + let mut fragment_start = LexerAtnState::new(4, AtnStateKind::RuleStart).with_rule_index(1); + fragment_start.add_transition(LexerTransition::Predicate { + target: 5, + rule_index: 1, + pred_index: 0, + context_dependent: false, + }); + atn.add_state(fragment_start); + atn.add_state(LexerAtnState::new(5, AtnStateKind::RuleStop).with_rule_index(1)); + + atn.set_rule_to_start_state(vec![1, 4]); + atn.set_rule_to_stop_state(vec![3, 5]); + atn.set_rule_to_token_type(vec![1, INVALID_TOKEN_TYPE]); + atn.add_mode_start_state(0); + atn + } + // `BLOCK_COMMENT: ('/**/' | '/*' ~[!] .*? '*/'); OTHER: .;` // // `#[rustfmt::skip]`: this serialized ATN is a positional stream emitted by @@ -1907,6 +2316,34 @@ mod tests { .expect("issue #106 lexer ATN should deserialize") } + // `COMMENT: '/*' (COMMENT | .)*? '*/' -> channel(HIDDEN);` + // `fragment Hidden: COMMENT; AT_PRE: Hidden '@'; OTHER: .;` + // + // This is the reduced issue #135 topology: recursive non-greedy comments + // remain speculative through another token rule, and the referenced token + // rule carries a built-in action. + #[rustfmt::skip] + fn recursive_comment_channel_atn() -> LexerAtn { + AtnDeserializer::new(&SerializedAtn::from_i32(&[ + 4, 0, 3, 31, 6, -1, 2, 0, 7, 0, 2, 1, 7, 1, 2, 2, 7, 2, 2, 3, 7, 3, + 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 5, 0, 15, 8, 0, 10, 0, 12, 0, 18, 9, 0, + 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 2, 1, 2, 1, 2, 1, 3, + 1, 3, 1, 16, 0, 4, 1, 1, 3, 0, 5, 2, 7, 3, 1, 0, 0, 31, 0, 1, 1, 0, + 0, 0, 0, 5, 1, 0, 0, 0, 0, 7, 1, 0, 0, 0, 1, 9, 1, 0, 0, 0, 3, 24, + 1, 0, 0, 0, 5, 26, 1, 0, 0, 0, 7, 29, 1, 0, 0, 0, 9, 10, 5, 47, 0, + 0, 10, 11, 5, 42, 0, 0, 11, 16, 1, 0, 0, 0, 12, 15, 3, 1, 0, 0, 13, + 15, 9, 0, 0, 0, 14, 12, 1, 0, 0, 0, 14, 13, 1, 0, 0, 0, 15, 18, 1, + 0, 0, 0, 16, 17, 1, 0, 0, 0, 16, 14, 1, 0, 0, 0, 17, 19, 1, 0, 0, + 0, 18, 16, 1, 0, 0, 0, 19, 20, 5, 42, 0, 0, 20, 21, 5, 47, 0, 0, 21, + 22, 1, 0, 0, 0, 22, 23, 6, 0, 0, 0, 23, 2, 1, 0, 0, 0, 24, 25, 3, 1, + 0, 0, 25, 4, 1, 0, 0, 0, 26, 27, 3, 3, 1, 0, 27, 28, 5, 64, 0, 0, 28, + 6, 1, 0, 0, 0, 29, 30, 9, 0, 0, 0, 30, 8, 1, 0, 0, 0, 3, 0, 14, 16, + 1, 0, 1, 0, + ])) + .deserialize() + .expect("recursive-comment lexer ATN should deserialize") + } + #[derive(Clone, Copy, Debug)] enum TestMatchStrategy { Interpreted, @@ -2865,6 +3302,35 @@ mod tests { assert!(lexer.cached_lexer_dfa_state(plain_state).is_some()); } + #[test] + fn predicate_can_inspect_and_clear_the_active_dfa() { + let atn = predicate_atn(); + let mut lexer = BaseLexer::new(InputStream::new("a"), recognizer_data()); + let mut store = TokenStore::new(lexer.source_text(), lexer.source_name()); + let mut sink = TokenSink::new(&mut store); + let mut predicate_calls = 0; + + let id = next_token_with_cache( + &mut lexer, + &mut sink, + &atn, + |_, _| {}, + |lexer, _| { + predicate_calls += 1; + let _ = lexer.lexer_dfa_string(); + lexer.clear_dfa(); + true + }, + |_, _, _| {}, + ) + .expect("predicate token should fit"); + let token = sink.view(id).expect("predicate token should exist"); + + assert_eq!(token.token_type(), 1); + assert_eq!(token.text(), "a"); + assert_eq!(predicate_calls, 1); + } + #[test] fn lexer_action_hook_context_can_change_mode() { // A custom-action hook receives a mutable lexer borrow, so it can push / @@ -2967,6 +3433,83 @@ mod tests { } } + #[test] + fn recursive_comment_contexts_and_actions_stay_bounded_on_all_paths() { + let atn = recursive_comment_channel_atn(); + let compiled = CompiledLexerDfa::compile(&atn); + let serialized = compiled.serialize(); + let compiled = + CompiledLexerDfa::from_serialized(&serialized).expect("compiled DFA round trip"); + let first = "/* first /* nested */ tail */"; + let mut source = format!("{first}x"); + for _ in 0..40 { + source.push_str("/* comment */x"); + } + + for strategy in [ + TestMatchStrategy::Interpreted, + TestMatchStrategy::Cached, + TestMatchStrategy::Compiled, + ] { + let data = RecognizerData::new( + "RecursiveComment", + Vocabulary::new( + [None::<&str>, None, None, None], + [None, Some("COMMENT"), Some("AT_PRE"), Some("OTHER")], + [None::<&str>, None, None, None], + ), + ); + let mut lexer = BaseLexer::new(InputStream::new(source.clone()), data); + let mut store = TokenStore::new(lexer.source_text(), lexer.source_name()); + let mut sink = TokenSink::new(&mut store); + let mut token_count = 0; + let mut contexts_after_warmup = 0; + loop { + let token = match strategy { + TestMatchStrategy::Interpreted => next_token_with_hooks( + &mut lexer, + &mut sink, + &atn, + |_, _| {}, + |_, _| true, + |_, _, _| {}, + ), + TestMatchStrategy::Cached => next_token(&mut lexer, &mut sink, &atn), + TestMatchStrategy::Compiled => { + next_token_compiled(&mut lexer, &mut sink, &atn, &compiled) + } + } + .expect("recursive comment source should lex"); + let token = sink.view(token).expect("token should exist"); + if token_count == 0 { + assert_eq!(token.token_type(), 1, "{strategy:?}"); + assert_eq!(token.text(), first, "{strategy:?}"); + assert_eq!(token.channel(), HIDDEN_CHANNEL, "{strategy:?}"); + } else if token_count == 3 { + contexts_after_warmup = lexer.lexer_dfa_cache_shape().3; + } + token_count += 1; + if token.token_type() == TOKEN_EOF { + break; + } + } + assert_eq!(token_count, 83, "{strategy:?}"); + + let (cached_states, cached_transitions, max_configs, contexts) = + lexer.lexer_dfa_cache_shape(); + assert!(max_configs <= 16, "{strategy:?}: {max_configs} configs"); + assert!(contexts <= 1024, "{strategy:?}: {contexts} contexts"); + assert!( + contexts <= contexts_after_warmup + 16, + "{strategy:?}: contexts grew from {contexts_after_warmup} to {contexts}" + ); + if !matches!(strategy, TestMatchStrategy::Interpreted) { + assert!(cached_states > 0, "{strategy:?}"); + assert!(cached_transitions > 0, "{strategy:?}"); + } + } + } + #[test] fn compiled_resume_rejects_untrusted_continuation_payloads() { let atn = trailing_action_atn(&['a'], 1, vec![LexerAction::Skip]); @@ -2975,7 +3518,7 @@ mod tests { consumed_eof: false, alt_rule_index: Some(0), passed_non_greedy: false, - stack: Vec::new(), + context: 0, actions: vec![CompiledLexerActionTrace { action_index: 0, rule_index: 0, @@ -2983,6 +3526,7 @@ mod tests { }], }; let valid = |config| CompiledLexerContinuation { + contexts: Vec::new(), configs: vec![config], }; let matches = |continuation: &CompiledLexerContinuation| { @@ -3000,6 +3544,7 @@ mod tests { assert!(matches(&valid(valid_config.clone()))); assert!(!matches(&CompiledLexerContinuation { + contexts: Vec::new(), configs: Vec::new(), })); @@ -3016,9 +3561,17 @@ mod tests { assert!(!matches(&valid(invalid))); let mut invalid = valid_config.clone(); - invalid.stack.push(usize::MAX); + invalid.context = u32::MAX; assert!(!matches(&valid(invalid))); + assert!(!matches(&CompiledLexerContinuation { + contexts: vec![CompiledLexerContext::Singleton { + parent: 1, + return_state: 2, + }], + configs: vec![valid_config.clone()], + })); + let mut invalid = valid_config.clone(); invalid.actions[0].action_index = 1; assert!(!matches(&valid(invalid))); diff --git a/src/atn/lexer_dfa.rs b/src/atn/lexer_dfa.rs index 7b9395c1..a8799a7d 100644 --- a/src/atn/lexer_dfa.rs +++ b/src/atn/lexer_dfa.rs @@ -9,7 +9,7 @@ //! //! 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 +//! time), grows an unbounded caller-context path (recursive lexer rules such as //! nested comments), or exceeds the state budget is compiled as an *escape* //! edge. Dynamic closures carry their pre-closure configs so the interpreter //! resumes from the narrowed edge instead of re-matching from the token start. @@ -31,8 +31,11 @@ use crate::atn::lexer::{ }; use crate::atn::{LexerAtn, LexerTransition}; use crate::int_stream::EOF; -use crate::lexer::{LexerDfaActionKey, LexerDfaConfigKey, LexerDfaKey}; -use crate::prediction::PredictionFxHasher; +use crate::lexer::{ + EMPTY_LEXER_CONTEXT, LexerContextArena, LexerContextId, LexerContextNode, LexerDfaActionKey, + LexerDfaConfigKey, LexerDfaKey, +}; +use crate::prediction::{PredictionFxHasher, PredictionWorkspace}; #[allow(clippy::disallowed_types)] type FxHashMap = HashMap>; @@ -50,9 +53,9 @@ pub(super) const ESCAPE_STATE: u16 = u16::MAX - 1; /// 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 +/// Caller-context paths 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; +const MAX_CONTEXT_DEPTH: usize = 32; /// Configs whose surviving action trace grows past this escape to the /// interpreter: a custom action crossed inside a loop is genuinely @@ -266,9 +269,17 @@ struct CompiledLexerEscapeRange { /// its dynamic epsilon closure. #[derive(Clone, Debug)] pub(super) struct CompiledLexerContinuation { + pub(super) contexts: Vec, pub(super) configs: Vec, } +/// One ordered caller-context node in a continuation-local topological table. +#[derive(Clone, Copy, Debug)] +pub(super) enum CompiledLexerContext { + Singleton { parent: u32, return_state: usize }, + Union { left: u32, right: u32 }, +} + /// Position-independent ATN config stored in an escape continuation. #[derive(Clone, Debug)] pub(super) struct CompiledLexerConfig { @@ -276,7 +287,7 @@ pub(super) struct CompiledLexerConfig { pub(super) consumed_eof: bool, pub(super) alt_rule_index: Option, pub(super) passed_non_greedy: bool, - pub(super) stack: Vec, + pub(super) context: u32, pub(super) actions: Vec, } @@ -458,11 +469,12 @@ impl CompiledLexerDfa { .continuations .iter() .map(|continuation| { - 1 + continuation - .configs - .iter() - .map(|config| 6 + config.stack.len() + config.actions.len() * 3) - .sum::() + 2 + continuation.contexts.len() * 3 + + continuation + .configs + .iter() + .map(|config| 6 + config.actions.len() * 3) + .sum::() }) .sum(); let capacity = 9 @@ -529,16 +541,34 @@ impl CompiledLexerDfa { } out.push(self.continuations.len() as u32); for continuation in &self.continuations { + out.push(continuation.contexts.len() as u32); + for context in &continuation.contexts { + match context { + CompiledLexerContext::Singleton { + parent, + return_state, + } => { + out.push(0); + out.push(*parent); + out.push( + u32::try_from(*return_state) + .expect("lexer context return state must fit in u32"), + ); + } + CompiledLexerContext::Union { left, right } => { + out.push(1); + out.push(*left); + out.push(*right); + } + } + } out.push(continuation.configs.len() as u32); for config in &continuation.configs { out.push(config.state as u32); out.push(config.alt_rule_index.map_or(u32::MAX, |rule| rule as u32)); out.push(u32::from(config.consumed_eof)); out.push(u32::from(config.passed_non_greedy)); - out.push(config.stack.len() as u32); - for &state in &config.stack { - out.push(state as u32); - } + out.push(config.context); out.push(config.actions.len() as u32); for action in &config.actions { out.push(action.action_index as u32); @@ -648,6 +678,10 @@ impl CompiledLexerDfa { .iter() .all(|range| continuation_ok(range.continuation)) }) + && self + .continuations + .iter() + .all(compiled_continuation_contexts_are_valid) } #[cfg(feature = "perf-counters")] @@ -658,6 +692,29 @@ impl CompiledLexerDfa { } } +fn compiled_continuation_contexts_are_valid(continuation: &CompiledLexerContinuation) -> bool { + let contexts_valid = continuation + .contexts + .iter() + .enumerate() + .all(|(index, context)| { + let local_id = u32::try_from(index + 1).ok(); + local_id.is_some_and(|local_id| match context { + CompiledLexerContext::Singleton { parent, .. } => *parent < local_id, + CompiledLexerContext::Union { left, right } => { + *left < local_id && *right < local_id + } + }) + }); + contexts_valid + && u32::try_from(continuation.contexts.len()).is_ok_and(|max_context| { + continuation + .configs + .iter() + .all(|config| config.context <= max_context) + }) +} + /// 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. @@ -673,7 +730,7 @@ fn escape_row_is_searchable(row: &[CompiledLexerEscapeRange]) -> bool { /// Version tag guarding embedded tables against format or construction-semantic /// drift. -const SERIALIZED_TAG: u32 = 0x4C58_4405; +const SERIALIZED_TAG: u32 = 0x4C58_4406; /// Cursor over a serialized DFA stream. struct SerializedReader<'a> { @@ -802,6 +859,24 @@ impl SerializedReader<'_> { let count = self.next_len()?; let mut continuations = Vec::with_capacity(count.min(self.data.len())); for _ in 0..count { + let context_count = self.next_len()?; + let mut contexts = Vec::with_capacity(context_count.min(self.data.len())); + for _ in 0..context_count { + let kind = self.next()?; + let first = self.next()?; + let second = self.next()?; + contexts.push(match kind { + 0 => CompiledLexerContext::Singleton { + parent: first, + return_state: usize::try_from(second).ok()?, + }, + 1 => CompiledLexerContext::Union { + left: first, + right: second, + }, + _ => return None, + }); + } let config_count = self.next_len()?; let mut configs = Vec::with_capacity(config_count.min(self.data.len())); for _ in 0..config_count { @@ -814,11 +889,7 @@ impl SerializedReader<'_> { }; let consumed_eof = self.next()? != 0; let passed_non_greedy = self.next()? != 0; - let stack_count = self.next_len()?; - let mut stack = Vec::with_capacity(stack_count.min(self.data.len())); - for _ in 0..stack_count { - stack.push(self.next_len()?); - } + let context = self.next()?; let action_count = self.next_len()?; let mut actions = Vec::with_capacity(action_count.min(self.data.len())); for _ in 0..action_count { @@ -833,11 +904,11 @@ impl SerializedReader<'_> { consumed_eof, alt_rule_index, passed_non_greedy, - stack, + context, actions, }); } - continuations.push(CompiledLexerContinuation { configs }); + continuations.push(CompiledLexerContinuation { contexts, configs }); } Some(continuations) } @@ -882,6 +953,8 @@ impl RowPools { struct ModeBuild { base: usize, continuation_base: usize, + contexts: LexerContextArena, + workspace: PredictionWorkspace, ids: FxHashMap, configs: Vec>, steps: Vec, @@ -916,6 +989,8 @@ impl ModeBuild { Self { base, continuation_base, + contexts: LexerContextArena::new(), + workspace: PredictionWorkspace::default(), ids: FxHashMap::default(), configs: Vec::new(), steps: Vec::new(), @@ -962,26 +1037,36 @@ impl ModeBuild { let Ok(id) = u32::try_from(id) else { return u32::MAX; }; + let mut context_ids = FxHashMap::default(); + context_ids.insert(EMPTY_LEXER_CONTEXT, 0); + let mut contexts = Vec::new(); + let compiled_configs = configs + .iter() + .map(|config| CompiledLexerConfig { + state: config.state, + consumed_eof: config.consumed_eof, + alt_rule_index: config.alt_rule_index, + passed_non_greedy: config.passed_non_greedy, + context: compile_context( + &self.contexts, + config.context, + &mut context_ids, + &mut contexts, + ), + actions: config + .actions + .iter() + .map(|action| CompiledLexerActionTrace { + action_index: action.action_index, + rule_index: action.rule_index, + behind: step.saturating_sub(action.position), + }) + .collect(), + }) + .collect(); self.continuations.push(CompiledLexerContinuation { - configs: configs - .iter() - .map(|config| CompiledLexerConfig { - state: config.state, - consumed_eof: config.consumed_eof, - alt_rule_index: config.alt_rule_index, - passed_non_greedy: config.passed_non_greedy, - stack: config.stack.clone(), - actions: config - .actions - .iter() - .map(|action| CompiledLexerActionTrace { - action_index: action.action_index, - rule_index: action.rule_index, - behind: step.saturating_sub(action.position), - }) - .collect(), - }) - .collect(), + contexts, + configs: compiled_configs, }); id } @@ -1000,7 +1085,7 @@ fn relative_config_key(config: &LexerConfig, step: usize) -> LexerDfaConfigKey { config.alt_rule_index, config.consumed_eof, config.passed_non_greedy, - config.stack.clone(), + config.context, config .actions .iter() @@ -1052,13 +1137,14 @@ fn build_mode( let mut build = ModeBuild::new(dfa.states.len(), dfa.continuations.len()); let start_configs = closed_configs( atn, + &mut build, vec![LexerConfig { state: start_state, position: 0, consumed_eof: false, alt_rule_index: None, passed_non_greedy: false, - stack: Vec::new(), + context: EMPTY_LEXER_CONTEXT, actions: Vec::new(), }], )?; @@ -1082,12 +1168,26 @@ fn build_mode( /// `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: &LexerAtn, moved: Vec) -> Option> { - let closure = epsilon_closure(atn, moved, &mut |_| true); +fn closed_configs( + atn: &LexerAtn, + build: &mut ModeBuild, + moved: Vec, +) -> Option> { + let closure = epsilon_closure( + atn, + moved, + &mut build.contexts, + &mut build.workspace, + &mut |_| true, + ); if closure.has_semantic_context { return None; } - if closure.configs.iter().any(has_recursive_stack) { + if closure + .configs + .iter() + .any(|config| has_recursive_context(config, &build.contexts)) + { return None; } let mut configs = closure.configs; @@ -1117,18 +1217,109 @@ fn prune_dead_action_traces(atn: &LexerAtn, config: &mut LexerConfig) { .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; +/// Copies one arena context into a continuation-local topological table. +fn compile_context( + contexts: &LexerContextArena, + context: LexerContextId, + ids: &mut FxHashMap, + compiled: &mut Vec, +) -> u32 { + if let Some(&id) = ids.get(&context) { + return id; } - stack - .iter() - .enumerate() - .any(|(index, follow)| stack[..index].contains(follow)) + + enum Frame { + Visit(LexerContextId), + Finish(LexerContextId), + } + + let mut pending = vec![Frame::Visit(context)]; + while let Some(frame) = pending.pop() { + match frame { + Frame::Visit(current) => { + if ids.contains_key(¤t) { + continue; + } + let node = contexts.node(current); + match node { + LexerContextNode::Empty => { + ids.insert(current, 0); + } + LexerContextNode::Singleton { parent, .. } => { + pending.push(Frame::Finish(current)); + pending.push(Frame::Visit(parent)); + } + LexerContextNode::Union { left, right } => { + pending.push(Frame::Finish(current)); + pending.push(Frame::Visit(right)); + pending.push(Frame::Visit(left)); + } + } + } + Frame::Finish(current) => { + let node = match contexts.node(current) { + LexerContextNode::Empty => unreachable!("empty contexts finish immediately"), + LexerContextNode::Singleton { + parent, + return_state, + } => CompiledLexerContext::Singleton { + parent: ids[&parent], + return_state, + }, + LexerContextNode::Union { left, right } => CompiledLexerContext::Union { + left: ids[&left], + right: ids[&right], + }, + }; + let id = u32::try_from(compiled.len() + 1) + .expect("compiled lexer context table overflow"); + compiled.push(node); + ids.insert(current, id); + } + } + } + + ids[&context] +} + +/// Detects a repeated return state on any caller-context path. Such recursive +/// paths grow without bound and therefore cannot be represented by a finite +/// compiled DFA. +fn has_recursive_context(config: &LexerConfig, contexts: &LexerContextArena) -> bool { + enum Frame { + Visit(LexerContextId), + LeaveRule, + } + + let mut path = Vec::with_capacity(MAX_CONTEXT_DEPTH); + let mut pending = vec![Frame::Visit(config.context)]; + while let Some(frame) = pending.pop() { + match frame { + Frame::Visit(context) => match contexts.node(context) { + LexerContextNode::Empty => {} + LexerContextNode::Union { left, right } => { + pending.push(Frame::Visit(right)); + pending.push(Frame::Visit(left)); + } + LexerContextNode::Singleton { + parent, + return_state, + } => { + if path.len() >= MAX_CONTEXT_DEPTH || path.contains(&return_state) { + return true; + } + path.push(return_state); + pending.push(Frame::LeaveRule); + pending.push(Frame::Visit(parent)); + } + }, + Frame::LeaveRule => { + path.pop().expect("leave frame must match an entered rule"); + } + } + } + + false } /// Computes every outgoing edge of one interned DFA state. @@ -1295,7 +1486,7 @@ fn move_target( moved.push(advanced); } let continuation_configs = moved.clone(); - let Some(active) = closed_configs(atn, moved) else { + let Some(active) = closed_configs(atn, build, moved) else { return EdgeTarget { state: ESCAPE_STATE, continuation: build.add_continuation(&continuation_configs, step + 1), @@ -1333,7 +1524,7 @@ fn eof_move( return EdgeTarget::DEAD; } let continuation_configs = moved.clone(); - let Some(active) = closed_configs(atn, moved) else { + let Some(active) = closed_configs(atn, build, moved) else { return EdgeTarget { state: ESCAPE_STATE, continuation: build.add_continuation(&continuation_configs, step), @@ -1786,6 +1977,71 @@ mod tests { serialized_run_offset(dfa, state) } + #[test] + fn recursive_context_detection_enforces_depth_and_cycle_backstops() { + let mut contexts = LexerContextArena::new(); + let mut bounded = EMPTY_LEXER_CONTEXT; + for return_state in 0..MAX_CONTEXT_DEPTH { + bounded = contexts.singleton(bounded, return_state); + } + let config = |context| LexerConfig { + state: 0, + position: 0, + consumed_eof: false, + alt_rule_index: Some(0), + passed_non_greedy: false, + context, + actions: Vec::new(), + }; + + assert!(!has_recursive_context(&config(bounded), &contexts)); + + let too_deep = contexts.singleton(bounded, MAX_CONTEXT_DEPTH); + assert!(has_recursive_context(&config(too_deep), &contexts)); + + let first = contexts.singleton(EMPTY_LEXER_CONTEXT, 7); + let recursive = contexts.singleton(first, 7); + assert!(has_recursive_context(&config(recursive), &contexts)); + } + + #[test] + fn deep_union_context_operations_fit_on_a_small_native_stack() { + let mut contexts = LexerContextArena::new(); + let mut workspace = PredictionWorkspace::default(); + let mut context = contexts.singleton(EMPTY_LEXER_CONTEXT, 0); + for return_state in 1..2048 { + let branch = contexts.singleton(EMPTY_LEXER_CONTEXT, return_state); + context = contexts.merge(context, branch, &mut workspace); + } + let expected_contexts = contexts.len() - 1; + let config = LexerConfig { + state: 0, + position: 0, + consumed_eof: false, + alt_rule_index: Some(0), + passed_non_greedy: false, + context, + actions: Vec::new(), + }; + + std::thread::Builder::new() + .stack_size(64 * 1024) + .spawn(move || { + assert!(!has_recursive_context(&config, &contexts)); + + let mut ids = FxHashMap::default(); + ids.insert(EMPTY_LEXER_CONTEXT, 0); + let mut compiled = Vec::new(); + let compiled_context = + compile_context(&contexts, config.context, &mut ids, &mut compiled); + assert_eq!(compiled.len(), expected_contexts); + assert_eq!(compiled_context as usize, expected_contexts); + }) + .expect("small-stack context thread should start") + .join() + .expect("deep context traversal should not overflow"); + } + #[test] fn ascii_run_classifies_and_scans_only_exact_self_loops() { let state = 7; @@ -1980,19 +2236,24 @@ mod tests { let mut lexer = BaseLexer::new(InputStream::new(" ab"), recognizer_data()); let mut store = TokenStore::new(lexer.source_text(), lexer.source_name()); let mut sink = TokenSink::new(&mut store); + let mut true_predicate_calls = 0; let id = next_token_compiled_with_hooks( &mut lexer, &mut sink, &atn, &dfa, |_, _| {}, - |_, _| true, + |_, _| { + true_predicate_calls += 1; + true_predicate_calls == 1 + }, |_, _, _| {}, ) .expect("test token should fit"); let token = sink.view(id).expect("emitted token should exist"); assert_eq!(token.token_type(), 1); assert_eq!(token.text(), "ab"); + assert_eq!(true_predicate_calls, 1); let mut compiled = BaseLexer::new(InputStream::new("ab"), recognizer_data()); let mut interpreted = BaseLexer::new(InputStream::new("ab"), recognizer_data()); @@ -2001,25 +2262,35 @@ mod tests { TokenStore::new(interpreted.source_text(), interpreted.source_name()); let mut compiled_sink = TokenSink::new(&mut compiled_store); let mut interpreted_sink = TokenSink::new(&mut interpreted_store); + let mut compiled_predicate_calls = 0; let compiled_id = next_token_compiled_with_hooks( &mut compiled, &mut compiled_sink, &atn, &dfa, |_, _| {}, - |_, _| false, + |_, _| { + compiled_predicate_calls += 1; + false + }, |_, _, _| {}, ) .expect("false predicate should recover to EOF"); + let mut interpreted_predicate_calls = 0; let interpreted_id = next_token_with_hooks( &mut interpreted, &mut interpreted_sink, &atn, |_, _| {}, - |_, _| false, + |_, _| { + interpreted_predicate_calls += 1; + false + }, |_, _, _| {}, ) .expect("interpreted false predicate should recover to EOF"); + assert_eq!(compiled_predicate_calls, interpreted_predicate_calls); + assert_eq!(compiled_predicate_calls, 1); assert_eq!( compiled_sink .view(compiled_id) diff --git a/src/lexer.rs b/src/lexer.rs index 79ce6446..0976000d 100644 --- a/src/lexer.rs +++ b/src/lexer.rs @@ -1,4 +1,4 @@ -use std::cell::RefCell; +use std::cell::{RefCell, RefMut}; use std::collections::{BTreeSet, HashMap, VecDeque}; use std::hash::BuildHasherDefault; use std::rc::Rc; @@ -6,7 +6,9 @@ use std::rc::Rc; use crate::atn::LexerAtn; use crate::char_stream::{CharStream, TextInterval}; use crate::int_stream::EOF; -use crate::prediction::PredictionFxHasher; +use crate::prediction::{ + ContextArena, ContextId, EMPTY_CONTEXT, PredictionFxHasher, PredictionWorkspace, +}; use crate::recognizer::{Recognizer, RecognizerData}; use crate::token::{ DEFAULT_CHANNEL, INVALID_TOKEN_TYPE, TokenId, TokenSink, TokenSourceError, TokenSpec, @@ -595,8 +597,9 @@ pub struct BaseLexer { /// re-simulates them instead of trusting their cached data, so the cache can /// be shared across lexer instances (and inputs) for the same ATN — see /// [`BaseLexer::with_shared_dfa`]. -#[derive(Clone, Debug, Default)] +#[derive(Debug, Default)] struct LexerDfaCache { + prediction: LexerPredictionStore, state_numbers: FxHashMap, accept_predictions: FxHashMap, /// `showDFA` edge trace. Lives with the tables it describes, so a lexer @@ -615,6 +618,147 @@ struct LexerDfaCache { mode_starts: FxHashMap, } +/// Canonical caller contexts paired with the learned lexer DFA that stores +/// their IDs. +#[derive(Debug, Default)] +pub(crate) struct LexerPredictionStore { + pub(crate) contexts: LexerContextArena, + pub(crate) workspace: PredictionWorkspace, +} + +/// Store-local identity for one ordered lexer caller-context node. +#[repr(transparent)] +#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +pub(crate) struct LexerContextId(u32); + +pub(crate) const EMPTY_LEXER_CONTEXT: LexerContextId = LexerContextId(0); + +/// One node in an ordered graph of lexer caller stacks. +/// +/// `Union` preserves ATN traversal priority. The paired unordered prediction +/// context detects when a later union adds no stack paths, which keeps cyclic +/// lexer closures finite without flattening their priority order. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) enum LexerContextNode { + Empty, + Singleton { + parent: LexerContextId, + return_state: usize, + }, + Union { + left: LexerContextId, + right: LexerContextId, + }, +} + +#[derive(Clone, Copy, Debug)] +struct LexerContextRecord { + node: LexerContextNode, + path_set: ContextId, +} + +/// Canonical ordered caller-context DAG for one learned or compiled lexer DFA. +#[derive(Debug)] +pub(crate) struct LexerContextArena { + records: Vec, + ids: FxHashMap, + path_sets: ContextArena, +} + +impl LexerContextArena { + pub(crate) fn new() -> Self { + let mut ids = FxHashMap::default(); + ids.insert(LexerContextNode::Empty, EMPTY_LEXER_CONTEXT); + Self { + records: vec![LexerContextRecord { + node: LexerContextNode::Empty, + path_set: EMPTY_CONTEXT, + }], + ids, + path_sets: ContextArena::new(), + } + } + + pub(crate) fn singleton( + &mut self, + parent: LexerContextId, + return_state: usize, + ) -> LexerContextId { + self.assert_valid(parent); + let node = LexerContextNode::Singleton { + parent, + return_state, + }; + if let Some(&context) = self.ids.get(&node) { + return context; + } + let path_set = self + .path_sets + .singleton(self.record(parent).path_set, return_state); + self.intern(node, path_set) + } + + pub(crate) fn merge( + &mut self, + left: LexerContextId, + right: LexerContextId, + workspace: &mut PredictionWorkspace, + ) -> LexerContextId { + self.assert_valid(left); + self.assert_valid(right); + if left == right { + return left; + } + let left_set = self.record(left).path_set; + let right_set = self.record(right).path_set; + let path_set = self.path_sets.merge(left_set, right_set, false, workspace); + if path_set == left_set { + return left; + } + let node = LexerContextNode::Union { left, right }; + if let Some(&context) = self.ids.get(&node) { + return context; + } + self.intern(node, path_set) + } + + pub(crate) fn node(&self, context: LexerContextId) -> LexerContextNode { + self.record(context).node + } + + #[cfg(test)] + pub(crate) const fn len(&self) -> usize { + self.records.len() + } + + fn intern(&mut self, node: LexerContextNode, path_set: ContextId) -> LexerContextId { + let context = LexerContextId( + u32::try_from(self.records.len()).expect("lexer context arena must fit in u32"), + ); + self.records.push(LexerContextRecord { node, path_set }); + self.ids.insert(node, context); + context + } + + fn record(&self, context: LexerContextId) -> &LexerContextRecord { + self.assert_valid(context); + &self.records[usize::try_from(context.0).expect("u32 lexer context ID fits in usize")] + } + + fn assert_valid(&self, context: LexerContextId) { + assert!( + usize::try_from(context.0).is_ok_and(|index| index < self.records.len()), + "lexer context ID does not belong to this store" + ); + } +} + +impl Default for LexerContextArena { + fn default() -> Self { + Self::new() + } +} + /// Dense-row width: ASCII, matching the reference runtimes' DFA edge arrays. const DENSE_EDGE_SYMBOLS: usize = 128; @@ -654,7 +798,7 @@ pub(crate) struct LexerDfaConfigKey { pub(crate) alt_rule_index: Option, pub(crate) consumed_eof: bool, pub(crate) passed_non_greedy: bool, - pub(crate) stack: Vec, + pub(crate) context: LexerContextId, pub(crate) actions: Vec, } @@ -671,7 +815,7 @@ impl LexerDfaConfigKey { alt_rule_index: Option, consumed_eof: bool, passed_non_greedy: bool, - stack: Vec, + context: LexerContextId, actions: Vec, ) -> Self { Self { @@ -679,7 +823,7 @@ impl LexerDfaConfigKey { alt_rule_index, consumed_eof, passed_non_greedy, - stack, + context, actions, } } @@ -800,7 +944,14 @@ where /// unaffected. Any path that falls back to ATN interpretation relearns its /// dynamic DFA from an empty cache after this call. pub fn clear_dfa(&self) { - *self.dfa_cache.borrow_mut() = LexerDfaCache::default(); + let mut cache = self.dfa_cache.borrow_mut(); + // In-flight predicate evaluation may clear the DFA while its configs + // still hold store-local context IDs. + let prediction = std::mem::take(&mut cache.prediction); + *cache = LexerDfaCache { + prediction, + ..LexerDfaCache::default() + }; } pub const fn input(&self) -> &I { @@ -1316,6 +1467,44 @@ where std::mem::take(self.errors.get_mut()) } + /// Borrows the canonical caller-context store paired with this lexer's + /// learned DFA. + pub(crate) fn lexer_prediction_store(&self) -> RefMut<'_, LexerPredictionStore> { + RefMut::map(self.dfa_cache.borrow_mut(), |cache| &mut cache.prediction) + } + + /// Starts a fresh token prediction while retaining bounded scratch + /// allocations for subsequent matches. + pub(crate) fn reset_lexer_prediction_workspace(&self) { + self.dfa_cache.borrow_mut().prediction.workspace.reset(); + } + + #[cfg(test)] + pub(crate) fn lexer_dfa_cache_shape(&self) -> (usize, usize, usize, usize) { + let cache = self.dfa_cache.borrow(); + let cached_states = cache.cached_states.iter().flatten().count(); + let cached_transitions = cache + .dense_edges + .iter() + .flatten() + .map(|row| { + row.iter() + .filter(|transition| transition.target_state != usize::MAX) + .count() + }) + .sum::() + + cache.sparse_edges.len(); + let max_configs = cache + .cached_states + .iter() + .flatten() + .map(|state| state.configs.len()) + .max() + .unwrap_or(0); + let contexts = cache.prediction.contexts.len(); + (cached_states, cached_transitions, max_configs, contexts) + } + /// Returns the stable state number for a normalized lexer DFA config set, /// creating one if this input path has not reached it before. pub(crate) fn lexer_dfa_state(