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 force_interpreted: bool,
117 errors: Vec<TokenSourceError>,
118 dfa_cache: Rc<RefCell<LexerDfaCache>>,
119}
120
121#[derive(Clone, Debug, Default)]
129struct LexerDfaCache {
130 state_numbers: FxHashMap<LexerDfaKey, usize>,
131 accept_predictions: FxHashMap<usize, i32>,
132 edges: BTreeSet<LexerDfaEdge>,
136 cached_states: Vec<Option<Rc<LexerDfaCachedState>>>,
138 dense_edges: Vec<Option<Box<DenseEdgeRow>>>,
143 sparse_edges: FxHashMap<(usize, i32), LexerDfaCachedTransition>,
145 mode_starts: FxHashMap<i32, usize>,
146}
147
148const DENSE_EDGE_SYMBOLS: usize = 128;
150
151type DenseEdgeRow = [LexerDfaCachedTransition; DENSE_EDGE_SYMBOLS];
152
153const EMPTY_DENSE_EDGE: LexerDfaCachedTransition = LexerDfaCachedTransition {
156 target_state: usize::MAX,
157 position_delta: 0,
158};
159
160thread_local! {
161 static SHARED_LEXER_DFA_CACHES: RefCell<HashMap<usize, Rc<RefCell<LexerDfaCache>>>> =
164 RefCell::new(HashMap::new());
165}
166
167#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
169pub(crate) struct LexerDfaKey {
170 configs: Vec<LexerDfaConfigKey>,
171}
172
173impl LexerDfaKey {
174 pub(crate) fn new(mut configs: Vec<LexerDfaConfigKey>) -> Self {
175 configs.sort_unstable();
176 Self { configs }
177 }
178}
179
180#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
182pub(crate) struct LexerDfaConfigKey {
183 pub(crate) state: usize,
184 pub(crate) alt_rule_index: Option<usize>,
185 pub(crate) consumed_eof: bool,
186 pub(crate) passed_non_greedy: bool,
187 pub(crate) stack: Vec<usize>,
188 pub(crate) actions: Vec<LexerDfaActionKey>,
189}
190
191#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
192pub(crate) struct LexerDfaActionKey {
193 pub(crate) action_index: usize,
194 pub(crate) position_delta: usize,
195 pub(crate) rule_index: usize,
196}
197
198impl LexerDfaConfigKey {
199 pub(crate) const fn new(
200 state: usize,
201 alt_rule_index: Option<usize>,
202 consumed_eof: bool,
203 passed_non_greedy: bool,
204 stack: Vec<usize>,
205 actions: Vec<LexerDfaActionKey>,
206 ) -> Self {
207 Self {
208 state,
209 alt_rule_index,
210 consumed_eof,
211 passed_non_greedy,
212 stack,
213 actions,
214 }
215 }
216}
217
218#[derive(Clone, Copy, Debug)]
219pub(crate) struct LexerDfaCachedTransition {
220 pub(crate) target_state: usize,
221 pub(crate) position_delta: usize,
222}
223
224#[derive(Clone, Debug)]
225pub(crate) struct LexerDfaCachedAccept {
226 pub(crate) position_delta: usize,
227 pub(crate) rule_index: usize,
228 pub(crate) consumed_eof: bool,
229 pub(crate) actions: Vec<LexerDfaActionKey>,
230}
231
232#[derive(Clone, Debug)]
233pub(crate) struct LexerDfaCachedState {
234 pub(crate) has_semantic_context: bool,
235 pub(crate) configs: Vec<LexerDfaConfigKey>,
236 pub(crate) accept: Option<LexerDfaCachedAccept>,
237}
238
239#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
242struct LexerDfaEdge {
243 from: usize,
244 symbol: i32,
245 to: usize,
246}
247
248impl<I> BaseLexer<I>
249where
250 I: CharStream,
251{
252 pub fn new(input: I, data: RecognizerData) -> Self {
254 Self::with_factory(input, data, CommonTokenFactory)
255 }
256}
257
258impl<I, F> BaseLexer<I, F>
259where
260 I: CharStream,
261 F: TokenFactory,
262{
263 pub fn with_factory(input: I, data: RecognizerData, factory: F) -> Self {
265 Self {
266 input,
267 data,
268 factory,
269 mode: DEFAULT_MODE,
270 mode_stack: Vec::new(),
271 token_start: 0,
272 token_start_line: 1,
273 token_start_column: 0,
274 line: 1,
275 column: 0,
276 hit_eof: false,
277 force_interpreted: false,
278 errors: Vec::new(),
279 dfa_cache: Rc::new(RefCell::new(LexerDfaCache::default())),
280 }
281 }
282
283 #[must_use]
293 pub fn with_shared_dfa(mut self, atn: &'static Atn) -> Self {
294 let ptr: *const Atn = atn;
295 let key = ptr as usize;
296 self.dfa_cache = SHARED_LEXER_DFA_CACHES.with(|caches| {
297 Rc::clone(caches.borrow_mut().entry(key).or_insert_with(Rc::default))
298 });
299 self
300 }
301
302 pub const fn input(&self) -> &I {
303 &self.input
304 }
305
306 pub const fn input_mut(&mut self) -> &mut I {
307 &mut self.input
308 }
309
310 pub fn begin_token(&mut self) {
313 self.token_start = self.input.index();
314 self.token_start_line = self.line;
315 self.token_start_column = self.column;
316 }
317
318 pub const fn token_start(&self) -> usize {
320 self.token_start
321 }
322
323 pub const fn token_start_line(&self) -> usize {
325 self.token_start_line
326 }
327
328 pub const fn token_start_column(&self) -> usize {
330 self.token_start_column
331 }
332
333 pub fn consume_char(&mut self) {
340 let la = self.input.la(1);
341 if la == EOF {
342 return;
343 }
344 self.input.consume();
345 if char::from_u32(la.cast_unsigned()) == Some('\n') {
346 self.line += 1;
347 self.column = 0;
348 } else {
349 self.column += 1;
350 }
351 }
352
353 pub fn reset_accept_position(&mut self, index: usize) {
360 let target = index.max(self.token_start);
361 self.input.seek(self.token_start);
362 self.line = self.token_start_line;
363 self.column = self.token_start_column;
364 while self.input.index() < target && self.input.la(1) != EOF {
365 self.consume_char();
366 }
367 }
368
369 pub fn emit(&self, token_type: i32, channel: i32, text: Option<String>) -> CommonToken {
377 let stop = self.input.index().checked_sub(1).unwrap_or(usize::MAX);
378 self.emit_with_stop(token_type, channel, stop, text)
379 }
380
381 pub fn emit_with_stop(
387 &self,
388 token_type: i32,
389 channel: i32,
390 stop: usize,
391 text: Option<String>,
392 ) -> CommonToken {
393 let text = text.or_else(|| {
394 if stop == usize::MAX {
395 Some("<EOF>".to_owned())
396 } else {
397 None
398 }
399 });
400 let source_interval = if text.is_none() && stop != usize::MAX && self.token_start <= stop {
401 self.input
402 .text_source_interval(TextInterval::new(self.token_start, stop))
403 } else {
404 None
405 };
406 let source_text = source_interval
407 .as_ref()
408 .and_then(|(input, start_byte, stop_byte)| {
409 Some(crate::token::TokenSourceText {
410 input: Rc::clone(input),
411 start_byte: u32::try_from(*start_byte).ok()?,
412 stop_byte: u32::try_from(*stop_byte).ok()?,
413 })
414 });
415 let source_byte_span = source_text
416 .as_ref()
417 .map(|source_text| (source_text.start_byte, source_text.stop_byte));
418 let text = text.or_else(|| {
419 source_text
420 .is_none()
421 .then(|| self.input.text(TextInterval::new(self.token_start, stop)))
422 });
423 let mut token = self.factory.create(TokenSpec {
424 token_type,
425 channel,
426 start: self.token_start,
427 stop,
428 line: self.token_start_line,
429 column: self.token_start_column,
430 text,
431 source_text,
432 source_name: self.input.source_name(),
433 });
434 if let Some((start_byte, stop_byte)) =
435 source_byte_span.or_else(|| self.token_byte_span(stop))
436 {
437 token = token.with_byte_span(start_byte, stop_byte);
438 }
439 token
440 }
441
442 pub fn token_text(&self) -> String {
445 self.token_text_until(self.input.index())
446 }
447
448 pub fn token_text_until(&self, stop_exclusive: usize) -> String {
456 if stop_exclusive <= self.token_start {
457 return String::new();
458 }
459 self.input
460 .text(TextInterval::new(self.token_start, stop_exclusive - 1))
461 }
462
463 pub fn column_at(&self, position: usize) -> usize {
466 let mut column = self.token_start_column;
467 if position <= self.token_start {
468 return column;
469 }
470 for ch in self
471 .input
472 .text(TextInterval::new(self.token_start, position - 1))
473 .chars()
474 {
475 if ch == '\n' {
476 column = 0;
477 } else {
478 column += 1;
479 }
480 }
481 column
482 }
483
484 pub fn eof_token(&self) -> CommonToken {
486 let token = CommonToken::eof(
487 self.input.source_name(),
488 self.input.index(),
489 self.line,
490 self.column,
491 );
492 match self.eof_byte_offset() {
493 Some(byte_offset) => token.with_byte_span(byte_offset, byte_offset),
494 None => token,
495 }
496 }
497
498 fn eof_byte_offset(&self) -> Option<u32> {
499 self.byte_offset_at(self.input.index())
500 }
501
502 fn token_byte_span(&self, stop: usize) -> Option<(u32, u32)> {
503 if stop != usize::MAX && self.token_start <= stop {
504 let (_, start_byte, stop_byte) = self
505 .input
506 .text_source_interval(TextInterval::new(self.token_start, stop))?;
507 return Some((
508 u32::try_from(start_byte).ok()?,
509 u32::try_from(stop_byte).ok()?,
510 ));
511 }
512 let byte_offset = self.byte_offset_at(self.token_start)?;
513 Some((byte_offset, byte_offset))
514 }
515
516 fn byte_offset_at(&self, index: usize) -> Option<u32> {
517 let byte_offset = if index == 0 {
518 0
519 } else {
520 let previous = TextInterval::new(index - 1, index - 1);
521 self.input.text_source_interval(previous)?.2
522 };
523 u32::try_from(byte_offset).ok()
524 }
525}
526
527impl<I, F> Recognizer for BaseLexer<I, F>
528where
529 I: CharStream,
530 F: TokenFactory,
531{
532 fn data(&self) -> &RecognizerData {
533 &self.data
534 }
535
536 fn data_mut(&mut self) -> &mut RecognizerData {
537 &mut self.data
538 }
539}
540
541impl<I, F> Lexer for BaseLexer<I, F>
542where
543 I: CharStream,
544 F: TokenFactory,
545{
546 fn mode(&self) -> i32 {
547 self.mode
548 }
549
550 fn set_mode(&mut self, mode: i32) {
551 self.mode = mode;
552 }
553
554 fn push_mode(&mut self, mode: i32) {
555 self.mode_stack.push(self.mode);
556 self.mode = mode;
557 }
558
559 fn pop_mode(&mut self) -> Option<i32> {
560 let mode = self.mode_stack.pop()?;
561 self.mode = mode;
562 Some(mode)
563 }
564}
565
566impl<I, F> BaseLexer<I, F>
567where
568 I: CharStream,
569 F: TokenFactory,
570{
571 pub const fn line(&self) -> usize {
572 self.line
573 }
574
575 pub const fn column(&self) -> usize {
576 self.column
577 }
578
579 pub fn source_name(&self) -> &str {
580 self.input.source_name()
581 }
582
583 pub const fn hit_eof(&self) -> bool {
584 self.hit_eof
585 }
586
587 pub const fn set_hit_eof(&mut self, hit_eof: bool) {
588 self.hit_eof = hit_eof;
589 }
590
591 pub const fn set_force_interpreted(&mut self, force_interpreted: bool) {
599 self.force_interpreted = force_interpreted;
600 }
601
602 pub const fn force_interpreted(&self) -> bool {
604 self.force_interpreted
605 }
606
607 pub fn record_error(&mut self, line: usize, column: usize, message: impl Into<String>) {
610 self.errors
611 .push(TokenSourceError::new(line, column, message));
612 }
613
614 pub fn drain_errors(&mut self) -> Vec<TokenSourceError> {
616 std::mem::take(&mut self.errors)
617 }
618
619 pub(crate) fn lexer_dfa_state(
622 &self,
623 key: LexerDfaKey,
624 accept_prediction: Option<i32>,
625 ) -> usize {
626 let mut cache = self.dfa_cache.borrow_mut();
627 let next = cache.state_numbers.len();
628 let state = *cache.state_numbers.entry(key).or_insert(next);
629 if let Some(prediction) = accept_prediction {
630 cache.accept_predictions.insert(state, prediction);
631 }
632 state
633 }
634
635 pub fn record_lexer_dfa_edge(&self, from: usize, symbol: i32, to: usize) {
637 self.dfa_cache
638 .borrow_mut()
639 .edges
640 .insert(LexerDfaEdge { from, symbol, to });
641 }
642
643 pub(crate) fn cached_lexer_dfa_transition(
644 &self,
645 state: usize,
646 symbol: i32,
647 ) -> Option<LexerDfaCachedTransition> {
648 let cache = self.dfa_cache.borrow();
649 if let Ok(sym) = usize::try_from(symbol)
650 && sym < DENSE_EDGE_SYMBOLS
651 {
652 let transition = cache.dense_edges.get(state)?.as_ref()?[sym];
653 return (transition.target_state != usize::MAX).then_some(transition);
654 }
655 cache.sparse_edges.get(&(state, symbol)).copied()
656 }
657
658 pub(crate) fn cache_lexer_dfa_transition(
659 &self,
660 state: usize,
661 symbol: i32,
662 transition: LexerDfaCachedTransition,
663 ) {
664 let mut cache = self.dfa_cache.borrow_mut();
665 if let Ok(sym) = usize::try_from(symbol)
666 && sym < DENSE_EDGE_SYMBOLS
667 {
668 if cache.dense_edges.len() <= state {
669 cache.dense_edges.resize_with(state + 1, || None);
670 }
671 let row = cache.dense_edges[state]
672 .get_or_insert_with(|| Box::new([EMPTY_DENSE_EDGE; DENSE_EDGE_SYMBOLS]));
673 if row[sym].target_state == usize::MAX {
675 row[sym] = transition;
676 }
677 return;
678 }
679 cache.sparse_edges.entry((state, symbol)).or_insert(transition);
680 }
681
682 pub(crate) fn cached_lexer_dfa_state(&self, state: usize) -> Option<Rc<LexerDfaCachedState>> {
683 self.dfa_cache
684 .borrow()
685 .cached_states
686 .get(state)
687 .cloned()
688 .flatten()
689 }
690
691 pub(crate) fn cache_lexer_dfa_state(
692 &self,
693 state: usize,
694 cached_state: LexerDfaCachedState,
695 ) {
696 let mut cache = self.dfa_cache.borrow_mut();
697 if cache.cached_states.len() <= state {
698 cache.cached_states.resize_with(state + 1, || None);
699 }
700 cache.cached_states[state].get_or_insert_with(|| Rc::new(cached_state));
701 }
702
703 pub(crate) fn cached_lexer_mode_start(&self, mode: i32) -> Option<usize> {
704 self.dfa_cache.borrow().mode_starts.get(&mode).copied()
705 }
706
707 pub(crate) fn cache_lexer_mode_start(&self, mode: i32, state: usize) {
708 self.dfa_cache
709 .borrow_mut()
710 .mode_starts
711 .entry(mode)
712 .or_insert(state);
713 }
714
715 pub fn lexer_dfa_string(&self) -> String {
717 let mut out = String::new();
718 let cache = self.dfa_cache.borrow();
719 for edge in &cache.edges {
720 let Some(label) = lexer_dfa_edge_label(edge.symbol) else {
721 continue;
722 };
723 out.push_str(&self.lexer_dfa_state_string(edge.from));
724 out.push('-');
725 out.push_str(&label);
726 out.push_str("->");
727 out.push_str(&self.lexer_dfa_state_string(edge.to));
728 out.push('\n');
729 }
730 out
731 }
732
733 fn lexer_dfa_state_string(&self, state: usize) -> String {
734 self.dfa_cache
735 .borrow()
736 .accept_predictions
737 .get(&state)
738 .map_or_else(
739 || format!("s{state}"),
740 |prediction| format!(":s{state}=>{prediction}"),
741 )
742 }
743}
744
745fn lexer_dfa_edge_label(symbol: i32) -> Option<String> {
746 char::from_u32(symbol.cast_unsigned()).map(|ch| format!("'{ch}'"))
747}
748
749#[cfg(test)]
750mod tests {
751 use super::*;
752 use crate::char_stream::InputStream;
753 use crate::recognizer::RecognizerData;
754 use crate::token::{DEFAULT_CHANNEL, Token};
755 use crate::vocabulary::Vocabulary;
756
757 #[test]
758 fn eof_token_uses_utf8_byte_offset_after_non_ascii_input() {
759 let data = RecognizerData::new(
760 "T",
761 Vocabulary::new(
762 std::iter::empty::<Option<&str>>(),
763 std::iter::empty::<Option<&str>>(),
764 std::iter::empty::<Option<&str>>(),
765 ),
766 );
767 let mut lexer = BaseLexer::new(InputStream::new("β"), data);
768 lexer.consume_char();
769
770 let token = lexer.eof_token();
771
772 assert_eq!(token.start(), 1);
773 assert_eq!(token.stop(), 0);
774 assert_eq!(token.text(), Some("<EOF>"));
775 assert_eq!(token.byte_span(), 2..2);
776 }
777
778 #[test]
779 fn eof_rule_token_uses_utf8_byte_offset_after_non_ascii_input() {
780 let data = RecognizerData::new(
781 "T",
782 Vocabulary::new(
783 std::iter::empty::<Option<&str>>(),
784 std::iter::empty::<Option<&str>>(),
785 std::iter::empty::<Option<&str>>(),
786 ),
787 );
788 let mut lexer = BaseLexer::new(InputStream::new("β"), data);
789 lexer.consume_char();
790 lexer.begin_token();
791
792 let token = lexer.emit_with_stop(1, DEFAULT_CHANNEL, 0, Some("<EOF>".to_owned()));
793
794 assert_eq!(token.start(), 1);
795 assert_eq!(token.stop(), 0);
796 assert_eq!(token.text(), Some("<EOF>"));
797 assert_eq!(token.byte_span(), 2..2);
798 }
799
800 #[test]
801 fn emit_implicit_text_uses_utf8_byte_span_for_non_ascii_input() {
802 let data = RecognizerData::new(
803 "T",
804 Vocabulary::new(
805 std::iter::empty::<Option<&str>>(),
806 std::iter::empty::<Option<&str>>(),
807 std::iter::empty::<Option<&str>>(),
808 ),
809 );
810 let mut lexer = BaseLexer::new(InputStream::new("β"), data);
811 lexer.begin_token();
812 lexer.consume_char();
813
814 let token = lexer.emit(1, DEFAULT_CHANNEL, None);
815
816 assert_eq!(token.start(), 0);
817 assert_eq!(token.stop(), 0);
818 assert_eq!(token.text(), Some("β"));
819 assert_eq!(token.byte_span(), 0..2);
820 }
821}