Skip to main content

mago_syntax/parser/
stream.rs

1use std::fmt::Debug;
2
3use mago_allocator::prelude::*;
4
5use mago_database::file::FileId;
6use mago_database::file::HasFileId;
7use mago_span::Position;
8use mago_span::Span;
9use mago_syntax_core::parser::LookaheadBuf;
10
11use crate::cst::sequence::Sequence;
12use crate::cst::trivia::Trivia;
13use crate::cst::trivia::TriviaKind;
14use crate::error::Expected;
15use crate::error::ParseError;
16use crate::error::SyntaxError;
17use crate::lexer::Lexer;
18use crate::token::Token;
19use crate::token::TokenKind;
20
21#[derive(Debug)]
22pub struct TokenStream<'input, 'arena, A>
23where
24    'input: 'arena,
25    A: Arena,
26{
27    arena: &'arena A,
28    lexer: Lexer<'input>,
29    buffer: LookaheadBuf<Token<'input>, 4>,
30    trivia: Vec<'arena, Trivia<'input>, A>,
31    position: Position,
32    file_id: FileId,
33}
34
35impl<'input, 'arena, A> TokenStream<'input, 'arena, A>
36where
37    A: Arena,
38{
39    pub fn new(arena: &'arena A, lexer: Lexer<'input>) -> TokenStream<'input, 'arena, A> {
40        let position = lexer.current_position();
41        let file_id_cached = lexer.file_id();
42
43        TokenStream {
44            arena,
45            lexer,
46            buffer: LookaheadBuf::new(),
47            trivia: Vec::new_in(arena),
48            position,
49            file_id: file_id_cached,
50        }
51    }
52
53    /// Returns the current position of the stream within the source file.
54    ///
55    /// This position represents the end location of the most recently
56    /// consumed significant token via `advance()` or `consume()`.
57    #[inline]
58    #[must_use]
59    pub const fn current_position(&self) -> Position {
60        self.position
61    }
62
63    /// Returns whether the stream has consumed all tokens up to EOF.
64    ///
65    /// # Errors
66    ///
67    /// Returns a [`SyntaxError`] if the lexer fails to produce the next token.
68    #[inline]
69    pub fn has_reached_eof(&mut self) -> Result<bool, SyntaxError> {
70        Ok(self.fill_buffer(1)?.is_none())
71    }
72
73    /// Consumes and returns the next significant token.
74    ///
75    /// # Errors
76    ///
77    /// Returns a [`ParseError`] if EOF is reached or a lexer error occurs.
78    #[inline]
79    pub fn consume(&mut self) -> Result<Token<'input>, ParseError> {
80        match self.advance() {
81            Some(Ok(token)) => Ok(token),
82            Some(Err(error)) => Err(error.into()),
83            None => Err(self.unexpected(None, &[])),
84        }
85    }
86
87    /// Consumes the next token only if it matches the expected kind.
88    ///
89    /// Returns the token if it matches, otherwise returns an error.
90    ///
91    /// # Errors
92    ///
93    /// Returns a [`ParseError`] if the next token's kind does not match `kind`, or if EOF is reached.
94    #[inline]
95    pub fn eat(&mut self, kind: TokenKind) -> Result<Token<'input>, ParseError> {
96        // Fast path: head already buffered. Avoids the Result<Option<...>>
97        // round trip from `peek_kind` plus a follow-up `lookahead` on the
98        // happy path.
99        if let Some(token) = self.buffer.get(0) {
100            if token.kind == kind {
101                let _ = self.buffer.pop_front();
102
103                self.position = Position::new(token.start.offset + token.value.len() as u32);
104                return Ok(token);
105            }
106
107            return Err(self.unexpected_kind(Some(token), kind));
108        }
109
110        // Slow path: buffer empty, fill it.
111        let current_kind = self.peek_kind(0)?;
112        match current_kind {
113            Some(k) if k == kind => self.consume(),
114            Some(_) => match self.lookahead(0)? {
115                Some(token) => Err(self.unexpected_kind(Some(token), kind)),
116                None => Err(self.unexpected_kind(None, kind)),
117            },
118            None => Err(self.unexpected_kind(None, kind)),
119        }
120    }
121
122    /// Consumes and returns the span of the next significant token.
123    ///
124    /// This is a convenience method equivalent to `consume()?.span_for(file_id())`.
125    ///
126    /// # Errors
127    ///
128    /// Returns a [`ParseError`] if EOF is reached or a lexer error occurs.
129    #[inline]
130    pub fn consume_span(&mut self) -> Result<Span, ParseError> {
131        let file_id = self.file_id();
132        self.consume().map(|t| t.span_for(file_id))
133    }
134
135    /// Consumes the next token only if it matches the expected kind, returning its span.
136    ///
137    /// This is a convenience method equivalent to `eat(kind)?.span_for(file_id())`.
138    ///
139    /// # Errors
140    ///
141    /// Returns a [`ParseError`] if the next token's kind does not match `kind`, or if EOF is reached.
142    #[inline]
143    pub fn eat_span(&mut self, kind: TokenKind) -> Result<Span, ParseError> {
144        let file_id = self.file_id();
145        self.eat(kind).map(|t| t.span_for(file_id))
146    }
147
148    /// Advances the stream to the next token in the input source code and returns it.
149    ///
150    /// If the stream has already read the entire input source code, this method will return `None`.
151    ///
152    /// # Returns
153    ///
154    /// The next token in the input source code, or `None` if the lexer has reached the end of the input.
155    #[inline]
156    pub fn advance(&mut self) -> Option<Result<Token<'input>, SyntaxError>> {
157        match self.fill_buffer(1) {
158            Ok(Some(_)) => {
159                if let Some(token) = self.buffer.pop_front() {
160                    // Compute end position from start + value length
161                    self.position = Position::new(token.start.offset + token.value.len() as u32);
162                    Some(Ok(token))
163                } else {
164                    None
165                }
166            }
167            Ok(None) => None,
168            Err(error) => Some(Err(error)),
169        }
170    }
171
172    /// Checks if the next token matches the given kind without consuming it.
173    ///
174    /// Returns `false` if at EOF.
175    ///
176    /// # Errors
177    ///
178    /// Returns a [`ParseError`] if the lexer fails to produce the next token.
179    #[inline]
180    pub fn is_at(&mut self, kind: TokenKind) -> Result<bool, ParseError> {
181        if let Some(token) = self.buffer.get(0) {
182            return Ok(token.kind == kind);
183        }
184
185        Ok(self.peek_kind(0)? == Some(kind))
186    }
187
188    /// Peeks at the nth (0-indexed) significant token ahead without consuming it.
189    ///
190    /// Returns `Ok(None)` if EOF is reached before the nth token.
191    ///
192    /// # Errors
193    ///
194    /// Returns a [`ParseError`] if the lexer fails to produce a token while filling the lookahead buffer.
195    #[inline]
196    pub fn lookahead(&mut self, n: usize) -> Result<Option<Token<'input>>, ParseError> {
197        if n < self.buffer.len() {
198            return Ok(self.buffer.get(n));
199        }
200
201        match self.fill_buffer(n + 1) {
202            Ok(Some(_)) => Ok(self.buffer.get(n)),
203            Ok(None) => Ok(None),
204            Err(error) => Err(error.into()),
205        }
206    }
207
208    /// Peeks at the kind of the nth (0-indexed) significant token ahead.
209    ///
210    /// More efficient than `lookahead(n)?.map(|t| t.kind)` as it avoids
211    /// copying the full token when only the kind is needed.
212    ///
213    /// # Errors
214    ///
215    /// Returns a [`ParseError`] if the lexer fails to produce a token while filling the lookahead buffer.
216    #[inline]
217    pub fn peek_kind(&mut self, n: usize) -> Result<Option<TokenKind>, ParseError> {
218        if n < self.buffer.len() {
219            return Ok(self.buffer.get(n).map(|t| t.kind));
220        }
221
222        match self.fill_buffer(n + 1) {
223            Ok(Some(_)) => Ok(self.buffer.get(n).map(|t| t.kind)),
224            Ok(None) => Ok(None),
225            Err(error) => Err(error.into()),
226        }
227    }
228
229    /// Creates a `ParseError` for an unexpected token or EOF, given one or more expected kinds.
230    #[inline]
231    #[must_use]
232    pub fn unexpected(&self, found: Option<Token<'_>>, expected: &'static [TokenKind]) -> ParseError {
233        self.unexpected_with(found, Expected::OneOf(expected))
234    }
235
236    /// Creates a `ParseError` for an unexpected token or EOF when a single, runtime-known kind was expected.
237    #[inline]
238    #[must_use]
239    pub fn unexpected_kind(&self, found: Option<Token<'_>>, expected: TokenKind) -> ParseError {
240        self.unexpected_with(found, Expected::Exactly(expected))
241    }
242
243    #[inline]
244    #[must_use]
245    fn unexpected_with(&self, found: Option<Token<'_>>, expected: Expected) -> ParseError {
246        if let Some(token) = found {
247            ParseError::UnexpectedToken(expected, token.kind, token.span_for(self.file_id()))
248        } else {
249            ParseError::UnexpectedEndOfFile(expected, self.file_id(), self.current_position())
250        }
251    }
252
253    /// Consumes the comments collected by the lexer and returns them.
254    #[inline]
255    pub fn get_trivia(&mut self) -> Sequence<'arena, Trivia<'arena>> {
256        let mut trivia = Vec::new_in(self.arena);
257        std::mem::swap(&mut self.trivia, &mut trivia);
258
259        Sequence::new(trivia)
260    }
261
262    /// Fills the token buffer until at least `n` tokens are available, unless the lexer returns EOF.
263    ///
264    /// Trivia tokens are collected separately and are not stored in the main token buffer.
265    #[inline]
266    fn fill_buffer(&mut self, n: usize) -> Result<Option<usize>, SyntaxError> {
267        if self.buffer.len() >= n {
268            return Ok(Some(n));
269        }
270
271        self.fill_buffer_slow(n)
272    }
273
274    #[inline(never)]
275    fn fill_buffer_slow(&mut self, n: usize) -> Result<Option<usize>, SyntaxError> {
276        while self.buffer.len() < n {
277            match self.lexer.advance() {
278                Some(result) => {
279                    let token = result?;
280                    let trivia_kind = match token.kind {
281                        TokenKind::Whitespace => Some(TriviaKind::WhiteSpace),
282                        TokenKind::HashComment => Some(TriviaKind::HashComment),
283                        TokenKind::SingleLineComment => Some(TriviaKind::SingleLineComment),
284                        TokenKind::MultiLineComment => Some(TriviaKind::MultiLineComment),
285                        TokenKind::DocBlockComment => Some(TriviaKind::DocBlockComment),
286                        _ => None,
287                    };
288
289                    if let Some(kind) = trivia_kind {
290                        self.trivia.push(Trivia { kind, span: token.span_for(self.file_id), value: token.value });
291                        continue;
292                    }
293
294                    self.buffer.push_back(token);
295                }
296                None => return Ok(None),
297            }
298        }
299
300        Ok(Some(n))
301    }
302}
303
304impl<A> HasFileId for TokenStream<'_, '_, A>
305where
306    A: Arena,
307{
308    #[inline]
309    fn file_id(&self) -> FileId {
310        self.file_id
311    }
312}