From 8a4846a39962392ebada8d913e38cb058a37e97f Mon Sep 17 00:00:00 2001 From: Konstantin Vyatkin Date: Fri, 17 Jul 2026 10:19:09 +0200 Subject: [PATCH 1/2] Support dynamic lexer token emission --- README.md | 6 +- src/atn/lexer.rs | 280 ++++++++++++++++++++++++++++++++++-------- src/lexer.rs | 312 ++++++++++++++++++++++++++++++++++++++++++++--- src/parser.rs | 10 +- 4 files changed, 540 insertions(+), 68 deletions(-) diff --git a/README.md b/README.md index 7688f0f9..86092ffc 100644 --- a/README.md +++ b/README.md @@ -306,7 +306,11 @@ adapter (`MyParserHooks` plus `MyParserTypedHooks`) that maps stable manifest coordinates to named Rust methods. Lexer callers can use `LexerSemCtx` with `atn::lexer::next_token_with_semantic_hooks` or the compiled-DFA variant to route lexer predicates/actions through the same -`SemanticHooks` trait. +`SemanticHooks` trait. On the committed action path, `LexerSemCtx` exposes the +pending token type/channel, character lookahead and consumption, and mode +mutators. Actions can also queue a prefix token and advance the current token +start, allowing one lexer match to return multiple tokens while each +`TokenSource::next_token` call still appends exactly one token. Generated lexers also own optional hook state and emit typed lexer adapters when a semantic pattern maps lexer helper calls to hooks. The official diff --git a/src/atn/lexer.rs b/src/atn/lexer.rs index dcb22f22..94511c1b 100644 --- a/src/atn/lexer.rs +++ b/src/atn/lexer.rs @@ -13,7 +13,7 @@ use crate::lexer::{ }; use crate::parser::{SemanticHooks, UnknownSemanticPolicy}; use crate::prediction::PredictionFxHasher; -use crate::token::{DEFAULT_CHANNEL, INVALID_TOKEN_TYPE, TokenId, TokenSink, TokenStoreError}; +use crate::token::{INVALID_TOKEN_TYPE, TokenId, TokenSink, TokenStoreError}; #[allow(clippy::disallowed_types)] type FxHashSet = HashSet>; @@ -63,45 +63,26 @@ pub(super) struct ClosureResult { pub(super) has_semantic_context: bool, } -/// Mutable emission state produced by executing lexer actions for one token. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -struct LexerActionResult { - token_type: i32, - channel: i32, - skip: bool, - more: bool, -} - -impl LexerActionResult { - /// Starts action execution with the token type chosen by the accepted rule - /// and the default channel. - const fn new(token_type: i32, channel: i32) -> Self { - Self { - token_type, - channel, - skip: false, - more: false, - } - } - - /// Applies one deserialized lexer action to this token emission result and - /// to the lexer mode stack when the action changes modes. - fn apply(&mut self, action: &LexerAction, lexer: &mut BaseLexer) - where - I: CharStream, - { - match action { - LexerAction::Channel(channel) => self.channel = *channel, - LexerAction::Custom { .. } => {} - LexerAction::Mode(mode) => lexer.set_mode(*mode), - LexerAction::More => self.more = true, - LexerAction::PopMode => { - lexer.pop_mode(); - } - LexerAction::PushMode(mode) => lexer.push_mode(*mode), - LexerAction::Skip => self.skip = true, - LexerAction::Type(token_type) => self.token_type = *token_type, +/// Applies one deserialized lexer action to the shared in-progress token state. +/// +/// Keeping type and channel on [`BaseLexer`] gives portable commands and custom +/// hooks the same mutation surface, matching ANTLR's `Lexer._type` / +/// `Lexer._channel` model. +fn apply_lexer_action(action: &LexerAction, lexer: &mut BaseLexer) +where + I: CharStream, +{ + match action { + LexerAction::Channel(channel) => lexer.set_channel(*channel), + LexerAction::Custom { .. } => {} + LexerAction::Mode(mode) => lexer.set_mode(*mode), + LexerAction::More => lexer.more(), + LexerAction::PopMode => { + lexer.pop_mode(); } + LexerAction::PushMode(mode) => lexer.push_mode(*mode), + LexerAction::Skip => lexer.skip(), + LexerAction::Type(token_type) => lexer.set_type(*token_type), } } @@ -626,10 +607,14 @@ where P: FnMut(&BaseLexer, LexerPredicate) -> bool, E: FnMut(&mut BaseLexer, i32, usize), { + if let Some(token) = lexer.emit_pending_token(sink)? { + return Ok(token); + } + let mut continuing_more = false; loop { if lexer.hit_eof() { - return lexer.eof_token(sink); + return lexer.emit_eof_or_pending(sink); } if !continuing_more { @@ -645,7 +630,7 @@ where lexer.input_mut().seek(start); if lexer.input_mut().la(1) == EOF { lexer.set_hit_eof(true); - return lexer.eof_token(sink); + return lexer.emit_eof_or_pending(sink); } record_token_recognition_error(lexer, start, stop); while lexer.input().index() < stop { @@ -669,7 +654,7 @@ where .get(accept.rule_index) .copied() .unwrap_or(INVALID_TOKEN_TYPE); - let mut result = LexerActionResult::new(token_type, DEFAULT_CHANNEL); + lexer.set_type(token_type); for trace in accept.actions { if !lexer_action_belongs_to_accept(atn, accept.rule_index, trace.rule_index) { continue; @@ -685,21 +670,21 @@ where LexerCustomAction::new(*rule_index, *action_index, trace.position), ); } - other => result.apply(other, lexer), + other => apply_lexer_action(other, lexer), } } } - if result.skip { + if lexer.token_type() == crate::lexer::SKIP { continuing_more = false; continue; } - if result.more { + if lexer.token_type() == crate::lexer::MORE { continuing_more = true; continue; } - accept_adjuster(lexer, result.token_type, accept.position); + accept_adjuster(lexer, lexer.token_type(), accept.position); let emit_position = lexer.input().index(); let stop = emit_position.checked_sub(1).unwrap_or(usize::MAX); let text = if accept.consumed_eof && start == emit_position { @@ -707,7 +692,7 @@ where } else { None }; - return lexer.emit_with_stop(sink, result.token_type, result.channel, stop, text); + return lexer.emit_or_enqueue_with_stop(sink, stop, text); } } @@ -1495,9 +1480,10 @@ where mod tests { use super::*; use crate::atn::serialized::{AtnDeserializer, SerializedAtn}; + use crate::atn::{LexerAtnState, LexerTransition}; use crate::char_stream::InputStream; use crate::recognizer::RecognizerData; - use crate::token::{TOKEN_EOF, Token, TokenStore}; + use crate::token::{DEFAULT_CHANNEL, HIDDEN_CHANNEL, TOKEN_EOF, Token, TokenStore, TokenView}; use crate::vocabulary::Vocabulary; #[derive(Debug)] @@ -1508,6 +1494,65 @@ mod tests { text: String, } + fn recognizer_data() -> RecognizerData { + RecognizerData::new( + "T", + Vocabulary::new([None, Some("T")], [None, Some("T")], [None::<&str>, None]), + ) + } + + fn trailing_action_atn( + labels: &[char], + token_type: i32, + actions: Vec, + ) -> LexerAtn { + assert!(!labels.is_empty()); + let stop = labels.len() + actions.len() + 1; + let mut atn = LexerAtn::new(token_type); + for state in 0..=stop { + let kind = match state { + 0 => AtnStateKind::TokenStart, + 1 => AtnStateKind::RuleStart, + value if value == stop => AtnStateKind::RuleStop, + _ => AtnStateKind::Basic, + }; + let state = if state == 0 { + LexerAtnState::new(state, kind) + } else { + LexerAtnState::new(state, kind).with_rule_index(0) + }; + atn.add_state(state); + } + atn.state_mut(0) + .expect("token start") + .add_transition(LexerTransition::Epsilon { target: 1 }); + for (index, label) in labels.iter().enumerate() { + atn.state_mut(index + 1) + .expect("label source") + .add_transition(LexerTransition::Atom { + target: index + 2, + label: u32::from(*label).cast_signed(), + }); + } + for index in 0..actions.len() { + atn.state_mut(labels.len() + index + 1) + .expect("action source") + .add_transition(LexerTransition::Action { + target: labels.len() + index + 2, + rule_index: 0, + action_index: Some(index), + context_dependent: false, + }); + } + atn.set_rule_to_start_state(vec![1]); + atn.set_rule_to_stop_state(vec![stop]); + atn.set_rule_to_token_type(vec![token_type]); + atn.add_mode_start_state(0); + atn.add_decision_state(0); + atn.set_lexer_actions(actions); + atn + } + fn lex_one(lexer: &mut BaseLexer, atn: &LexerAtn) -> TokenSnapshot { let mut store = TokenStore::new(lexer.source_text(), lexer.source_name()); let mut sink = TokenSink::new(&mut store); @@ -1521,6 +1566,145 @@ mod tests { } } + #[derive(Debug, Default)] + struct FunctionTokenHooks { + emitted: Vec<(i32, i32, String)>, + } + + impl SemanticHooks for FunctionTokenHooks { + fn lexer_action( + &mut self, + ctx: &mut LexerSemCtx<'_, I>, + _action: LexerCustomAction, + ) -> bool + where + I: CharStream, + { + assert_eq!(ctx.text_so_far(), "count"); + while matches!(ctx.la(1), value if value == ' ' as i32 || value == '\t' as i32) { + assert!(ctx.consume()); + assert!(ctx.set_channel(HIDDEN_CHANNEL)); + } + assert_eq!(ctx.la(1), '(' as i32); + assert!(ctx.set_type(7)); + true + } + + fn lexer_token_emitted(&mut self, token: TokenView<'_>) { + self.emitted + .push((token.token_type(), token.channel(), token.text().to_owned())); + } + } + + #[test] + fn lexer_action_can_override_pending_type_and_channel() { + let atn = trailing_action_atn( + &['c', 'o', 'u', 'n', 't'], + 1, + vec![LexerAction::Custom { + rule_index: 0, + action_index: 0, + }], + ); + let mut lexer = BaseLexer::new(InputStream::new("count \t("), recognizer_data()); + let mut hooks = FunctionTokenHooks::default(); + let mut store = TokenStore::new(lexer.source_text(), lexer.source_name()); + let mut sink = TokenSink::new(&mut store); + + let id = next_token_with_semantic_hooks(&mut lexer, &mut sink, &atn, &mut hooks) + .expect("dynamic token should fit"); + let token = sink.view(id).expect("dynamic token should exist"); + assert_eq!(token.token_type(), 7); + assert_eq!(token.channel(), HIDDEN_CHANNEL); + assert_eq!(token.text(), "count \t"); + assert_eq!(lexer.la(1), '(' as i32); + assert_eq!(hooks.emitted, [(7, HIDDEN_CHANNEL, "count \t".to_owned())]); + } + + #[derive(Debug, Default)] + struct DotSplitHooks { + emitted_types: Vec, + } + + impl SemanticHooks for DotSplitHooks { + fn lexer_action( + &mut self, + ctx: &mut LexerSemCtx<'_, I>, + _action: LexerCustomAction, + ) -> bool + where + I: CharStream, + { + assert_eq!(ctx.text_so_far(), ".β"); + let dot = ctx.token_start(); + assert!(ctx.enqueue_token(1, dot)); + assert!(ctx.set_token_start(dot + 1)); + true + } + + fn lexer_token_emitted(&mut self, token: TokenView<'_>) { + self.emitted_types.push(token.token_type()); + } + } + + #[test] + fn lexer_action_can_queue_prefix_before_automatic_token() { + let atn = trailing_action_atn( + &['.', 'β'], + 3, + vec![ + LexerAction::Custom { + rule_index: 0, + action_index: 0, + }, + LexerAction::Type(2), + ], + ); + let dfa = CompiledLexerDfa::compile(&atn); + let mut lexer = BaseLexer::new(InputStream::new(".β"), recognizer_data()); + let mut hooks = DotSplitHooks::default(); + let mut store = TokenStore::new(lexer.source_text(), lexer.source_name()); + let mut sink = TokenSink::new(&mut store); + + let dot = + next_token_compiled_with_semantic_hooks(&mut lexer, &mut sink, &atn, &dfa, &mut hooks) + .expect("dot token should fit"); + assert_eq!(sink.token_count(), 1); + let identifier = + next_token_compiled_with_semantic_hooks(&mut lexer, &mut sink, &atn, &dfa, &mut hooks) + .expect("identifier token should fit"); + assert_eq!(sink.token_count(), 2); + let eof = + next_token_compiled_with_semantic_hooks(&mut lexer, &mut sink, &atn, &dfa, &mut hooks) + .expect("EOF token should fit"); + assert_eq!(sink.token_count(), 3); + + let dot = sink.view(dot).expect("dot token should exist"); + assert_eq!(dot.token_type(), 1); + assert_eq!(dot.channel(), DEFAULT_CHANNEL); + assert_eq!(dot.text(), "."); + assert_eq!(dot.start(), 0); + assert_eq!(dot.stop(), 0); + assert_eq!(dot.byte_span(), 0..1); + assert_eq!((dot.line(), dot.column()), (1, 0)); + + let identifier = sink + .view(identifier) + .expect("identifier token should exist"); + assert_eq!(identifier.token_type(), 2); + assert_eq!(identifier.text(), "β"); + assert_eq!(identifier.start(), 1); + assert_eq!(identifier.stop(), 1); + assert_eq!(identifier.byte_span(), 1..3); + assert_eq!((identifier.line(), identifier.column()), (1, 1)); + + assert_eq!( + sink.view(eof).expect("EOF token should exist").token_type(), + TOKEN_EOF + ); + assert_eq!(hooks.emitted_types, [1, 2, TOKEN_EOF]); + } + #[test] fn predicate_sensitive_lexer_state_is_not_replay_cached() { let atn = LexerAtn::new(1); diff --git a/src/lexer.rs b/src/lexer.rs index 2614efa4..58c67b82 100644 --- a/src/lexer.rs +++ b/src/lexer.rs @@ -1,5 +1,5 @@ use std::cell::RefCell; -use std::collections::{BTreeSet, HashMap}; +use std::collections::{BTreeSet, HashMap, VecDeque}; use std::hash::BuildHasherDefault; use std::rc::Rc; @@ -8,7 +8,10 @@ use crate::char_stream::{CharStream, TextInterval}; use crate::int_stream::EOF; use crate::prediction::PredictionFxHasher; use crate::recognizer::{Recognizer, RecognizerData}; -use crate::token::{TokenId, TokenSink, TokenSourceError, TokenSpec, TokenStoreError}; +use crate::token::{ + DEFAULT_CHANNEL, INVALID_TOKEN_TYPE, TokenId, TokenSink, TokenSourceError, TokenSpec, + TokenStoreError, +}; #[allow(clippy::disallowed_types)] type FxHashMap = HashMap>; @@ -96,7 +99,7 @@ impl LexerPredicate { /// Lexer reference held by [`LexerSemCtx`]. A semantic *predicate* is evaluated /// speculatively and gets a shared borrow; a *custom action* runs on the /// committed path and gets a mutable borrow so a hook can change lexer state -/// (mode stack), matching the closure-based `custom_action` API. +/// and pending token emission, matching the closure-based `custom_action` API. #[derive(Debug)] enum LexerRef<'a, I> where @@ -149,7 +152,7 @@ where } /// Builds a context with a mutable lexer borrow, for a custom-action hook - /// that may change lexer state (mode stack). See [`Self::push_mode`] etc. + /// that may change lexer and pending-token state. pub(crate) const fn new_mut( lexer: &'a mut BaseLexer, rule_index: usize, @@ -212,6 +215,133 @@ where self.lexer.get().token_text_until(self.position) } + /// Character at a one-based lookahead/lookbehind offset. + /// + /// Predicates read relative to their speculative ATN coordinate. Actions + /// read relative to the committed input cursor, including characters + /// consumed by an earlier action. + pub fn la(&mut self, offset: isize) -> i32 { + match &mut self.lexer { + LexerRef::Shared(lexer) => lexer.lookahead_at(self.position, offset), + LexerRef::Mut(lexer) => lexer.input_mut().la(offset), + } + } + + /// Absolute source index where the current token begins. + #[must_use] + pub const fn token_start(&self) -> usize { + self.lexer.get().token_start() + } + + /// Pending type of the token being matched. + #[must_use] + pub const fn token_type(&self) -> i32 { + self.lexer.get().token_type() + } + + /// Pending channel of the token being matched. + #[must_use] + pub const fn channel(&self) -> i32 { + self.lexer.get().channel() + } + + /// Sets the pending emitted token type. Action context only; see + /// [`Self::set_mode`] for the return value. + pub const fn set_type(&mut self, token_type: i32) -> bool { + match &mut self.lexer { + LexerRef::Mut(lexer) => { + lexer.set_type(token_type); + true + } + LexerRef::Shared(_) => false, + } + } + + /// Sets the pending emitted token channel. Action context only; see + /// [`Self::set_mode`] for the return value. + pub const fn set_channel(&mut self, channel: i32) -> bool { + match &mut self.lexer { + LexerRef::Mut(lexer) => { + lexer.set_channel(channel); + true + } + LexerRef::Shared(_) => false, + } + } + + /// Consumes one input character and updates source position tracking. + /// Action context only; returns whether the operation was available. + pub fn consume(&mut self) -> bool { + match &mut self.lexer { + LexerRef::Mut(lexer) => { + lexer.consume_char(); + true + } + LexerRef::Shared(_) => false, + } + } + + /// Marks the current match as skipped. Action context only. + pub const fn skip(&mut self) -> bool { + self.set_type(SKIP) + } + + /// Extends the current token with another lexer-rule match. Action context + /// only. + pub const fn more(&mut self) -> bool { + self.set_type(MORE) + } + + /// Repositions the committed accept cursor. Action context only. + pub fn reset_accept_position(&mut self, index: usize) -> bool { + match &mut self.lexer { + LexerRef::Mut(lexer) => { + lexer.reset_accept_position(index); + true + } + LexerRef::Shared(_) => false, + } + } + + /// Moves the current token start forward within the committed match. + /// + /// This is used after queueing a prefix token so automatic emission covers + /// only the remaining suffix. Returns `false` for predicate contexts or an + /// index outside the current token span. + pub fn set_token_start(&mut self, index: usize) -> bool { + match &mut self.lexer { + LexerRef::Mut(lexer) => lexer.set_token_start(index), + LexerRef::Shared(_) => false, + } + } + + /// Queues an additional token on the current channel. + /// + /// The queued token spans the current token start through `stop` + /// (inclusive) and is returned before the match's automatically emitted + /// token. Action context only. + pub fn enqueue_token(&mut self, token_type: i32, stop: usize) -> bool { + let channel = self.channel(); + self.enqueue_token_with_channel(token_type, channel, stop) + } + + /// Queues an additional token on an explicit channel. See + /// [`Self::enqueue_token`]. + pub fn enqueue_token_with_channel( + &mut self, + token_type: i32, + channel: i32, + stop: usize, + ) -> bool { + match &mut self.lexer { + LexerRef::Mut(lexer) => { + lexer.enqueue_token(token_type, channel, stop, None); + true + } + LexerRef::Shared(_) => false, + } + } + /// Sets the current lexer mode. Available only from a custom-action hook /// (the mutable-borrow context); a no-op with a warning path for the /// speculative predicate context, where mutating lexer state is invalid. @@ -265,6 +395,8 @@ pub struct BaseLexer { has_source_text: bool, mode: i32, mode_stack: Vec, + token_type: i32, + channel: i32, token_start: usize, token_start_line: usize, token_start_column: usize, @@ -274,6 +406,7 @@ pub struct BaseLexer { force_interpreted: bool, errors: RefCell>, semantic_error_coordinates: RefCell>, + pending_tokens: VecDeque, dfa_cache: Rc>, } @@ -416,6 +549,8 @@ where has_source_text, mode: DEFAULT_MODE, mode_stack: Vec::new(), + token_type: INVALID_TOKEN_TYPE, + channel: DEFAULT_CHANNEL, token_start: 0, token_start_line: 1, token_start_column: 0, @@ -425,6 +560,7 @@ where force_interpreted: false, errors: RefCell::new(Vec::new()), semantic_error_coordinates: RefCell::new(BTreeSet::new()), + pending_tokens: VecDeque::new(), dfa_cache: Rc::new(RefCell::new(LexerDfaCache::default())), } } @@ -459,6 +595,8 @@ where /// being matched. pub fn begin_token(&mut self) { self.semantic_error_coordinates.get_mut().clear(); + self.token_type = INVALID_TOKEN_TYPE; + self.channel = DEFAULT_CHANNEL; self.token_start = self.input.index(); self.token_start_line = self.line; self.token_start_column = self.column; @@ -479,6 +617,64 @@ where self.token_start_column } + /// Returns the pending type of the token being matched. + pub const fn token_type(&self) -> i32 { + self.token_type + } + + /// Overrides the pending type of the token being matched. + pub const fn set_type(&mut self, token_type: i32) { + self.token_type = token_type; + } + + /// Returns the pending channel of the token being matched. + pub const fn channel(&self) -> i32 { + self.channel + } + + /// Overrides the pending channel of the token being matched. + pub const fn set_channel(&mut self, channel: i32) { + self.channel = channel; + } + + /// Marks the current match as skipped. + pub const fn skip(&mut self) { + self.set_type(SKIP); + } + + /// Extends the current token with another lexer-rule match. + pub const fn more(&mut self) { + self.set_type(MORE); + } + + /// Reads a character at a one-based lookahead/lookbehind offset from the + /// committed input cursor without moving it. + pub fn la(&mut self, offset: isize) -> i32 { + self.input.la(offset) + } + + fn lookahead_at(&self, position: usize, offset: isize) -> i32 { + if offset == 0 { + return 0; + } + let absolute = if offset > 0 { + position.checked_add((offset - 1).cast_unsigned()) + } else { + offset + .checked_neg() + .and_then(|distance| usize::try_from(distance).ok()) + .and_then(|distance| position.checked_sub(distance)) + }; + let Some(index) = absolute.filter(|index| *index < self.input.size()) else { + return EOF; + }; + self.input + .text(TextInterval::new(index, index)) + .chars() + .next() + .map_or(EOF, |ch| u32::from(ch).cast_signed()) + } + /// Consumes one character from the input stream and updates lexer line and /// column counters. /// @@ -515,6 +711,22 @@ where } } + /// Moves the current token start forward within the consumed input span. + /// + /// Source line and column are advanced with the start, so a subsequently + /// emitted suffix token carries the same coordinates it would have had if + /// lexed independently. + pub fn set_token_start(&mut self, index: usize) -> bool { + if index < self.token_start || index > self.input.index() { + return false; + } + let (line, column) = self.position_at(index); + self.token_start = index; + self.token_start_line = line; + self.token_start_column = column; + true + } + /// Builds a token spanning from the current token start to the character /// before the input cursor. /// @@ -546,6 +758,16 @@ where stop: usize, text: Option, ) -> Result { + sink.push(self.token_spec_with_stop(token_type, channel, stop, text)) + } + + fn token_spec_with_stop( + &self, + token_type: i32, + channel: i32, + stop: usize, + text: Option, + ) -> TokenSpec { let text = text.or_else(|| { if stop == usize::MAX { Some("".to_owned()) @@ -571,7 +793,7 @@ where let (start_byte, stop_byte) = source_interval .or_else(|| self.token_byte_span(stop)) .unwrap_or((self.token_start, self.token_start)); - sink.push(TokenSpec { + TokenSpec { token_type, channel, start: self.token_start, @@ -582,7 +804,56 @@ where column: self.token_start_column, text, source_backed: source_interval.is_some(), - }) + } + } + + /// Queues an additional token to be returned before the current match's + /// automatic token. + /// + /// The token spans the current token start through `stop` (inclusive). + /// `text = None` keeps the token source-backed when the input supports it. + pub fn enqueue_token( + &mut self, + token_type: i32, + channel: i32, + stop: usize, + text: Option, + ) { + let token = self.token_spec_with_stop(token_type, channel, stop, text); + self.pending_tokens.push_back(token); + } + + pub(crate) fn emit_pending_token( + &mut self, + sink: &mut TokenSink<'_>, + ) -> Result, TokenStoreError> { + self.pending_tokens + .pop_front() + .map(|token| sink.push(token)) + .transpose() + } + + pub(crate) fn emit_or_enqueue_with_stop( + &mut self, + sink: &mut TokenSink<'_>, + stop: usize, + text: Option, + ) -> Result { + let token = self.token_spec_with_stop(self.token_type, self.channel, stop, text); + self.emit_or_enqueue(sink, token) + } + + fn emit_or_enqueue( + &mut self, + sink: &mut TokenSink<'_>, + token: TokenSpec, + ) -> Result { + if self.pending_tokens.is_empty() { + return sink.push(token); + } + self.pending_tokens.push_back(token); + self.emit_pending_token(sink)? + .ok_or_else(|| unreachable!("the pending-token queue was just populated")) } /// Returns the current token text from the token start through the input @@ -609,9 +880,14 @@ where /// Computes the zero-based source column at an absolute input position /// reached during prediction of the current token. pub fn column_at(&self, position: usize) -> usize { + self.position_at(position).1 + } + + fn position_at(&self, position: usize) -> (usize, usize) { + let mut line = self.token_start_line; let mut column = self.token_start_column; if position <= self.token_start { - return column; + return (line, column); } for ch in self .input @@ -619,23 +895,31 @@ where .chars() { if ch == '\n' { + line += 1; column = 0; } else { column += 1; } } - column + (line, column) } /// Builds the synthetic EOF token at the current input cursor. pub fn eof_token(&self, sink: &mut TokenSink<'_>) -> Result { + sink.push(self.eof_token_spec()) + } + + pub(crate) fn emit_eof_or_pending( + &mut self, + sink: &mut TokenSink<'_>, + ) -> Result { + let token = self.eof_token_spec(); + self.emit_or_enqueue(sink, token) + } + + fn eof_token_spec(&self) -> TokenSpec { let byte_offset = self.eof_byte_offset().unwrap_or_else(|| self.input.index()); - sink.push(TokenSpec::eof( - self.input.index(), - byte_offset, - self.line, - self.column, - )) + TokenSpec::eof(self.input.index(), byte_offset, self.line, self.column) } fn eof_byte_offset(&self) -> Option { diff --git a/src/parser.rs b/src/parser.rs index fd00e59f..35aad62a 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -548,11 +548,11 @@ pub trait SemanticHooks { /// the hook handled the action. /// /// The action runs post-accept, so `ctx` carries a mutable lexer borrow: a - /// hook may change lexer state — [`LexerSemCtx::push_mode`], - /// [`LexerSemCtx::pop_mode`], [`LexerSemCtx::set_mode`] — just like the - /// closure-based `custom_action` API. (The speculative predicate context in - /// [`Self::lexer_sempred`] is a shared borrow, so those mutators are inert - /// there.) + /// hook may change lexer state, including [`LexerSemCtx::set_type`], + /// [`LexerSemCtx::set_channel`], mode changes, input consumption, and + /// queued prefix tokens, just like the closure-based `custom_action` API. + /// (The speculative predicate context in [`Self::lexer_sempred`] is a shared + /// borrow, so those mutators are inert there.) fn lexer_action(&mut self, ctx: &mut LexerSemCtx<'_, I>, action: LexerCustomAction) -> bool where I: CharStream, From adc3f407502b009250bc0dc71a3d69b5fe80c276 Mon Sep 17 00:00:00 2001 From: Konstantin Vyatkin Date: Fri, 17 Jul 2026 12:38:55 +0200 Subject: [PATCH 2/2] Preserve lexer suffix after EOF rewind --- src/atn/lexer.rs | 154 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 143 insertions(+), 11 deletions(-) diff --git a/src/atn/lexer.rs b/src/atn/lexer.rs index 94511c1b..e7d19f90 100644 --- a/src/atn/lexer.rs +++ b/src/atn/lexer.rs @@ -86,6 +86,15 @@ where } } +fn refresh_hit_eof(lexer: &mut BaseLexer) -> bool +where + I: CharStream, +{ + let hit_eof = lexer.input().index() >= lexer.input().size(); + lexer.set_hit_eof(hit_eof); + hit_eof +} + /// Accumulates one epsilon-closure expansion, including whether predicate /// evaluation made the closure input-position-sensitive. struct ClosureState { @@ -645,9 +654,6 @@ where while lexer.input().index() < accept.position { lexer.consume_char(); } - if accept.consumed_eof { - lexer.set_hit_eof(true); - } let token_type = atn .rule_to_token_type() @@ -675,19 +681,24 @@ where } } - if lexer.token_type() == crate::lexer::SKIP { - continuing_more = false; - continue; - } - if lexer.token_type() == crate::lexer::MORE { - continuing_more = true; + let token_type = lexer.token_type(); + if token_type == crate::lexer::SKIP || token_type == crate::lexer::MORE { + if accept.consumed_eof || lexer.input().index() != accept.position { + refresh_hit_eof(lexer); + } + continuing_more = token_type == crate::lexer::MORE; continue; } - accept_adjuster(lexer, lexer.token_type(), accept.position); + accept_adjuster(lexer, token_type, accept.position); let emit_position = lexer.input().index(); + let hit_eof = if accept.consumed_eof || emit_position != accept.position { + refresh_hit_eof(lexer) + } else { + false + }; let stop = emit_position.checked_sub(1).unwrap_or(usize::MAX); - let text = if accept.consumed_eof && start == emit_position { + let text = if hit_eof && accept.consumed_eof && start == emit_position { Some("".to_owned()) } else { None @@ -1553,6 +1564,74 @@ mod tests { atn } + fn eof_rewind_action_atn() -> LexerAtn { + let mut atn = LexerAtn::new(2); + for (state_number, kind, rule_index) in [ + (0, AtnStateKind::TokenStart, None), + (1, AtnStateKind::RuleStart, Some(0)), + (2, AtnStateKind::Basic, Some(0)), + (3, AtnStateKind::Basic, Some(0)), + (4, AtnStateKind::Basic, Some(0)), + (5, AtnStateKind::RuleStop, Some(0)), + (6, AtnStateKind::RuleStart, Some(1)), + (7, AtnStateKind::RuleStop, Some(1)), + ] { + let mut state = LexerAtnState::new(state_number, kind); + if let Some(rule_index) = rule_index { + state = state.with_rule_index(rule_index); + } + atn.add_state(state); + } + atn.state_mut(0) + .expect("token start") + .add_transition(LexerTransition::Epsilon { target: 1 }); + atn.state_mut(0) + .expect("token start") + .add_transition(LexerTransition::Epsilon { target: 6 }); + atn.state_mut(1) + .expect("prefix rule start") + .add_transition(LexerTransition::Atom { + target: 2, + label: 'a' as i32, + }); + atn.state_mut(2) + .expect("prefix rule body") + .add_transition(LexerTransition::Atom { + target: 3, + label: 'b' as i32, + }); + atn.state_mut(3) + .expect("prefix rule EOF") + .add_transition(LexerTransition::Atom { + target: 4, + label: EOF, + }); + atn.state_mut(4) + .expect("prefix rule action") + .add_transition(LexerTransition::Action { + target: 5, + rule_index: 0, + action_index: Some(0), + context_dependent: false, + }); + atn.state_mut(6) + .expect("suffix rule start") + .add_transition(LexerTransition::Atom { + target: 7, + label: 'b' as i32, + }); + atn.set_rule_to_start_state(vec![1, 6]); + atn.set_rule_to_stop_state(vec![5, 7]); + atn.set_rule_to_token_type(vec![1, 2]); + atn.add_mode_start_state(0); + atn.add_decision_state(0); + atn.set_lexer_actions(vec![LexerAction::Custom { + rule_index: 0, + action_index: 0, + }]); + atn + } + fn lex_one(lexer: &mut BaseLexer, atn: &LexerAtn) -> TokenSnapshot { let mut store = TokenStore::new(lexer.source_text(), lexer.source_name()); let mut sink = TokenSink::new(&mut store); @@ -1705,6 +1784,59 @@ mod tests { assert_eq!(hooks.emitted_types, [1, 2, TOKEN_EOF]); } + #[derive(Debug, Default)] + struct RewindAtEofHooks { + action_count: usize, + } + + impl SemanticHooks for RewindAtEofHooks { + fn lexer_action( + &mut self, + ctx: &mut LexerSemCtx<'_, I>, + _action: LexerCustomAction, + ) -> bool + where + I: CharStream, + { + assert_eq!(ctx.la(1), EOF); + let suffix_start = ctx.token_start() + 1; + assert!(ctx.reset_accept_position(suffix_start)); + self.action_count += 1; + true + } + } + + #[test] + fn lexer_action_rewind_from_eof_preserves_suffix() { + let atn = eof_rewind_action_atn(); + let dfa = CompiledLexerDfa::compile(&atn); + let mut lexer = BaseLexer::new(InputStream::new("ab"), recognizer_data()); + let mut hooks = RewindAtEofHooks::default(); + let mut store = TokenStore::new(lexer.source_text(), lexer.source_name()); + let mut sink = TokenSink::new(&mut store); + + let prefix = + next_token_compiled_with_semantic_hooks(&mut lexer, &mut sink, &atn, &dfa, &mut hooks) + .expect("prefix token should fit"); + assert!(!lexer.hit_eof()); + let suffix = + next_token_compiled_with_semantic_hooks(&mut lexer, &mut sink, &atn, &dfa, &mut hooks) + .expect("suffix token should fit"); + let eof = + next_token_compiled_with_semantic_hooks(&mut lexer, &mut sink, &atn, &dfa, &mut hooks) + .expect("EOF token should fit"); + + let prefix = sink.view(prefix).expect("prefix token should exist"); + assert_eq!((prefix.token_type(), prefix.text()), (1, "a")); + let suffix = sink.view(suffix).expect("suffix token should exist"); + assert_eq!((suffix.token_type(), suffix.text()), (2, "b")); + assert_eq!( + sink.view(eof).expect("EOF token should exist").token_type(), + TOKEN_EOF + ); + assert_eq!(hooks.action_count, 1); + } + #[test] fn predicate_sensitive_lexer_state_is_not_replay_cached() { let atn = LexerAtn::new(1);