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

antlr4_runtime/
token_stream.rs
