el-retired/engrams/el-lexer/src/lexer.rs

//! Single-pass lexer implementation.

use crate::error::{LexError, LexErrorKind};
use crate::token::{Span, Spanned, Token};

/// Tokenize an Engram source string.
///
/// Returns a `Vec` ending with a single [`Token::Eof`]. On the first
/// unrecognised character the function returns an error; partial output
/// is discarded.
pub fn tokenize(source: &str) -> Result<Vec<Spanned<Token>>, LexError> {
    let mut lex = Lexer::new(source);
    lex.scan_all()
}

// ── Lexer state ──────────────────────────────────────────────────────────────

struct Lexer<'src> {
    src: &'src str,
    /// Current byte position into `src`.
    pos: usize,
    line: u32,
    /// Byte position of the start of the current line (for computing columns).
    line_start: usize,
}

impl<'src> Lexer<'src> {
    fn new(src: &'src str) -> Self {
        Self { src, pos: 0, line: 1, line_start: 0 }
    }

    // ── Utilities ─────────────────────────────────────────────────────────────

    fn col_at(&self, byte_pos: usize) -> u32 {
        (byte_pos - self.line_start + 1) as u32
    }

    fn span_from(&self, start: usize) -> Span {
        Span::new(start, self.pos, self.line, self.col_at(start))
    }

    fn peek(&self) -> Option<char> {
        self.src[self.pos..].chars().next()
    }

    fn peek2(&self) -> Option<char> {
        let mut chars = self.src[self.pos..].chars();
        chars.next();
        chars.next()
    }

    fn advance(&mut self) -> Option<char> {
        let ch = self.peek()?;
        self.pos += ch.len_utf8();
        if ch == '\n' {
            self.line += 1;
            self.line_start = self.pos;
        }
        Some(ch)
    }

    /// Consume the next character only if it matches `expected`.
    fn eat(&mut self, expected: char) -> bool {
        if self.peek() == Some(expected) {
            self.advance();
            true
        } else {
            false
        }
    }

    fn at_end(&self) -> bool {
        self.pos >= self.src.len()
    }

    fn spanned(&self, tok: Token, start: usize) -> Spanned<Token> {
        Spanned::new(tok, self.span_from(start))
    }

    // ── Main scan loop ────────────────────────────────────────────────────────

    fn scan_all(&mut self) -> Result<Vec<Spanned<Token>>, LexError> {
        let mut tokens = Vec::new();
        loop {
            self.skip_whitespace_and_comments();
            if self.at_end() {
                let span = Span::point(self.pos, self.line, self.col_at(self.pos));
                tokens.push(Spanned::new(Token::Eof, span));
                break;
            }
            let tok = self.scan_token()?;
            tokens.push(tok);
        }
        Ok(tokens)
    }

    fn skip_whitespace_and_comments(&mut self) {
        loop {
            // Skip whitespace
            while let Some(ch) = self.peek() {
                if ch.is_whitespace() {
                    self.advance();
                } else {
                    break;
                }
            }
            // Skip line comments `// ...`
            if self.peek() == Some('/') && self.peek2() == Some('/') {
                self.advance(); // first /
                self.advance(); // second /
                while let Some(ch) = self.peek() {
                    self.advance();
                    if ch == '\n' {
                        break;
                    }
                }
            } else {
                break;
            }
        }
    }

