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 #[inline]
58 #[must_use]
59 pub const fn current_position(&self) -> Position {
60 self.position
61 }
62
63 #[inline]
69 pub fn has_reached_eof(&mut self) -> Result<bool, SyntaxError> {
70 Ok(self.fill_buffer(1)?.is_none())
71 }
72
73 #[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 #[inline]
95 pub fn eat(&mut self, kind: TokenKind) -> Result<Token<'input>, ParseError> {
96 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 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 #[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 #[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 #[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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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}