1use std::cell::RefCell;
2use std::collections::{BTreeSet, HashMap};
3use std::hash::BuildHasherDefault;
4use std::rc::Rc;
5
6use crate::atn::Atn;
7use crate::char_stream::{CharStream, TextInterval};
8use crate::int_stream::EOF;
9use crate::prediction::PredictionFxHasher;
10use crate::recognizer::{Recognizer, RecognizerData};
11use crate::token::{CommonToken, CommonTokenFactory, TokenFactory, TokenSourceError, TokenSpec};
12
13#[allow(clippy::disallowed_types)]
14type FxHashMap<K, V> = HashMap<K, V, BuildHasherDefault<PredictionFxHasher>>;
15
16pub const SKIP: i32 = -3;
17pub const MORE: i32 = -2;
18pub const DEFAULT_MODE: i32 = 0;
19
20#[derive(Clone, Copy, Debug, Eq, PartialEq)]
21pub struct LexerMode(pub i32);
22
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
30pub struct LexerCustomAction {
31 rule_index: i32,
32 action_index: i32,
33 position: usize,
34}
35
36impl LexerCustomAction {
37 pub const fn new(rule_index: i32, action_index: i32, position: usize) -> Self {
39 Self {
40 rule_index,
41 action_index,
42 position,
43 }
44 }
45
46 pub const fn rule_index(self) -> i32 {
48 self.rule_index
49 }
50
51 pub const fn action_index(self) -> i32 {
53 self.action_index
54 }
55
56 pub const fn position(self) -> usize {
58 self.position
59 }
60}
61
62#[derive(Clone, Copy, Debug, Eq, PartialEq)]
64pub struct LexerPredicate {
65 rule_index: usize,
66 pred_index: usize,
67 position: usize,
68}
69
70impl LexerPredicate {
71 pub const fn new(rule_index: usize, pred_index: usize, position: usize) -> Self {
73 Self {
74 rule_index,
75 pred_index,
76 position,
77 }
78 }
79
80 pub const fn rule_index(self) -> usize {
82 self.rule_index
83 }
84
85 pub const fn pred_index(self) -> usize {
87 self.pred_index
88 }
89
90 pub const fn position(self) -> usize {
92 self.position
93 }
94}
95
96pub trait Lexer: Recognizer {
97 fn mode(&self) -> i32;
98 fn set_mode(&mut self, mode: i32);
99 fn push_mode(&mut self, mode: i32);
100 fn pop_mode(&mut self) -> Option<i32>;
101}
102
103#[derive(Clone, Debug)]
104pub struct BaseLexer<I, F = CommonTokenFactory> {
105 input: I,
106 data: RecognizerData,
107 factory: F,
108 mode: i32,
109 mode_stack: Vec<i32>,
110 token_start: usize,
111 token_start_line: usize,
112 token_start_column: usize,
113 line: usize,
114 column: usize,
115 hit_eof: bool,
116 errors: Vec<TokenSourceError>,
117 dfa_cache: Rc<RefCell<LexerDfaCache>>,
118}
119
120#[derive(Clone, Debug, Default)]
128struct LexerDfaCache {
129 state_numbers: FxHashMap<LexerDfaKey, usize>,
130 accept_predictions: FxHashMap<usize, i32>,
131 edges: BTreeSet<LexerDfaEdge>,
135 cached_states: Vec<Option<Rc<LexerDfaCachedState>>>,
137 dense_edges: Vec<Option<Box<DenseEdgeRow>>>,
142 sparse_edges: FxHashMap<(usize, i32), LexerDfaCachedTransition>,
144 mode_starts: FxHashMap<i32, usize>,
145}
146
147const DENSE_EDGE_SYMBOLS: usize = 128;
149
150type DenseEdgeRow = [LexerDfaCachedTransition; DENSE_EDGE_SYMBOLS];
151
152const EMPTY_DENSE_EDGE: LexerDfaCachedTransition = LexerDfaCachedTransition {
155 target_state: usize::MAX,
156 position_delta: 0,
157};
158
159thread_local! {
160 static SHARED_LEXER_DFA_CACHES: RefCell<HashMap<usize, Rc<RefCell<LexerDfaCache>>>> =
163 RefCell::new(HashMap::new());
164}
165
166#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
168pub(crate) struct LexerDfaKey {
169 configs: Vec<LexerDfaConfigKey>,
170}
171
172impl LexerDfaKey {
173 pub(crate) fn new(mut configs: Vec<LexerDfaConfigKey>) -> Self {
174 configs.sort_unstable();
175 Self { configs }
176 }
177}
178
179#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
181pub(crate) struct LexerDfaConfigKey {
182 pub(crate) state: usize,
183 pub(crate) alt_rule_index: Option<usize>,
184 pub(crate) consumed_eof: bool,
185 pub(crate) passed_non_greedy: bool,
186 pub(crate) stack: Vec<usize>,
187 pub(crate) actions: Vec<LexerDfaActionKey>,
188}
189
190#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
191pub(crate) struct LexerDfaActionKey {
192 pub(crate) action_index: usize,
193 pub(crate) position_delta: usize,
194 pub(crate) rule_index: usize,
195}
196
197impl LexerDfaConfigKey {
198 pub(crate) const fn new(
199 state: usize,
200 alt_rule_index: Option<usize>,
201 consumed_eof: bool,
202 passed_non_greedy: bool,
203 stack: Vec<usize>,
204 actions: Vec<LexerDfaActionKey>,
205 ) -> Self {
206 Self {
207 state,
208 alt_rule_index,
209 consumed_eof,
210 passed_non_greedy,
211 stack,
212 actions,
213 }
214 }
215}
216
217#[derive(Clone, Copy, Debug)]
218pub(crate) struct LexerDfaCachedTransition {
219 pub(crate) target_state: usize,
220 pub(crate) position_delta: usize,
221}
222
223#[derive(Clone, Debug)]
224pub(crate) struct LexerDfaCachedAccept {
225 pub(crate) position_delta: usize,
226 pub(crate) rule_index: usize,
227 pub(crate) consumed_eof: bool,
228 pub(crate) actions: Vec<LexerDfaActionKey>,
229}
230
231#[derive(Clone, Debug)]
232pub(crate) struct LexerDfaCachedState {
233 pub(crate) has_semantic_context: bool,
234 pub(crate) configs: Vec<LexerDfaConfigKey>,
235 pub(crate) accept: Option<LexerDfaCachedAccept>,
236}
237
238#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
241struct LexerDfaEdge {
242 from: usize,
243 symbol: i32,
244 to: usize,
245}
246
247impl<I> BaseLexer<I>
248where
249 I: CharStream,
250{
251 pub fn new(input: I, data: RecognizerData) -> Self {
253 Self::with_factory(input, data, CommonTokenFactory)
254 }
255}
256
257impl<I, F> BaseLexer<I, F>
258where
259 I: CharStream,
260 F: TokenFactory,
261{
262 pub fn with_factory(input: I, data: RecognizerData, factory: F) -> Self {
264 Self {
265 input,
266 data,
267 factory,
268 mode: DEFAULT_MODE,
269 mode_stack: Vec::new(),
270 token_start: 0,
271 token_start_line: 1,
272 token_start_column: 0,
273 line: 1,
274 column: 0,
275 hit_eof: false,
276 errors: Vec::new(),
277 dfa_cache: Rc::new(RefCell::new(LexerDfaCache::default())),
278 }
279 }
280
281 #[must_use]
291 pub fn with_shared_dfa(mut self, atn: &'static Atn) -> Self {
292 let ptr: *const Atn = atn;
293 let key = ptr as usize;
294 self.dfa_cache = SHARED_LEXER_DFA_CACHES.with(|caches| {
295 Rc::clone(caches.borrow_mut().entry(key).or_insert_with(Rc::default))
296 });
297 self
298 }
299
300 pub const fn input(&self) -> &I {
301 &self.input
302 }
303
304 pub const fn input_mut(&mut self) -> &mut I {
305 &mut self.input
306 }
307
308 pub fn begin_token(&mut self) {
311 self.token_start = self.input.index();
312 self.token_start_line = self.line;
313 self.token_start_column = self.column;
314 }
315
316 pub const fn token_start(&self) -> usize {
318 self.token_start
319 }
320
321 pub const fn token_start_line(&self) -> usize {
323 self.token_start_line
324 }
325
326 pub const fn token_start_column(&self) -> usize {
328 self.token_start_column
329 }
330
331 pub fn consume_char(&mut self) {
338 let la = self.input.la(1);
339 if la == EOF {
340 return;
341 }
342 self.input.consume();
343 if char::from_u32(la.cast_unsigned()) == Some('\n') {
344 self.line += 1;
345 self.column = 0;
346 } else {
347 self.column += 1;
348 }
349 }
350
351 pub fn reset_accept_position(&mut self, index: usize) {
358 let target = index.max(self.token_start);
359 self.input.seek(self.token_start);
360 self.line = self.token_start_line;
361 self.column = self.token_start_column;
362 while self.input.index() < target && self.input.la(1) != EOF {
363 self.consume_char();
364 }
365 }
366
367 pub fn emit(&self, token_type: i32, channel: i32, text: Option<String>) -> CommonToken {
375 let stop = self.input.index().checked_sub(1).unwrap_or(usize::MAX);
376 self.emit_with_stop(token_type, channel, stop, text)
377 }
378
379 pub fn emit_with_stop(
385 &self,
386 token_type: i32,
387 channel: i32,
388 stop: usize,
389 text: Option<String>,
390 ) -> CommonToken {
391 let text = text.or_else(|| {
392 if stop == usize::MAX {
393 Some("<EOF>".to_owned())
394 } else {
395 None
396 }
397 });
398 let source_interval = if text.is_none() && stop != usize::MAX && self.token_start <= stop {
399 self.input
400 .text_source_interval(TextInterval::new(self.token_start, stop))
401 } else {
402 None
403 };
404 let source_text = source_interval
405 .as_ref()
406 .and_then(|(input, start_byte, stop_byte)| {
407 Some(crate::token::TokenSourceText {
408 input: Rc::clone(input),
409 start_byte: u32::try_from(*start_byte).ok()?,
410 stop_byte: u32::try_from(*stop_byte).ok()?,
411 })
412 });
413 let source_byte_span = source_text
414 .as_ref()
415 .map(|source_text| (source_text.start_byte, source_text.stop_byte));
416 let text = text.or_else(|| {
417 source_text
418 .is_none()
419 .then(|| self.input.text(TextInterval::new(self.token_start, stop)))
420 });
421 let mut token = self.factory.create(TokenSpec {
422 token_type,
423 channel,
424 start: self.token_start,
425 stop,
426 line: self.token_start_line,
427 column: self.token_start_column,
428 text,
429 source_text,
430 source_name: self.input.source_name(),
431 });
432 if let Some((start_byte, stop_byte)) =
433 source_byte_span.or_else(|| self.token_byte_span(stop))
434 {
435 token = token.with_byte_span(start_byte, stop_byte);
436 }
437 token
438 }
439
440 pub fn token_text(&self) -> String {
443 self.token_text_until(self.input.index())
444 }
445
446 pub fn token_text_until(&self, stop_exclusive: usize) -> String {
454 if stop_exclusive <= self.token_start {
455 return String::new();
456 }
457 self.input
458 .text(TextInterval::new(self.token_start, stop_exclusive - 1))
459 }
460
461 pub fn column_at(&self, position: usize) -> usize {
464 let mut column = self.token_start_column;
465 if position <= self.token_start {
466 return column;
467 }
468 for ch in self
469 .input
470 .text(TextInterval::new(self.token_start, position - 1))
471 .chars()
472 {
473 if ch == '\n' {
474 column = 0;
475 } else {
476 column += 1;
477 }
478 }
479 column
480 }
481
482 pub fn eof_token(&self) -> CommonToken {
484 let token = CommonToken::eof(
485 self.input.source_name(),
486 self.input.index(),
487 self.line,
488 self.column,
489 );
490 match self.eof_byte_offset() {
491 Some(byte_offset) => token.with_byte_span(byte_offset, byte_offset),
492 None => token,
493 }
494 }
495
496 fn eof_byte_offset(&self) -> Option<u32> {
497 self.byte_offset_at(self.input.index())
498 }
499
500 fn token_byte_span(&self, stop: usize) -> Option<(u32, u32)> {
501 if stop != usize::MAX && self.token_start <= stop {
502 let (_, start_byte, stop_byte) = self
503 .input
504 .text_source_interval(TextInterval::new(self.token_start, stop))?;
505 return Some((
506 u32::try_from(start_byte).ok()?,
507 u32::try_from(stop_byte).ok()?,
508 ));
509 }
510 let byte_offset = self.byte_offset_at(self.token_start)?;
511 Some((byte_offset, byte_offset))
512 }
513
514 fn byte_offset_at(&self, index: usize) -> Option<u32> {
515 let byte_offset = if index == 0 {
516 0
517 } else {
518 let previous = TextInterval::new(index - 1, index - 1);
519 self.input.text_source_interval(previous)?.2
520 };
521 u32::try_from(byte_offset).ok()
522 }
523}
524
525impl<I, F> Recognizer for BaseLexer<I, F>
526where
527 I: CharStream,
528 F: TokenFactory,
529{
530 fn data(&self) -> &RecognizerData {
531 &self.data
532 }
533
534 fn data_mut(&mut self) -> &mut RecognizerData {
535 &mut self.data
536 }
537}
538
539impl<I, F> Lexer for BaseLexer<I, F>
540where
541 I: CharStream,
542 F: TokenFactory,
543{
544 fn mode(&self) -> i32 {
545 self.mode
546 }
547
548 fn set_mode(&mut self, mode: i32) {
549 self.mode = mode;
550 }
551
552 fn push_mode(&mut self, mode: i32) {
553 self.mode_stack.push(self.mode);
554 self.mode = mode;
555 }
556
557 fn pop_mode(&mut self) -> Option<i32> {
558 let mode = self.mode_stack.pop()?;
559 self.mode = mode;
560 Some(mode)
561 }
562}
563
564impl<I, F> BaseLexer<I, F>
565where
566 I: CharStream,
567 F: TokenFactory,
568{
569 pub const fn line(&self) -> usize {
570 self.line
571 }
572
573 pub const fn column(&self) -> usize {
574 self.column
575 }
576
577 pub fn source_name(&self) -> &str {
578 self.input.source_name()
579 }
580
581 pub const fn hit_eof(&self) -> bool {
582 self.hit_eof
583 }
584
585 pub const fn set_hit_eof(&mut self, hit_eof: bool) {
586 self.hit_eof = hit_eof;
587 }
588
589 pub fn record_error(&mut self, line: usize, column: usize, message: impl Into<String>) {
592 self.errors
593 .push(TokenSourceError::new(line, column, message));
594 }
595
596 pub fn drain_errors(&mut self) -> Vec<TokenSourceError> {
598 std::mem::take(&mut self.errors)
599 }
600
601 pub(crate) fn lexer_dfa_state(
604 &self,
605 key: LexerDfaKey,
606 accept_prediction: Option<i32>,
607 ) -> usize {
608 let mut cache = self.dfa_cache.borrow_mut();
609 let next = cache.state_numbers.len();
610 let state = *cache.state_numbers.entry(key).or_insert(next);
611 if let Some(prediction) = accept_prediction {
612 cache.accept_predictions.insert(state, prediction);
613 }
614 state
615 }
616
617 pub fn record_lexer_dfa_edge(&self, from: usize, symbol: i32, to: usize) {
619 self.dfa_cache
620 .borrow_mut()
621 .edges
622 .insert(LexerDfaEdge { from, symbol, to });
623 }
624
625 pub(crate) fn cached_lexer_dfa_transition(
626 &self,
627 state: usize,
628 symbol: i32,
629 ) -> Option<LexerDfaCachedTransition> {
630 let cache = self.dfa_cache.borrow();
631 if let Ok(sym) = usize::try_from(symbol)
632 && sym < DENSE_EDGE_SYMBOLS
633 {
634 let transition = cache.dense_edges.get(state)?.as_ref()?[sym];
635 return (transition.target_state != usize::MAX).then_some(transition);
636 }
637 cache.sparse_edges.get(&(state, symbol)).copied()
638 }
639
640 pub(crate) fn cache_lexer_dfa_transition(
641 &self,
642 state: usize,
643 symbol: i32,
644 transition: LexerDfaCachedTransition,
645 ) {
646 let mut cache = self.dfa_cache.borrow_mut();
647 if let Ok(sym) = usize::try_from(symbol)
648 && sym < DENSE_EDGE_SYMBOLS
649 {
650 if cache.dense_edges.len() <= state {
651 cache.dense_edges.resize_with(state + 1, || None);
652 }
653 let row = cache.dense_edges[state]
654 .get_or_insert_with(|| Box::new([EMPTY_DENSE_EDGE; DENSE_EDGE_SYMBOLS]));
655 if row[sym].target_state == usize::MAX {
657 row[sym] = transition;
658 }
659 return;
660 }
661 cache.sparse_edges.entry((state, symbol)).or_insert(transition);
662 }
663
664 pub(crate) fn cached_lexer_dfa_state(&self, state: usize) -> Option<Rc<LexerDfaCachedState>> {
665 self.dfa_cache
666 .borrow()
667 .cached_states
668 .get(state)
669 .cloned()
670 .flatten()
671 }
672
673 pub(crate) fn cache_lexer_dfa_state(
674 &self,
675 state: usize,
676 cached_state: LexerDfaCachedState,
677 ) {
678 let mut cache = self.dfa_cache.borrow_mut();
679 if cache.cached_states.len() <= state {
680 cache.cached_states.resize_with(state + 1, || None);
681 }
682 cache.cached_states[state].get_or_insert_with(|| Rc::new(cached_state));
683 }
684
685 pub(crate) fn cached_lexer_mode_start(&self, mode: i32) -> Option<usize> {
686 self.dfa_cache.borrow().mode_starts.get(&mode).copied()
687 }
688
689 pub(crate) fn cache_lexer_mode_start(&self, mode: i32, state: usize) {
690 self.dfa_cache
691 .borrow_mut()
692 .mode_starts
693 .entry(mode)
694 .or_insert(state);
695 }
696
697 pub fn lexer_dfa_string(&self) -> String {
699 let mut out = String::new();
700 let cache = self.dfa_cache.borrow();
701 for edge in &cache.edges {
702 let Some(label) = lexer_dfa_edge_label(edge.symbol) else {
703 continue;
704 };
705 out.push_str(&self.lexer_dfa_state_string(edge.from));
706 out.push('-');
707 out.push_str(&label);
708 out.push_str("->");
709 out.push_str(&self.lexer_dfa_state_string(edge.to));
710 out.push('\n');
711 }
712 out
713 }
714
715 fn lexer_dfa_state_string(&self, state: usize) -> String {
716 self.dfa_cache
717 .borrow()
718 .accept_predictions
719 .get(&state)
720 .map_or_else(
721 || format!("s{state}"),
722 |prediction| format!(":s{state}=>{prediction}"),
723 )
724 }
725}
726
727fn lexer_dfa_edge_label(symbol: i32) -> Option<String> {
728 char::from_u32(symbol.cast_unsigned()).map(|ch| format!("'{ch}'"))
729}
730
731#[cfg(test)]
732mod tests {
733 use super::*;
734 use crate::char_stream::InputStream;
735 use crate::recognizer::RecognizerData;
736 use crate::token::{DEFAULT_CHANNEL, Token};
737 use crate::vocabulary::Vocabulary;
738
739 #[test]
740 fn eof_token_uses_utf8_byte_offset_after_non_ascii_input() {
741 let data = RecognizerData::new(
742 "T",
743 Vocabulary::new(
744 std::iter::empty::<Option<&str>>(),
745 std::iter::empty::<Option<&str>>(),
746 std::iter::empty::<Option<&str>>(),
747 ),
748 );
749 let mut lexer = BaseLexer::new(InputStream::new("β"), data);
750 lexer.consume_char();
751
752 let token = lexer.eof_token();
753
754 assert_eq!(token.start(), 1);
755 assert_eq!(token.stop(), 0);
756 assert_eq!(token.text(), Some("<EOF>"));
757 assert_eq!(token.byte_span(), 2..2);
758 }
759
760 #[test]
761 fn eof_rule_token_uses_utf8_byte_offset_after_non_ascii_input() {
762 let data = RecognizerData::new(
763 "T",
764 Vocabulary::new(
765 std::iter::empty::<Option<&str>>(),
766 std::iter::empty::<Option<&str>>(),
767 std::iter::empty::<Option<&str>>(),
768 ),
769 );
770 let mut lexer = BaseLexer::new(InputStream::new("β"), data);
771 lexer.consume_char();
772 lexer.begin_token();
773
774 let token = lexer.emit_with_stop(1, DEFAULT_CHANNEL, 0, Some("<EOF>".to_owned()));
775
776 assert_eq!(token.start(), 1);
777 assert_eq!(token.stop(), 0);
778 assert_eq!(token.text(), Some("<EOF>"));
779 assert_eq!(token.byte_span(), 2..2);
780 }
781
782 #[test]
783 fn emit_implicit_text_uses_utf8_byte_span_for_non_ascii_input() {
784 let data = RecognizerData::new(
785 "T",
786 Vocabulary::new(
787 std::iter::empty::<Option<&str>>(),
788 std::iter::empty::<Option<&str>>(),
789 std::iter::empty::<Option<&str>>(),
790 ),
791 );
792 let mut lexer = BaseLexer::new(InputStream::new("β"), data);
793 lexer.begin_token();
794 lexer.consume_char();
795
796 let token = lexer.emit(1, DEFAULT_CHANNEL, None);
797
798 assert_eq!(token.start(), 0);
799 assert_eq!(token.stop(), 0);
800 assert_eq!(token.text(), Some("β"));
801 assert_eq!(token.byte_span(), 0..2);
802 }
803}