    fn scan_token(&mut self) -> Result<Spanned<Token>, LexError> {
        let start = self.pos;
        let ch = self.advance().unwrap();

        let tok = match ch {
            // ── Delimiters ──────────────────────────────────────────────────
            '(' => Token::LParen,
            ')' => Token::RParen,
            '{' => Token::LBrace,
            '}' => Token::RBrace,
            '[' => Token::LBracket,
            ']' => Token::RBracket,
            ',' => Token::Comma,
            '.' => Token::Dot,
            ';' => Token::Semicolon,

            // ── Operators that may be multi-char ────────────────────────────
            '+' => Token::Plus,
            '*' => Token::Star,
            '-' => {
                if self.eat('>') {
                    Token::Arrow
                } else {
                    Token::Minus
                }
            }
            '/' => Token::Slash,
            '=' => {
                if self.eat('=') {
                    Token::EqEq
                } else if self.eat('>') {
                    Token::FatArrow
                } else {
                    Token::Eq
                }
            }
            '!' => {
                if self.eat('=') {
                    Token::NotEq
                } else {
                    Token::Not
                }
            }
            '<' => {
                if self.eat('=') {
                    Token::LtEq
                } else {
                    Token::Lt
                }
            }
            '>' => {
                if self.eat('=') {
                    Token::GtEq
                } else {
                    Token::Gt
                }
            }
            '&' => {
                if self.eat('&') {
                    Token::And
                } else {
                    return Err(LexError::new(
                        LexErrorKind::UnexpectedChar('&'),
                        self.span_from(start),
                    ));
                }
            }
            '|' => {
                if self.eat('|') {
                    Token::Or
                } else if self.eat('>') {
                    Token::PipeOp
                } else {
                    Token::Pipe
                }
            }
            '@' => Token::At,
            '?' => Token::QuestionMark,
            ':' => {
                if self.eat(':') {
                    Token::ColonColon
                } else {
                    Token::Colon
                }
            }

            // ── String literals ──────────────────────────────────────────────
            '"' => self.scan_string(start)?,

            // ── Numeric literals ─────────────────────────────────────────────
            c if c.is_ascii_digit() => self.scan_number(start, c)?,

            // ── Identifiers and keywords ─────────────────────────────────────
            c if c.is_alphabetic() || c == '_' => self.scan_ident_or_keyword(start, c),

            other => {
                return Err(LexError::new(
                    LexErrorKind::UnexpectedChar(other),
                    self.span_from(start),
                ))
            }
        };

        Ok(self.spanned(tok, start))
    }

    // ── String scanning ───────────────────────────────────────────────────────

    fn scan_string(&mut self, start: usize) -> Result<Token, LexError> {
        let mut s = String::new();
        loop {
            match self.peek() {
                None => {
                    return Err(LexError::new(
                        LexErrorKind::UnterminatedString,
                        self.span_from(start),
                    ))
                }
                Some('"') => {
                    self.advance();
                    return Ok(Token::StringLiteral(s));
                }
                Some('\\') => {
                    self.advance(); // consume backslash
                    match self.peek() {
                        Some('n') => { self.advance(); s.push('\n'); }
                        Some('t') => { self.advance(); s.push('\t'); }
                        Some('r') => { self.advance(); s.push('\r'); }
                        Some('"') => { self.advance(); s.push('"'); }
                        Some('\\') => { self.advance(); s.push('\\'); }
                        Some('0') => { self.advance(); s.push('\0'); }
                        Some(c) => {
                            let esc = c;
                            self.advance();
                            return Err(LexError::new(
                                LexErrorKind::InvalidEscape(esc),
                                self.span_from(start),
                            ));
                        }
                        None => {
                            return Err(LexError::new(
                                LexErrorKind::UnterminatedString,
                                self.span_from(start),
                            ))
                        }
                    }
                }
                Some(c) => {
                    s.push(c);
                    self.advance();
                }
            }
        }
    }

    // ── Numeric scanning ──────────────────────────────────────────────────────

    fn scan_number(&mut self, start: usize, first: char) -> Result<Token, LexError> {
        let mut raw = String::new();
        raw.push(first);

        let mut is_float = false;

        // Collect digits
        while let Some(c) = self.peek() {
            if c.is_ascii_digit() || c == '_' {
                raw.push(c);
                self.advance();
            } else if c == '.' && self.peek2().is_some_and(|d| d.is_ascii_digit()) {
                is_float = true;
                raw.push(c);
                self.advance();
            } else {
                break;
            }
        }

        if is_float {
            let clean: String = raw.chars().filter(|&c| c != '_').collect();
            let v: f64 = clean.parse().map_err(|_| LexError::new(
                LexErrorKind::InvalidNumeric,
                self.span_from(start),
            ))?;
            Ok(Token::FloatLiteral(v))
        } else {
            let clean: String = raw.chars().filter(|&c| c != '_').collect();
            let v: i64 = clean.parse().map_err(|_| LexError::new(
                LexErrorKind::IntegerOverflow,
                self.span_from(start),
            ))?;
            Ok(Token::IntLiteral(v))
        }
    }

    // ── Identifier / keyword scanning ─────────────────────────────────────────

    fn scan_ident_or_keyword(&mut self, _start: usize, first: char) -> Token {
        let mut s = String::new();
        s.push(first);
        while let Some(c) = self.peek() {
            if c.is_alphanumeric() || c == '_' {
                s.push(c);
                self.advance();
            } else {
                break;
            }
        }
        keyword_or_ident(s)
    }
}

// ── Keyword table ─────────────────────────────────────────────────────────────

