RustAst/src/ast/lexer.rs

use crate::ast::types::SourceLocation;
use std::iter::Peekable;
use std::rc::Rc;
use std::str::Chars;

#[derive(Debug, Clone, PartialEq)]
pub enum TokenKind {
    LeftParen,
    RightParen,
    LeftBracket,
    RightBracket,
    LeftBrace,
    RightBrace,
    Quote,
    Backtick,
    Tilde,
    At,
    Identifier(Rc<str>),
    Keyword(Rc<str>),
    Integer(i64),
    Float(f64),
    String(Rc<str>),
    EOF,
}

#[derive(Debug, Clone)]
pub struct Token {
    pub kind: TokenKind,
    pub location: SourceLocation,
}

pub struct Lexer<'a> {
    input: Peekable<Chars<'a>>,
    line: u32,
    col: u32,
}

impl<'a> Lexer<'a> {
    pub fn new(input: &'a str) -> Self {
        Self {
            input: input.chars().peekable(),
            line: 1,
            col: 1,
        }
    }

    pub fn next_token(&mut self) -> Result<Token, String> {
        self.skip_whitespace();

        let start_location = SourceLocation {
            line: self.line,
            col: self.col,
        };

        let char = match self.input.next() {
            Some(c) => {
                let current_char = c;
                self.col += 1;
                current_char
            }
            None => {
                return Ok(Token {
                    kind: TokenKind::EOF,
                    location: start_location,
                });
            }
        };

        let kind = match char {
            '(' => TokenKind::LeftParen,
            ')' => TokenKind::RightParen,
            '[' => TokenKind::LeftBracket,
            ']' => TokenKind::RightBracket,
            '{' => TokenKind::LeftBrace,
            '}' => TokenKind::RightBrace,
            '\'' => TokenKind::Quote,
            '`' => TokenKind::Backtick,
            '~' => TokenKind::Tilde,
            '@' => TokenKind::At,
            ':' => {
                let id = self.read_while(|c| !c.is_whitespace() && !"()[]{}\"'`~@;".contains(c));
                TokenKind::Keyword(Rc::from(id))
            }
            '"' => TokenKind::String(Rc::from(self.read_string()?)),
            c if c.is_ascii_digit()
                || (c == '-' && self.input.peek().is_some_and(|p| p.is_ascii_digit())) =>
            {
                let mut num_str = String::from(c);
                while let Some(&next) = self.peek() {
                    if next.is_ascii_digit() || next == '.' {
                        num_str.push(self.input.next().unwrap());
                        self.col += 1;
                    } else {
                        break;
                    }
                }

                if num_str.contains('.') {
                    TokenKind::Float(num_str.parse().map_err(|_| "Invalid float format")?)
                } else {
                    TokenKind::Integer(num_str.parse().map_err(|_| "Invalid integer format")?)
                }
            }
            c if is_ident_start(c) => {
                let mut id = String::from(c);
                id.push_str(&self.read_while(is_ident_char));
                TokenKind::Identifier(Rc::from(id))
            }
            _ => {
                return Err(format!(
                    "Unexpected character: {} at {}:{}",
                    char,
                    self.line,
                    self.col - 1
                ));
            }
        };

        Ok(Token {
            kind,
            location: start_location,
        })
    }

    fn peek(&mut self) -> Option<&char> {
        self.input.peek()
    }

    fn skip_whitespace(&mut self) {
        while let Some(&c) = self.peek() {
            if c.is_whitespace() || is_invisible(c) || c == ',' {
                if c == '\n' {
                    self.line += 1;
                    self.col = 1;
                } else {
                    self.col += 1;
                }
                self.input.next();
            } else if c == ';' {
                for c in self.input.by_ref() {
                    if c == '\n' {
                        self.line += 1;
                        self.col = 1;
                        break;
                    }
                }
            } else {
                break;
            }
        }
    }

    fn read_while<F>(&mut self, mut predicate: F) -> String
    where
        F: FnMut(char) -> bool,
    {
        let mut s = String::new();
        while let Some(&c) = self.peek() {
            if predicate(c) {
                s.push(self.input.next().unwrap());
                self.col += 1;
            } else {
                break;
            }
        }
        s
    }

    fn read_string(&mut self) -> Result<String, String> {
        let mut s = String::new();
        while let Some(c) = self.input.next() {
            self.col += 1;
            match c {
                '"' => return Ok(s),
                '\\' => {
                    let next = self.input.next().ok_or("Unterminated string escape")?;
                    self.col += 1;
                    match next {
                        'n' => s.push('\n'),
                        'r' => s.push('\r'),
                        't' => s.push('\t'),
                        '\\' => s.push('\\'),
                        '"' => s.push('"'),
                        _ => s.push(next),
                    }
                }
                '\n' => {
                    self.line += 1;
                    self.col = 1;
                    s.push(c);
                }
                _ => s.push(c),
            }
        }
        Err("Unterminated string".to_string())
    }
}

fn is_ident_start(c: char) -> bool {
    c.is_alphabetic() || "+-*/<=>!$%&_?.".contains(c)
}

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

/// Returns true if the character is an invisible control character
/// that should be ignored by the lexer (like BOM or Zero Width Space).
fn is_invisible(c: char) -> bool {
    match c {
        '\u{FEFF}' | // BOM
        '\u{200B}' | // Zero Width Space
        '\u{200C}' | // Zero Width Non-Joiner
        '\u{200D}' | // Zero Width Joiner
        '\u{2060}'   // Word Joiner
        => true,
        _ => false,
    }
}