Skip to main content

sim_codec_python/
lexer.rs

1//! Bounded Python 3.14 tokenizer with explicit layout and trivia.
2
3use crate::{Diagnostic, DiagnosticCode as Code, Limits, Span, Token, TokenKind};
4
5const KEYWORDS: &[&str] = &[
6    "False", "None", "True", "and", "as", "assert", "async", "await", "break", "class", "continue",
7    "def", "del", "elif", "else", "except", "finally", "for", "from", "global", "if", "import",
8    "in", "is", "lambda", "nonlocal", "not", "or", "pass", "raise", "return", "try", "while",
9    "with", "yield",
10];
11const OPS: &[&str] = &[
12    "**=", ">>=", "<<=", "//=", "...", "->", ":=", "==", "!=", "<=", ">=", "<<", ">>", "**", "//",
13    "+=", "-=", "*=", "/=", "%=", "@=", "&=", "|=", "^=", "=>", "+", "-", "*", "/", "%", "@", "&",
14    "|", "^", "~", ":", ",", ";", ".", "=", "(", ")", "[", "]", "{", "}", "<", ">",
15];
16
17/// Tokenize with default resource bounds.
18pub fn tokenize(source: &str) -> Result<Vec<Token>, Diagnostic> {
19    tokenize_with_limits(source, Limits::default())
20}
21
22/// Tokenize with explicit resource bounds.
23pub fn tokenize_with_limits(source: &str, limits: Limits) -> Result<Vec<Token>, Diagnostic> {
24    Lexer::new(source, limits).run()
25}
26
27struct Lexer<'a> {
28    source: &'a str,
29    limits: Limits,
30    pos: usize,
31    line: usize,
32    column: usize,
33    at_line_start: bool,
34    bracket_depth: usize,
35    indents: Vec<(usize, String)>,
36    out: Vec<Token>,
37}
38
39impl<'a> Lexer<'a> {
40    fn new(source: &'a str, limits: Limits) -> Self {
41        Self {
42            source,
43            limits,
44            pos: 0,
45            line: 1,
46            column: 0,
47            at_line_start: true,
48            bracket_depth: 0,
49            indents: vec![(0, String::new())],
50            out: Vec::new(),
51        }
52    }
53
54    fn run(mut self) -> Result<Vec<Token>, Diagnostic> {
55        if self.source.len() > self.limits.max_bytes {
56            return Err(self.error(Code::ResourceLimit, 0, "source byte limit exceeded"));
57        }
58        while self.pos < self.source.len() {
59            if self.line > self.limits.max_lines {
60                return Err(self.error(
61                    Code::ResourceLimit,
62                    self.pos,
63                    "physical line limit exceeded",
64                ));
65            }
66            if self.at_line_start && self.bracket_depth == 0 {
67                self.layout()?;
68            }
69            if self.pos >= self.source.len() {
70                break;
71            }
72            let start = self.pos;
73            let ch = self.peek().expect("position is in source");
74            if ch == '\n' || ch == '\r' {
75                self.newline()?;
76                continue;
77            }
78            if matches!(ch, ' ' | '\t' | '\u{0c}') {
79                self.take_while(|c| matches!(c, ' ' | '\t' | '\u{0c}'));
80                self.emit(TokenKind::Trivia, start)?;
81                continue;
82            }
83            if ch == '#' {
84                self.take_while(|c| !matches!(c, '\r' | '\n'));
85                self.emit(TokenKind::Trivia, start)?;
86                continue;
87            }
88            if ch == '\\'
89                && self
90                    .rest()
91                    .get(1..)
92                    .is_some_and(|s| s.starts_with('\n') || s.starts_with("\r\n"))
93            {
94                self.bump();
95                if self.peek() == Some('\r') {
96                    self.bump();
97                }
98                self.bump();
99                self.emit(TokenKind::Trivia, start)?;
100                continue;
101            }
102            if is_name_start(ch) {
103                self.name_or_string(start)?;
104                continue;
105            }
106            if ch.is_ascii_digit()
107                || (ch == '.'
108                    && self
109                        .rest()
110                        .get(1..)
111                        .is_some_and(|s| s.starts_with(|c: char| c.is_ascii_digit())))
112            {
113                self.number(start)?;
114                continue;
115            }
116            if matches!(ch, '\'' | '"') {
117                self.string(start, "", TokenKind::String)?;
118                continue;
119            }
120            if let Some(op) = OPS.iter().find(|op| self.rest().starts_with(**op)) {
121                let op = *op;
122                for _ in op.chars() {
123                    self.bump();
124                }
125                match op {
126                    "(" | "[" | "{" => {
127                        self.bracket_depth += 1;
128                        if self.bracket_depth > self.limits.max_nesting {
129                            return Err(self.error(
130                                Code::ResourceLimit,
131                                start,
132                                "delimiter nesting limit exceeded",
133                            ));
134                        }
135                    }
136                    ")" | "]" | "}" => self.bracket_depth = self.bracket_depth.saturating_sub(1),
137                    _ => {}
138                }
139                self.emit(TokenKind::Operator, start)?;
140                continue;
141            }
142            self.bump();
143            return Err(self.error(
144                Code::InvalidCharacter,
145                start,
146                "invalid Python source character",
147            ));
148        }
149        while self.indents.len() > 1 {
150            self.indents.pop();
151            self.emit_zero(TokenKind::Dedent)?;
152        }
153        self.emit_zero(TokenKind::End)?;
154        Ok(self.out)
155    }
156
157    fn layout(&mut self) -> Result<(), Diagnostic> {
158        let start = self.pos;
159        while matches!(self.peek(), Some(' ' | '\t' | '\u{0c}')) {
160            self.bump();
161        }
162        let prefix = &self.source[start..self.pos];
163        if self.peek().is_none() || matches!(self.peek(), Some('\r' | '\n' | '#')) {
164            if self.pos > start {
165                self.emit(TokenKind::Trivia, start)?;
166            }
167            self.at_line_start = false;
168            return Ok(());
169        }
170        let width = indent_width(prefix);
171        let (current, current_prefix) = self.indents.last().expect("base indent").clone();
172        if width == current
173            && prefix != current_prefix
174            && prefix.contains('\t') != current_prefix.contains('\t')
175        {
176            return Err(self.error(
177                Code::AmbiguousIndentation,
178                start,
179                "inconsistent tabs and spaces in indentation",
180            ));
181        }
182        if self.pos > start {
183            self.emit(TokenKind::Trivia, start)?;
184        }
185        if width > current {
186            if self.indents.len() >= self.limits.max_nesting {
187                return Err(self.error(
188                    Code::ResourceLimit,
189                    start,
190                    "indentation nesting limit exceeded",
191                ));
192            }
193            self.indents.push((width, prefix.to_owned()));
194            self.emit_zero(TokenKind::Indent)?;
195        } else if width < current {
196            while self.indents.last().is_some_and(|(w, _)| *w > width) {
197                self.indents.pop();
198                self.emit_zero(TokenKind::Dedent)?;
199            }
200            if self.indents.last().map(|x| x.0) != Some(width) {
201                return Err(self.error(
202                    Code::InvalidIndentation,
203                    start,
204                    "unindent does not match an outer level",
205                ));
206            }
207        }
208        self.at_line_start = false;
209        Ok(())
210    }
211
212    fn name_or_string(&mut self, start: usize) -> Result<(), Diagnostic> {
213        self.take_while(is_name_continue);
214        let word = &self.source[start..self.pos];
215        if self.peek().is_some_and(|c| matches!(c, '\'' | '"')) && is_string_prefix(word) {
216            let kind = if word.to_ascii_lowercase().contains('f') {
217                TokenKind::FString
218            } else if word.to_ascii_lowercase().contains('t') {
219                TokenKind::TemplateString
220            } else {
221                TokenKind::String
222            };
223            self.string(start, word, kind)
224        } else {
225            self.emit(
226                if KEYWORDS.contains(&word) {
227                    TokenKind::Keyword
228                } else {
229                    TokenKind::Name
230                },
231                start,
232            )
233        }
234    }
235
236    fn string(&mut self, start: usize, _prefix: &str, kind: TokenKind) -> Result<(), Diagnostic> {
237        let quote = self.peek().expect("string quote");
238        self.bump();
239        let triple = self.rest().starts_with(quote)
240            && self
241                .rest()
242                .get(quote.len_utf8()..)
243                .is_some_and(|s| s.starts_with(quote));
244        if triple {
245            self.bump();
246            self.bump();
247        }
248        let mut braces = 0usize;
249        loop {
250            let Some(ch) = self.peek() else {
251                return Err(self.error(
252                    Code::UnterminatedLiteral,
253                    start,
254                    "unterminated string literal",
255                ));
256            };
257            if ch == '\\' {
258                self.bump();
259                if self.peek().is_some() {
260                    self.bump();
261                }
262                continue;
263            }
264            if matches!(kind, TokenKind::FString | TokenKind::TemplateString) {
265                if ch == '{' && !self.rest().starts_with("{{") {
266                    braces += 1;
267                    if braces > self.limits.max_nesting {
268                        return Err(self.error(
269                            Code::ResourceLimit,
270                            self.pos,
271                            "interpolation nesting limit exceeded",
272                        ));
273                    }
274                }
275                if ch == '}' && !self.rest().starts_with("}}") {
276                    if braces == 0 {
277                        return Err(self.error(
278                            Code::InvalidSyntax,
279                            self.pos,
280                            "single closing brace in interpolated string",
281                        ));
282                    }
283                    braces -= 1;
284                }
285            }
286            if ch == quote {
287                self.bump();
288                if !triple
289                    || (self.peek() == Some(quote)
290                        && self
291                            .rest()
292                            .get(quote.len_utf8()..)
293                            .is_some_and(|s| s.starts_with(quote)))
294                {
295                    if triple {
296                        self.bump();
297                        self.bump();
298                    }
299                    if braces != 0 {
300                        return Err(self.error(
301                            Code::UnterminatedLiteral,
302                            start,
303                            "unterminated interpolation field",
304                        ));
305                    }
306                    return self.emit(kind, start);
307                }
308                continue;
309            }
310            if !triple && matches!(ch, '\r' | '\n') {
311                return Err(self.error(
312                    Code::UnterminatedLiteral,
313                    start,
314                    "unterminated string literal",
315                ));
316            }
317            self.bump();
318        }
319    }
320
321    fn number(&mut self, start: usize) -> Result<(), Diagnostic> {
322        self.take_while(|c| c.is_ascii_alphanumeric() || matches!(c, '_' | '.'));
323        if matches!(self.peek(), Some('+' | '-'))
324            && self.source[start..self.pos].ends_with(['e', 'E'])
325        {
326            self.bump();
327            self.take_while(|c| c.is_ascii_digit() || c == '_');
328        }
329        self.emit(TokenKind::Number, start)
330    }
331
332    fn newline(&mut self) -> Result<(), Diagnostic> {
333        let start = self.pos;
334        if self.peek() == Some('\r') {
335            self.bump();
336        }
337        if self.peek() == Some('\n') {
338            self.bump();
339        }
340        self.emit(TokenKind::Newline, start)?;
341        self.line += 1;
342        self.column = 0;
343        self.at_line_start = true;
344        Ok(())
345    }
346
347    fn emit(&mut self, kind: TokenKind, start: usize) -> Result<(), Diagnostic> {
348        let (line, column) = locate(self.source, start);
349        self.push(Token {
350            kind,
351            span: Span {
352                start,
353                end: self.pos,
354            },
355            line,
356            column,
357        })
358    }
359    fn emit_zero(&mut self, kind: TokenKind) -> Result<(), Diagnostic> {
360        self.push(Token {
361            kind,
362            span: Span {
363                start: self.pos,
364                end: self.pos,
365            },
366            line: self.line,
367            column: self.column,
368        })
369    }
370    fn push(&mut self, token: Token) -> Result<(), Diagnostic> {
371        if self.out.len() >= self.limits.max_tokens {
372            return Err(self.error(
373                Code::ResourceLimit,
374                token.span.start,
375                "token limit exceeded",
376            ));
377        }
378        self.out.push(token);
379        Ok(())
380    }
381    fn rest(&self) -> &str {
382        &self.source[self.pos..]
383    }
384    fn peek(&self) -> Option<char> {
385        self.rest().chars().next()
386    }
387    fn bump(&mut self) {
388        if let Some(c) = self.peek() {
389            self.pos += c.len_utf8();
390            self.column += 1;
391        }
392    }
393    fn take_while(&mut self, test: impl Fn(char) -> bool) {
394        while self.peek().is_some_and(&test) {
395            self.bump();
396        }
397    }
398    fn error(&self, code: Code, at: usize, message: &str) -> Diagnostic {
399        let (line, column) = locate(self.source, at);
400        Diagnostic {
401            code,
402            span: Span { start: at, end: at },
403            line,
404            column,
405            message: message.to_owned(),
406        }
407    }
408}
409
410fn is_name_start(c: char) -> bool {
411    c == '_' || c.is_alphabetic()
412}
413fn is_name_continue(c: char) -> bool {
414    c == '_' || c.is_alphanumeric()
415}
416fn is_string_prefix(s: &str) -> bool {
417    matches!(
418        s.to_ascii_lowercase().as_str(),
419        "r" | "u" | "b" | "br" | "rb" | "f" | "fr" | "rf" | "t" | "tr" | "rt"
420    )
421}
422fn indent_width(s: &str) -> usize {
423    s.chars().fold(0, |n, c| match c {
424        ' ' => n + 1,
425        '\t' => (n / 8 + 1) * 8,
426        '\u{0c}' => 0,
427        _ => n,
428    })
429}
430fn locate(source: &str, at: usize) -> (usize, usize) {
431    let before = &source[..at];
432    let line = before.bytes().filter(|b| *b == b'\n').count() + 1;
433    let col = before
434        .rsplit_once('\n')
435        .map_or(before, |(_, tail)| tail)
436        .chars()
437        .count();
438    (line, col)
439}