fn keyword_or_ident(s: String) -> Token {
    match s.as_str() {
        "let"      => Token::Let,
        "fn"       => Token::Fn,
        "type"     => Token::Type,
        "enum"     => Token::Enum,
        "match"    => Token::Match,
        "return"   => Token::Return,
        "activate" => Token::Activate,
        "where"    => Token::Where,
        "sealed"   => Token::Sealed,
        "if"       => Token::If,
        "else"     => Token::Else,
        "for"      => Token::For,
        "in"       => Token::In,
        "while"    => Token::While,
        "test"     => Token::Test,
        "seed"     => Token::Seed,
        "assert"   => Token::Assert,
        "target"   => Token::Target,
        "protocol" => Token::Protocol,
        "impl"     => Token::Impl,
        "import"   => Token::Import,
        "from"     => Token::From,
        "as"       => Token::As,
        "true"     => Token::BoolLiteral(true),
        "false"    => Token::BoolLiteral(false),
        "with"     => Token::With,
        "retry"    => Token::Retry,
        "times"    => Token::Times,
        "fallback" => Token::Fallback,
        "reason"   => Token::Reason,
        "parallel" => Token::Parallel,
        "trace"    => Token::Trace,
        "requires" => Token::Requires,
        "deploy"   => Token::Deploy,
        "to"       => Token::To,
        "via"      => Token::Via,
        _          => Token::Ident(s),
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::token::Token;

    fn toks(src: &str) -> Vec<Token> {
        tokenize(src).unwrap().into_iter().map(|s| s.node).collect()
    }

    #[test]
    fn test_empty_source() {
        let result = tokenize("").unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].node, Token::Eof);
    }

    #[test]
    fn test_keywords() {
        let src = "let fn type enum match return activate where sealed if else for in";
        let tokens = toks(src);
        assert_eq!(tokens[0], Token::Let);
        assert_eq!(tokens[1], Token::Fn);
        assert_eq!(tokens[2], Token::Type);
        assert_eq!(tokens[3], Token::Enum);
        assert_eq!(tokens[4], Token::Match);
        assert_eq!(tokens[5], Token::Return);
        assert_eq!(tokens[6], Token::Activate);
        assert_eq!(tokens[7], Token::Where);
        assert_eq!(tokens[8], Token::Sealed);
        assert_eq!(tokens[9], Token::If);
        assert_eq!(tokens[10], Token::Else);
        assert_eq!(tokens[11], Token::For);
        assert_eq!(tokens[12], Token::In);
    }

    #[test]
    fn test_bool_literals() {
        let tokens = toks("true false");
        assert_eq!(tokens[0], Token::BoolLiteral(true));
        assert_eq!(tokens[1], Token::BoolLiteral(false));
    }

    #[test]
    fn test_int_literal() {
        let tokens = toks("42 0 1_000_000");
        assert_eq!(tokens[0], Token::IntLiteral(42));
        assert_eq!(tokens[1], Token::IntLiteral(0));
        assert_eq!(tokens[2], Token::IntLiteral(1_000_000));
    }

    #[test]
    fn test_float_literal() {
        let tokens = toks("3.14 0.5");
        assert_eq!(tokens[0], Token::FloatLiteral(3.14));
        assert_eq!(tokens[1], Token::FloatLiteral(0.5));
    }

    #[test]
    fn test_string_literal() {
        let tokens = toks(r#""hello" "world\n""#);
        assert_eq!(tokens[0], Token::StringLiteral("hello".into()));
        assert_eq!(tokens[1], Token::StringLiteral("world\n".into()));
    }

    #[test]
    fn test_operators() {
        let src = "+ - * / = == != < > <= >= && || ! -> =>";
        let tokens = toks(src);
        assert_eq!(tokens[0], Token::Plus);
        assert_eq!(tokens[1], Token::Minus);
        assert_eq!(tokens[2], Token::Star);
        assert_eq!(tokens[3], Token::Slash);
        assert_eq!(tokens[4], Token::Eq);
        assert_eq!(tokens[5], Token::EqEq);
        assert_eq!(tokens[6], Token::NotEq);
        assert_eq!(tokens[7], Token::Lt);
        assert_eq!(tokens[8], Token::Gt);
        assert_eq!(tokens[9], Token::LtEq);
        assert_eq!(tokens[10], Token::GtEq);
        assert_eq!(tokens[11], Token::And);
        assert_eq!(tokens[12], Token::Or);
        assert_eq!(tokens[13], Token::Not);
        assert_eq!(tokens[14], Token::Arrow);
        assert_eq!(tokens[15], Token::FatArrow);
    }

    #[test]
    fn test_delimiters() {
        let src = "( ) { } [ ] , : :: . ;";
        let tokens = toks(src);
        assert_eq!(tokens[0], Token::LParen);
        assert_eq!(tokens[1], Token::RParen);
        assert_eq!(tokens[2], Token::LBrace);
        assert_eq!(tokens[3], Token::RBrace);
        assert_eq!(tokens[4], Token::LBracket);
        assert_eq!(tokens[5], Token::RBracket);
        assert_eq!(tokens[6], Token::Comma);
        assert_eq!(tokens[7], Token::Colon);
        assert_eq!(tokens[8], Token::ColonColon);
        assert_eq!(tokens[9], Token::Dot);
        assert_eq!(tokens[10], Token::Semicolon);
    }

    #[test]
    fn test_line_comment_skipped() {
        let tokens = toks("let // this is a comment\nfn");
        assert_eq!(tokens[0], Token::Let);
        assert_eq!(tokens[1], Token::Fn);
    }

    #[test]
    fn test_span_line_col() {
        let src = "let\nfn";
        let tokens = tokenize(src).unwrap();
        assert_eq!(tokens[0].span.line, 1);
        assert_eq!(tokens[0].span.col, 1);
        assert_eq!(tokens[1].span.line, 2);
        assert_eq!(tokens[1].span.col, 1);
    }

    #[test]
    fn test_unterminated_string_error() {
        let result = tokenize(r#""unterminated"#);
        assert!(result.is_err());
    }

    #[test]
    fn test_at_token() {
        let tokens = toks("@public");
        assert_eq!(tokens[0], Token::At);
        assert_eq!(tokens[1], Token::Ident("public".into()));
    }

    #[test]
    fn test_hello_world_program() {
        let src = r#"
fn greet(name: String) -> String {
    return "Hello, " + name
}

let msg: String = greet("Will")
"#;
        let tokens = tokenize(src).unwrap();
        // Verify it tokenizes without error and the last token is EOF
        assert_eq!(tokens.last().unwrap().node, Token::Eof);
        // Should have a reasonable number of tokens
        assert!(tokens.len() > 10);
    }

    #[test]
    fn test_activate_syntax() {
        let src = r#"activate User where "customer who purchased recently""#;
        let tokens = toks(src);
        assert_eq!(tokens[0], Token::Activate);
        assert_eq!(tokens[1], Token::Ident("User".into()));
        assert_eq!(tokens[2], Token::Where);
        assert_eq!(tokens[3], Token::StringLiteral("customer who purchased recently".into()));
    }

    #[test]
    fn test_colon_colon_path() {
        let tokens = toks("Status::Active");
        assert_eq!(tokens[0], Token::Ident("Status".into()));
        assert_eq!(tokens[1], Token::ColonColon);
        assert_eq!(tokens[2], Token::Ident("Active".into()));
    }

    #[test]
    fn test_new_keywords() {
        let tokens = toks("protocol impl import from as");
        assert_eq!(tokens[0], Token::Protocol);
        assert_eq!(tokens[1], Token::Impl);
        assert_eq!(tokens[2], Token::Import);
        assert_eq!(tokens[3], Token::From);
        assert_eq!(tokens[4], Token::As);
    }

    #[test]
    fn test_pipe_token() {
        let tokens = toks("|x: Int|");
        assert_eq!(tokens[0], Token::Pipe);
        assert_eq!(tokens[1], Token::Ident("x".into()));
        assert_eq!(tokens[4], Token::Pipe);
    }

    #[test]
    fn test_question_mark_token() {
        let tokens = toks("x?");
        assert_eq!(tokens[0], Token::Ident("x".into()));
        assert_eq!(tokens[1], Token::QuestionMark);
    }

    #[test]
    fn test_ident_with_underscore() {
        let tokens = toks("my_var _private __double");
        assert_eq!(tokens[0], Token::Ident("my_var".into()));
        assert_eq!(tokens[1], Token::Ident("_private".into()));
        assert_eq!(tokens[2], Token::Ident("__double".into()));
    }

    #[test]
    fn test_sealed_block_tokens() {
        // sealed { let x: String = "secret" }
        // [0]=Sealed [1]={ [2]=let [3]=x [4]=: [5]=String [6]==  [7]="secret" [8]=}
        let src = "sealed { let x: String = \"secret\" }";
        let tokens = toks(src);
        assert_eq!(tokens[0], Token::Sealed);
        assert_eq!(tokens[1], Token::LBrace);
        assert_eq!(tokens[2], Token::Let);
        assert_eq!(tokens[7], Token::StringLiteral("secret".into()));
        assert_eq!(tokens[8], Token::RBrace);
    }
}