antlr4_runtime/
token_stream.rs

1use crate::int_stream::{EOF, IntStream, UNKNOWN_SOURCE_NAME};
2use std::cell::Cell;
3
4use crate::token::{
5    DEFAULT_CHANNEL, TOKEN_EOF, Token, TokenId, TokenSink, TokenSource, TokenSourceError,
6    TokenSpec, TokenStore, TokenStoreError, TokenView,
7};
8
9#[derive(Debug)]
10struct BufferedSourceError {
11    token_index: usize,
12    error: TokenSourceError,
13}
14
15#[derive(Debug)]
16pub struct CommonTokenStream<S> {
17    source: S,
18    store: TokenStore,
19    source_token_count: usize,
20    next_visible_after: Vec<usize>,
21    cursor: usize,
22    channel: i32,
23    requested_token_count: Cell<usize>,
24    source_errors: Vec<BufferedSourceError>,
25}
26
27const UNKNOWN_NEXT_VISIBLE: usize = usize::MAX;
28
29fn buffer_token_source<S>(
30    source: &mut S,
31) -> Result<(TokenStore, Vec<BufferedSourceError>), TokenStoreError>
32where
33    S: TokenSource,
34{
35    let source_name = source.source_name().to_owned();
36    let mut store = TokenStore::new(source.source_text(), source_name);
37    let mut source_errors = Vec::new();
38    loop {
39        let expected_id = store.len();
40        let mut sink = TokenSink::new(&mut store);
41        let id = source.next_token(&mut sink)?;
42        let appended = sink.token_count().saturating_sub(expected_id);
43        if appended != 1 || id.index() != expected_id {
44            return Err(TokenStoreError::invalid_source_output(
45                expected_id,
46                id.index(),
47                appended,
48            ));
49        }
50        source_errors.extend(
51            source
52                .drain_errors()
53                .into_iter()
54                .map(|error| BufferedSourceError {
55                    token_index: id.index(),
56                    error,
57                }),
58        );
59        let token = sink
60            .view(id)
61            .expect("token source returned an ID it did not emit");
62        if token.token_type() == TOKEN_EOF {
63            break;
64        }
65    }
66    Ok((store, source_errors))
67}
68
69impl<S> CommonTokenStream<S>
70where
71    S: TokenSource,
72{
73    /// Creates and fills a token stream that filters lookahead to the default
74    /// channel.
75    ///
76    /// Use [`Self::try_new`] when token/source limit errors should be handled
77    /// instead of reported as a construction panic.
78    pub fn new(source: S) -> Self {
79        Self::try_new(source).unwrap_or_else(|error| panic!("failed to buffer tokens: {error}"))
80    }
81
82    pub fn try_new(source: S) -> Result<Self, TokenStoreError> {
83        Self::try_with_channel(source, DEFAULT_CHANNEL)
84    }
85
86    /// Creates and fills a token stream whose `LT/LA` operations see only
87    /// `channel`.
88    pub fn with_channel(source: S, channel: i32) -> Self {
89        Self::try_with_channel(source, channel)
90            .unwrap_or_else(|error| panic!("failed to buffer tokens: {error}"))
91    }
92
93    pub fn try_with_channel(mut source: S, channel: i32) -> Result<Self, TokenStoreError> {
94        let (store, source_errors) = buffer_token_source(&mut source)?;
95        let source_token_count = store.len();
96        let mut stream = Self {
97            source,
98            store,
99            source_token_count,
100            next_visible_after: vec![UNKNOWN_NEXT_VISIBLE; source_token_count],
101            cursor: 0,
102            channel,
103            requested_token_count: Cell::new(0),
104            source_errors,
105        };
106        stream.cursor = stream.adjust_seek_index(0);
107        stream.requested_token_count.set(0);
108        Ok(stream)
109    }
110
111    /// Replaces the token source and eagerly buffers it through EOF.
112    ///
113    /// The configured channel is retained; cursor, token storage, requested
114    /// lookahead, and buffered source errors are reset.
115    pub fn set_token_source(&mut self, source: S) {
116        self.try_set_token_source(source)
117            .unwrap_or_else(|error| panic!("failed to buffer tokens: {error}"));
118    }
119
120    /// Fallible form of [`Self::set_token_source`].
121    pub fn try_set_token_source(&mut self, source: S) -> Result<(), TokenStoreError> {
122        let replacement = Self::try_with_channel(source, self.channel)?;
123        *self = replacement;
124        Ok(())
125    }
126
127    /// Rebuffers the current token source after it has been reset or re-fed.
128    ///
129    /// This supports fully owned recognizer stacks: mutate the nested source
130    /// through [`Self::token_source_mut`], then call `refill` without moving the
131    /// lexer out of the stream.
132    pub fn refill(&mut self) {
133        self.try_refill()
134            .unwrap_or_else(|error| panic!("failed to buffer tokens: {error}"));
135    }
136
137    /// Fallible form of [`Self::refill`].
138    pub fn try_refill(&mut self) -> Result<(), TokenStoreError> {
139        let (store, source_errors) = buffer_token_source(&mut self.source)?;
140        self.store = store;
141        self.source_token_count = self.store.len();
142        self.next_visible_after = vec![UNKNOWN_NEXT_VISIBLE; self.source_token_count];
143        self.cursor = self.adjust_seek_index(0);
144        self.requested_token_count.set(0);
145        self.source_errors = source_errors;
146        Ok(())
147    }
148
149    /// Idempotent eager-buffering operation. Construction already buffers
150    /// through EOF so the store can be shared with CST nodes.
151    pub fn fill(&mut self) {
152        self.note_requested_count(self.source_token_count);
153        self.cursor = self.adjust_seek_index(self.cursor);
154    }
155
156    /// Returns a borrowing view of the token at an absolute buffered index.
157    pub fn get(&self, index: usize) -> Option<TokenView<'_>> {
158        self.get_id(index).and_then(|id| self.store.view(id))
159    }
160
161    /// Returns the compact ID at an absolute buffered index.
162    pub fn get_id(&self, index: usize) -> Option<TokenId> {
163        self.note_requested_count(index.saturating_add(1));
164        (index < self.source_token_count)
165            .then(|| TokenId::try_from(index).ok())
166            .flatten()
167    }
168
169    /// Returns the token at one-based lookahead/lookbehind offset, skipping
170    /// tokens outside the configured channel for positive offsets.
171    pub fn lt(&self, offset: isize) -> Option<TokenView<'_>> {
172        self.lt_id(offset).and_then(|id| self.store.view(id))
173    }
174
175    /// Returns the compact token ID at one-based lookahead/lookbehind offset.
176    pub fn lt_id(&self, offset: isize) -> Option<TokenId> {
177        if offset == 0 {
178            return None;
179        }
180        if offset < 0 {
181            return offset
182                .checked_neg()
183                .map(isize::cast_unsigned)
184                .and_then(|offset| self.lb_id(offset));
185        }
186
187        let mut index = self.next_token_on_channel(self.cursor, self.channel);
188        let mut remaining = offset;
189        while remaining > 1 {
190            index = self.next_token_on_channel(index + 1, self.channel);
191            remaining -= 1;
192        }
193        self.get_id(index)
194    }
195
196    pub fn lb(&self, offset: usize) -> Option<TokenView<'_>> {
197        self.lb_id(offset).and_then(|id| self.store.view(id))
198    }
199
200    fn lb_id(&self, offset: usize) -> Option<TokenId> {
201        if offset == 0 || self.cursor == 0 {
202            return None;
203        }
204        let mut index = self.cursor;
205        let mut remaining = offset;
206        while remaining > 0 {
207            index = self.previous_token_on_channel(index, self.channel)?;
208            remaining -= 1;
209        }
210        self.get_id(index)
211    }
212
213    pub const fn token_source(&self) -> &S {
214        &self.source
215    }
216
217    /// Returns the current source for in-place lexer re-feeding.
218    pub const fn token_source_mut(&mut self) -> &mut S {
219        &mut self.source
220    }
221
222    /// Iterates borrowing views of the original buffered token sequence.
223    pub fn tokens(&self) -> TokenIter<'_> {
224        self.note_requested_count(self.source_token_count);
225        TokenIter {
226            store: &self.store,
227            next: 0,
228            stop: self.source_token_count,
229        }
230    }
231
232    pub const fn token_count(&self) -> usize {
233        self.source_token_count
234    }
235
236    /// Returns the canonical token store owned by this stream.
237    #[must_use]
238    pub const fn token_store(&self) -> &TokenStore {
239        &self.store
240    }
241
242    /// Consumes the stream and returns its canonical token store.
243    #[must_use]
244    pub fn into_token_store(self) -> TokenStore {
245        self.store
246    }
247
248    pub(crate) fn token_view(&self, id: TokenId) -> Option<TokenView<'_>> {
249        self.store.view(id)
250    }
251
252    pub(crate) fn insert(&mut self, spec: TokenSpec) -> Result<TokenId, TokenStoreError> {
253        self.store.push(spec)
254    }
255
256    fn note_requested_count(&self, count: usize) {
257        self.requested_token_count.set(
258            self.requested_token_count
259                .get()
260                .max(count.min(self.source_token_count)),
261        );
262    }
263
264    /// Moves a raw token index to the next token visible on this stream's
265    /// channel.
266    fn adjust_seek_index(&self, index: usize) -> usize {
267        self.next_token_on_channel(index, self.channel)
268    }
269
270    /// Finds the next buffered token on `channel`.
271    fn next_token_on_channel(&self, mut index: usize, channel: i32) -> usize {
272        while let Some(id) = self.get_id(index) {
273            if self.store.token_type(id) == Some(TOKEN_EOF)
274                || self.store.channel(id) == Some(channel)
275            {
276                return index;
277            }
278            index += 1;
279        }
280        index
281    }
282
283    /// Finds the previous buffered token on `channel`.
284    fn previous_token_on_channel(&self, mut index: usize, channel: i32) -> Option<usize> {
285        while index > 0 {
286            index -= 1;
287            let id = self.get_id(index)?;
288            if self.store.token_type(id) == Some(TOKEN_EOF)
289                || self.store.channel(id) == Some(channel)
290            {
291                return Some(index);
292            }
293        }
294        None
295    }
296
297    /// Finds the previous buffered token visible to this stream before
298    /// `index`.
299    pub fn previous_visible_token_index(&self, index: usize) -> Option<usize> {
300        self.previous_token_on_channel(index, self.channel)
301    }
302}
303
304impl<S> IntStream for CommonTokenStream<S>
305where
306    S: TokenSource,
307{
308    fn consume(&mut self) {
309        if self.la(1) == EOF {
310            return;
311        }
312        let current = self.next_token_on_channel(self.cursor, self.channel);
313        self.cursor = self.adjust_seek_index(current + 1);
314    }
315
316    fn la(&mut self, offset: isize) -> i32 {
317        self.la_token(offset)
318    }
319
320    fn index(&self) -> usize {
321        self.cursor
322    }
323
324    fn seek(&mut self, index: usize) {
325        self.cursor = self.adjust_seek_index(index);
326    }
327
328    fn size(&self) -> usize {
329        self.source_token_count
330    }
331
332    fn source_name(&self) -> &str {
333        let source_name = self.source.source_name();
334        if source_name.is_empty() {
335            UNKNOWN_SOURCE_NAME
336        } else {
337            source_name
338        }
339    }
340}
341
342impl<S> CommonTokenStream<S>
343where
344    S: TokenSource,
345{
346    pub fn la_token(&self, offset: isize) -> i32 {
347        self.lt_id(offset)
348            .and_then(|id| self.store.token_type(id))
349            .unwrap_or(TOKEN_EOF)
350    }
351
352    /// Returns the token type at a buffered absolute index. Past-EOF reads are
353    /// reported as `TOKEN_EOF`.
354    pub fn token_type_at_index(&self, index: usize) -> i32 {
355        self.get_id(index)
356            .and_then(|id| self.store.token_type(id))
357            .unwrap_or(TOKEN_EOF)
358    }
359
360    /// Returns the token channel visible to `LT/LA` operations.
361    pub const fn channel(&self) -> i32 {
362        self.channel
363    }
364
365    /// Returns the next parser-visible token index after consuming the token
366    /// at `index`, skipping hidden-channel tokens.
367    pub fn next_visible_after(&mut self, index: usize) -> usize {
368        if let Some(cached) = self
369            .next_visible_after
370            .get(index)
371            .copied()
372            .filter(|cached| *cached != UNKNOWN_NEXT_VISIBLE)
373        {
374            return cached;
375        }
376
377        let mut next = index + 1;
378        let found = loop {
379            match self.get_id(next) {
380                Some(id)
381                    if self.store.token_type(id) != Some(TOKEN_EOF)
382                        && self.store.channel(id) != Some(self.channel) =>
383                {
384                    next += 1;
385                    continue;
386                }
387                _ => break next,
388            }
389        };
390        if let Some(slot) = self.next_visible_after.get_mut(index) {
391            *slot = found;
392        }
393        found
394    }
395
396    pub fn text(&self, start: usize, stop: usize) -> String {
397        if start > stop || start >= self.source_token_count {
398            return String::new();
399        }
400        (start..=stop.min(self.source_token_count.saturating_sub(1)))
401            .filter_map(|index| self.get(index))
402            .take_while(|token| token.token_type() != TOKEN_EOF)
403            .map(|token| token.text())
404            .collect()
405    }
406
407    /// Concatenated text of every buffered token except EOF.
408    pub fn text_all(&self) -> String {
409        self.tokens()
410            .filter(|token| token.token_type() != TOKEN_EOF)
411            .map(|token| token.text())
412            .collect()
413    }
414
415    /// Returns and clears diagnostics emitted while producing requested tokens.
416    pub fn drain_source_errors(&mut self) -> Vec<TokenSourceError> {
417        let requested = self.requested_token_count.get();
418        let ready = self
419            .source_errors
420            .partition_point(|buffered| buffered.token_index < requested);
421        self.source_errors
422            .drain(..ready)
423            .map(|buffered| buffered.error)
424            .collect()
425    }
426
427    pub const fn is_filled(&self) -> bool {
428        true
429    }
430}
431
432#[derive(Debug)]
433pub struct TokenIter<'a> {
434    store: &'a TokenStore,
435    next: usize,
436    stop: usize,
437}
438
439impl<'a> Iterator for TokenIter<'a> {
440    type Item = TokenView<'a>;
441
442    fn next(&mut self) -> Option<Self::Item> {
443        if self.next >= self.stop {
444            return None;
445        }
446        let id = TokenId::try_from(self.next).ok()?;
447        self.next += 1;
448        self.store.view(id)
449    }
450
451    fn size_hint(&self) -> (usize, Option<usize>) {
452        let remaining = self.stop - self.next;
453        (remaining, Some(remaining))
454    }
455}
456
457impl DoubleEndedIterator for TokenIter<'_> {
458    fn next_back(&mut self) -> Option<Self::Item> {
459        if self.next >= self.stop {
460            return None;
461        }
462        self.stop -= 1;
463        let id = TokenId::try_from(self.stop).ok()?;
464        self.store.view(id)
465    }
466}
467
468impl ExactSizeIterator for TokenIter<'_> {}
469
470#[cfg(test)]
471mod tests {
472    use super::*;
473    use crate::token::HIDDEN_CHANNEL;
474    use std::collections::VecDeque;
475
476    #[derive(Debug)]
477    struct VecTokenSource {
478        tokens: VecDeque<TokenSpec>,
479        index: usize,
480    }
481
482    impl TokenSource for VecTokenSource {
483        fn next_token(&mut self, sink: &mut TokenSink<'_>) -> Result<TokenId, TokenStoreError> {
484            let spec = self
485                .tokens
486                .pop_front()
487                .unwrap_or_else(|| TokenSpec::eof(self.index, self.index, 1, self.index));
488            self.index += 1;
489            sink.push(spec)
490        }
491
492        fn line(&self) -> usize {
493            1
494        }
495
496        fn column(&self) -> usize {
497            self.index
498        }
499
500        fn source_name(&self) -> &'static str {
501            "vec"
502        }
503    }
504
505    #[derive(Debug)]
506    struct ErrorTokenSource {
507        tokens: VecDeque<(TokenSpec, Vec<TokenSourceError>)>,
508        pending_errors: Vec<TokenSourceError>,
509        index: usize,
510    }
511
512    impl TokenSource for ErrorTokenSource {
513        fn next_token(&mut self, sink: &mut TokenSink<'_>) -> Result<TokenId, TokenStoreError> {
514            let (spec, errors) = self.tokens.pop_front().unwrap_or_else(|| {
515                (
516                    TokenSpec::eof(self.index, self.index, 1, self.index),
517                    Vec::new(),
518                )
519            });
520            self.index += 1;
521            self.pending_errors = errors;
522            sink.push(spec)
523        }
524
525        fn line(&self) -> usize {
526            1
527        }
528
529        fn column(&self) -> usize {
530            self.index
531        }
532
533        fn source_name(&self) -> &'static str {
534            "errors"
535        }
536
537        fn drain_errors(&mut self) -> Vec<TokenSourceError> {
538            std::mem::take(&mut self.pending_errors)
539        }
540    }
541
542    #[derive(Debug, Default)]
543    struct StaleIdTokenSource {
544        previous: Option<TokenId>,
545    }
546
547    impl TokenSource for StaleIdTokenSource {
548        fn next_token(&mut self, sink: &mut TokenSink<'_>) -> Result<TokenId, TokenStoreError> {
549            let emitted = sink.push(TokenSpec::explicit(1, "x"))?;
550            Ok(self.previous.replace(emitted).unwrap_or(emitted))
551        }
552
553        fn line(&self) -> usize {
554            1
555        }
556
557        fn column(&self) -> usize {
558            0
559        }
560
561        fn source_name(&self) -> &'static str {
562            "stale-id"
563        }
564    }
565
566    fn source(tokens: Vec<TokenSpec>) -> VecTokenSource {
567        VecTokenSource {
568            tokens: tokens.into(),
569            index: 0,
570        }
571    }
572
573    #[test]
574    fn stream_skips_hidden_channel_for_lookahead() {
575        let mut stream = CommonTokenStream::new(source(vec![
576            TokenSpec::explicit(1, "a"),
577            TokenSpec::explicit(2, " ").with_channel(HIDDEN_CHANNEL),
578            TokenSpec::explicit(3, "b"),
579            TokenSpec::eof(3, 3, 1, 3),
580        ]));
581        assert_eq!(stream.la_token(1), 1);
582        stream.consume();
583        assert_eq!(stream.la_token(1), 3);
584        assert_eq!(
585            stream
586                .lt(-1)
587                .expect("look-behind token should be buffered")
588                .token_type(),
589            1
590        );
591    }
592
593    #[test]
594    fn text_returns_empty_when_start_is_past_buffer() {
595        let stream = CommonTokenStream::new(source(vec![
596            TokenSpec::explicit(1, "a"),
597            TokenSpec::eof(1, 1, 1, 1),
598        ]));
599        assert_eq!(stream.text(10, 12), "");
600    }
601
602    #[test]
603    fn text_concatenates_borrowed_token_text() {
604        let stream = CommonTokenStream::new(source(vec![
605            TokenSpec::explicit(1, "a"),
606            TokenSpec::explicit(2, "b"),
607            TokenSpec::eof(2, 2, 1, 2),
608        ]));
609        assert_eq!(stream.text(0, 1), "ab");
610        assert_eq!(stream.text_all(), "ab");
611    }
612
613    #[test]
614    fn construction_rejects_stale_non_eof_token_id() {
615        let error = CommonTokenStream::try_new(StaleIdTokenSource::default())
616            .expect_err("a stale token ID must terminate buffering with an error");
617
618        assert!(error.to_string().contains("return ID 1"));
619        assert!(error.to_string().contains("returned ID 0"));
620    }
621
622    #[test]
623    fn source_errors_remain_hidden_until_their_token_is_requested() {
624        let suffix_error = TokenSourceError::new(1, 4, "token recognition error at: '@'");
625        let mut stream = CommonTokenStream::new(ErrorTokenSource {
626            tokens: [
627                (TokenSpec::explicit(1, "x"), Vec::new()),
628                (TokenSpec::explicit(2, "y"), Vec::new()),
629                (TokenSpec::eof(3, 3, 1, 3), vec![suffix_error.clone()]),
630            ]
631            .into(),
632            pending_errors: Vec::new(),
633            index: 0,
634        });
635
636        assert!(stream.drain_source_errors().is_empty());
637        assert_eq!(stream.token_type_at_index(1), 2);
638        assert!(stream.drain_source_errors().is_empty());
639
640        assert_eq!(stream.token_type_at_index(2), TOKEN_EOF);
641        assert_eq!(stream.drain_source_errors(), vec![suffix_error]);
642    }
643
644    #[test]
645    fn tokens_returns_borrowing_views() {
646        let stream = CommonTokenStream::new(source(vec![
647            TokenSpec::explicit(1, "a"),
648            TokenSpec::explicit(2, "b"),
649            TokenSpec::eof(2, 2, 1, 2),
650        ]));
651        assert_eq!(stream.tokens().len(), 3);
652        assert_eq!(
653            stream.tokens().next().map(|token| token.token_type()),
654            Some(1)
655        );
656        assert_eq!(
657            stream.tokens().next_back().map(|token| token.token_type()),
658            Some(TOKEN_EOF)
659        );
660    }
661
662    #[test]
663    fn set_token_source_replaces_buffer_and_preserves_channel() {
664        let mut stream = CommonTokenStream::with_channel(
665            source(vec![
666                TokenSpec::explicit(1, "old").with_channel(2),
667                TokenSpec::eof(3, 3, 1, 3),
668            ]),
669            2,
670        );
671
672        stream.set_token_source(source(vec![
673            TokenSpec::explicit(2, "new").with_channel(2),
674            TokenSpec::eof(3, 3, 1, 3),
675        ]));
676
677        assert_eq!(stream.channel(), 2);
678        assert_eq!(stream.index(), 0);
679        assert_eq!(stream.la_token(1), 2);
680        assert_eq!(stream.text_all(), "new");
681    }
682
683    #[test]
684    fn refill_reuses_mutated_token_source_in_place() {
685        let mut stream = CommonTokenStream::new(source(vec![
686            TokenSpec::explicit(1, "old"),
687            TokenSpec::eof(3, 3, 1, 3),
688        ]));
689        let source = stream.token_source_mut();
690        source.tokens = vec![TokenSpec::explicit(2, "new"), TokenSpec::eof(3, 3, 1, 3)].into();
691        source.index = 0;
692
693        stream.refill();
694
695        assert_eq!(stream.index(), 0);
696        assert_eq!(stream.la_token(1), 2);
697        assert_eq!(stream.text_all(), "new");
698    }
699}
antlr4_runtime/token_stream.rs

antlr4_runtime/
token_stream.rs
