aster/aster-core/src/lexer/lexer.rs

// src/lexer/lexer.rs

use super::{Token, TokenKind};
use crate::error::RuntimeError;

pub struct Lexer {
    src: Vec<char>,
    current: usize,
    line: usize,
    column: usize,
}

impl Lexer {
    pub fn new(input: &str) -> Self {
        Self {
            src: input.chars().collect(),
            current: 0,
            line: 1,
            column: 1,
        }
    }

    pub fn tokenize(mut self) -> (Vec<Token>, Vec<RuntimeError>) {
        let mut tokens = Vec::new();
        let mut errors = Vec::new();

        while !self.is_at_end() {
            if let Some(token) = self.next_token(&mut errors) {
                tokens.push(token);
            }
        }

        tokens.push(Token {
            kind: TokenKind::EOF,
            line: self.line,
            column: self.column,
        });

        (tokens, errors)
    }

    fn next_token(&mut self, errors: &mut Vec<RuntimeError>) -> Option<Token> {
        self.skip_whitespace();

        if self.is_at_end() {
            return None;
        }

        let line = self.line;
        let column = self.column;
        let c = self.advance();

        let kind = match c {
            '(' => TokenKind::LeftParen,
            ')' => TokenKind::RightParen,
            '{' => TokenKind::LeftBrace,
            '}' => TokenKind::RightBrace,
            ',' => TokenKind::Comma,
            ';' => TokenKind::Semicolon,

            '+' => {
                if self.match_char('=') {
                    TokenKind::PlusEqual
                } else {
                    TokenKind::Plus
                }
            }
            '-' => {
                if self.match_char('=') {
                    TokenKind::MinusEqual
                } else {
                    TokenKind::Minus
                }
            }
            '*' => {
                if self.match_char('=') {
                    TokenKind::StarEqual
                } else {
                    TokenKind::Star
                }
            }
            '.' => TokenKind::Dot,
            ':' => TokenKind::Colon,
            '?' => TokenKind::Question,
            '[' => TokenKind::LeftBracket,
            ']' => TokenKind::RightBracket,

            '/' => {
                if self.match_char('/') {
                    // 单行注释
                    while !self.is_at_end() && self.peek() != '\n' {
                        self.advance();
                    }
                    return None;
                } else if self.match_char('*') {
                    // 多行注释 /* ... */
                    return self.block_comment(line, column, errors);
                } else if self.match_char('=') {
                    TokenKind::SlashEqual
                } else {
                    TokenKind::Slash
                }
            }

            '%' => {
                if self.match_char('=') {
                    TokenKind::PercentEqual
                } else {
                    TokenKind::Percent
                }
            }

            '!' => {
                if self.match_char('=') {
                    TokenKind::BangEqual
                } else {
                    TokenKind::Bang
                }
            }

            '=' => {
                if self.match_char('=') {
                    TokenKind::EqualEqual
                } else {
                    TokenKind::Equal
                }
            }

            '>' => {
                if self.match_char('=') {
                    TokenKind::GreaterEqual
                } else {
                    TokenKind::Greater
                }
            }

            '<' => {
                if self.match_char('=') {
                    TokenKind::LessEqual
                } else {
                    TokenKind::Less
                }
            }

            '&' => {
                if self.match_char('&') {
                    TokenKind::AndAnd
                } else {
                    errors.push(RuntimeError::lex("Unexpected '&'. Did you mean '&&'?", line, column));
                    return None;
                }
            }

            '|' => {
                if self.match_char('|') {
                    TokenKind::OrOr
                } else {
                    errors.push(RuntimeError::lex("Unexpected '|'. Did you mean '||'?", line, column));
                    return None;
                }
            }

            '"' | '\'' => return self.string_literal(line, column, errors, c),

            c if c.is_ascii_digit() => {
                return Some(self.number_literal(c, line, column));
            }

            c if is_ident_start(c) => {
                return Some(self.identifier(c, line, column));
            }

            _ => {
                errors.push(RuntimeError::lex(
                    format!("Unexpected character '{}'", c),
                    line,
                    column,
                ));
                return None;
            }
        };

        Some(Token { kind, line, column })
    }

    // ---------------- helpers ----------------

    fn skip_whitespace(&mut self) {
        loop {
            if self.is_at_end() {
                return;
            }

            match self.peek() {
                ' ' | '\t' | '\r' => {
                    self.advance();
                }
                '\n' => {
                    self.advance();
                    self.line += 1;
                    self.column = 1;
                }
                _ => return,
            }
        }
    }

    fn advance(&mut self) -> char {
        let c = self.src[self.current];
        self.current += 1;
        self.column += 1;
        c
    }

    fn match_char(&mut self, expected: char) -> bool {
        if self.is_at_end() || self.peek() != expected {
            return false;
        }
        self.advance();
        true
    }

    fn peek(&self) -> char {
        self.src[self.current]
    }

    fn is_at_end(&self) -> bool {
        self.current >= self.src.len()
    }

    /// consume multi-line comment, tracking line numbers for error reporting
    fn block_comment(
        &mut self,
        start_line: usize,
        start_column: usize,
        errors: &mut Vec<RuntimeError>,
    ) -> Option<Token> {
        while !self.is_at_end() {
            let c = self.advance();
            if c == '\n' {
                self.line += 1;
                self.column = 1;
            } else if c == '*' && !self.is_at_end() && self.peek() == '/' {
                self.advance(); // consume closing '/'
                return None;
            }
        }
        // EOF in block comment
        errors.push(RuntimeError::lex(
            "Unterminated block comment (missing */)",
            start_line,
            start_column,
        ));
        None
    }

    fn string_literal(
        &mut self,
        line: usize,
        column: usize,
        errors: &mut Vec<RuntimeError>,
        quote: char,
    ) -> Option<Token> {
        let mut value = String::new();

        while !self.is_at_end() && self.peek() != quote {
            let c = self.advance();
            if c == '\\' {
                // Handle escape sequences
                if self.is_at_end() {
                    errors.push(RuntimeError::lex("Unterminated string literal", line, column));
                    return None;
                }
                match self.advance() {
                    'n' => value.push('\n'),
                    't' => value.push('\t'),
                    'r' => value.push('\r'),
                    '"' => value.push('"'),
                    '\'' => value.push('\''),
                    '\\' => value.push('\\'),
                    other => {
                        errors.push(RuntimeError::lex(
                            format!("Unknown escape sequence '\\{}'", other),
                            line,
                            column,
                        ));
                        // Consume the rest of the string to avoid cascading errors
                        while !self.is_at_end() && self.peek() != quote {
                            self.advance();
                        }
                        if !self.is_at_end() {
                            self.advance(); // consume closing quote
                        }
                        return None;
                    }
                }
            } else {
                value.push(c);
            }
        }

        if self.is_at_end() {
            errors.push(RuntimeError::lex("Unterminated string literal", line, column));
            return None;
        }

        self.advance(); // consume closing quote

        Some(Token {
            kind: TokenKind::String(value),
            line,
            column,
        })
    }

    fn number_literal(&mut self, first: char, line: usize, column: usize) -> Token {
        let mut s = String::new();
        s.push(first);

        while !self.is_at_end() && self.peek().is_ascii_digit() {
            s.push(self.advance());
        }

        if !self.is_at_end() && self.peek() == '.' {
            s.push(self.advance());
            while !self.is_at_end() && self.peek().is_ascii_digit() {
                s.push(self.advance());
            }
        }

        let value = s.parse::<f64>().unwrap();

        Token {
            kind: TokenKind::Number(value),
            line,
            column,
        }
    }

    fn identifier(&mut self, first: char, line: usize, column: usize) -> Token {
        let mut s = String::new();
        s.push(first);

        while !self.is_at_end() && is_ident_part(self.peek()) {
            s.push(self.advance());
        }

        let kind = match s.as_str() {
            "let" => TokenKind::Let,
            "const" => TokenKind::Const,
            "fn" => TokenKind::Fn,
            "if" => TokenKind::If,
            "else" => TokenKind::Else,
            "while" => TokenKind::While,
            "for" => TokenKind::For,
            "in" => TokenKind::In,
            "break" => TokenKind::Break,
            "continue" => TokenKind::Continue,
            "return" => TokenKind::Return,
            "true" => TokenKind::True,
            "false" => TokenKind::False,
            "nil" => TokenKind::Nil,
            _ => TokenKind::Identifier(s),
        };

        Token { kind, line, column }
    }
}

fn is_ident_start(c: char) -> bool {
    c.is_ascii_alphabetic() || c == '_'
}

fn is_ident_part(c: char) -> bool {
    is_ident_start(c) || c.is_ascii_digit()
}

#[cfg(test)]
#[path = "tests.rs"]
mod tests;