//! Tokenizer / lexer.

//! A simplified lexer that tokenizes arithmetic expressions into tokens,

//! then parses and evaluates them. Exercises: tagged unions with payloads,

//! match statements, optionals, if-let, while-let, for-in loops, records

//! with slice fields, and complex control flow.

/// Token types produced by the lexer.

union Token {

    /// Integer literal with its value.

    Number(i32),

    /// +

    Plus,

    /// -

    Minus,

    /// *

    Star,

    /// /

    Slash,

    /// (

    LParen,

    /// )

    RParen,

    /// End of input.

    Eof,

    /// Invalid character.

    Invalid(u8),

}

/// A lexer state over a byte slice.

record Lexer {

    source: *[u8],

    pos: u32,

}

/// A list of tokens with a fixed-size backing store.

record TokenList {

    tokens: *mut [Token],

    count: u32,

}

/// AST node for parsed expressions.

union Expr {

    Num(i32),

    BinOp(BinOpData),

    Neg(u32),

}

/// Binary operation data.

record BinOpData {

    op: u8,

    left: u32,

    right: u32,

}

/// Pool of AST nodes.

record ExprPool {

    nodes: *mut [Expr],

    count: u32,

}

/// Parser state.

record Parser {

    tokens: *[Token],

    tokenCount: u32,

    pos: u32,

    pool: *mut ExprPool,

}

/// Check if a byte is a digit.

fn isDigit(c: u8) -> bool {

    return c >= 48 and c <= 57;

}

/// Check if a byte is whitespace.

fn isSpace(c: u8) -> bool {

    return c == 32 or c == 9 or c == 10 or c == 13;

}

/// Peek at the current character, returning nil at end.

fn peek(lex: *Lexer) -> ?u8 {

    if lex.pos < lex.source.len {

        return lex.source[lex.pos];

    }

    return nil;

}

/// Advance the lexer by one character.

fn advance(lex: *mut Lexer) {

    if lex.pos < lex.source.len {

        lex.pos += 1;

    }

}

/// Skip whitespace characters.

fn skipWhitespace(lex: *mut Lexer) {

    while let ch = peek(lex); isSpace(ch) {

        advance(lex);

    }

}

/// Scan a number literal.

fn scanNumber(lex: *mut Lexer) -> i32 {

    let mut value: i32 = 0;

    while let ch = peek(lex); isDigit(ch) {

        value = value * 10 + (ch - 48) as i32;

        advance(lex);

    }

    return value;

}

/// Get the next token from the lexer.

fn nextToken(lex: *mut Lexer) -> Token {

    skipWhitespace(lex);

    if let ch = peek(lex) {

        if isDigit(ch) {

            return Token::Number(scanNumber(lex));

        }

        advance(lex);

        match ch {

            case 43 => { return Token::Plus; }

            case 45 => { return Token::Minus; }

            case 42 => { return Token::Star; }

            case 47 => { return Token::Slash; }

            case 40 => { return Token::LParen; }

            case 41 => { return Token::RParen; }

            else => { return Token::Invalid(ch); }

        }

    }

    return Token::Eof;

}

/// Tokenize the entire source into a token list.

fn tokenize(lex: *mut Lexer, list: *mut TokenList) -> bool {

    let mut done: bool = false;

    while not done {

        let tok = nextToken(lex);

        match tok {

            case Token::Invalid(_) => {

                return false;

            }

            case Token::Eof => {

                list.tokens[list.count] = tok;

                list.count += 1;

                done = true;

            }

            else => {

                if list.count >= list.tokens.len - 1 {

                    return false;

                }

                list.tokens[list.count] = tok;

                list.count += 1;

            }

        }

    }

    return true;

}

/// Allocate a new expression node.

fn newExpr(pool: *mut ExprPool, expr: Expr) -> u32 {

    let idx = pool.count;

    pool.nodes[idx] = expr;

    pool.count += 1;

    return idx;

}

/// Get the current token in the parser.

fn currentToken(p: *Parser) -> Token {

    if p.pos < p.tokenCount {

        return p.tokens[p.pos];

    }

    return Token::Eof;

}

/// Advance the parser to the next token.

fn advanceParser(p: *mut Parser) {

    if p.pos < p.tokenCount {

        p.pos += 1;

    }

}

/// Parse a primary expression (number, parenthesized expression, or unary minus).

fn parsePrimary(p: *mut Parser) -> ?u32 {

    let tok = currentToken(p);

    match tok {

        case Token::Number(n) => {

            advanceParser(p);

            return newExpr(p.pool, Expr::Num(n));

        }

        case Token::LParen => {

            advanceParser(p);

            let inner = parsePrimary(p) else {

                return nil;

            };

            // Actually parse a full expression inside parens.

            let expr = parseAddSub(p, inner);

            let close = currentToken(p);

            match close {

                case Token::RParen => {

                    advanceParser(p);

                    return expr;

                }

                else => {

                    return nil;

                }

            }

        }

        case Token::Minus => {

            advanceParser(p);

            let operand = parsePrimary(p) else {

                return nil;

            };

            return newExpr(p.pool, Expr::Neg(operand));

        }

        else => {

            return nil;

        }

    }

}

/// Parse multiplication and division (higher precedence).

fn parseMulDiv(p: *mut Parser, left: u32) -> u32 {

    let mut result: u32 = left;

    let mut cont: bool = true;

    while cont {

        let tok = currentToken(p);

        match tok {

            case Token::Star => {

                advanceParser(p);

                if let right = parsePrimary(p) {

                    result = newExpr(p.pool, Expr::BinOp(BinOpData { op: 42, left: result, right }));

                } else {

                    cont = false;

                }

            }

            case Token::Slash => {

                advanceParser(p);

                if let right = parsePrimary(p) {

                    result = newExpr(p.pool, Expr::BinOp(BinOpData { op: 47, left: result, right }));

                } else {

                    cont = false;

                }

            }

            else => {

                cont = false;

            }

        }

    }

    return result;

}

/// Parse addition and subtraction (lower precedence).

fn parseAddSub(p: *mut Parser, left: u32) -> u32 {

    let mut result: u32 = parseMulDiv(p, left);

    let mut cont: bool = true;

    while cont {

        let tok = currentToken(p);

        match tok {

            case Token::Plus => {

                advanceParser(p);

                if let right = parsePrimary(p) {

                    let rhs = parseMulDiv(p, right);

                    result = newExpr(p.pool, Expr::BinOp(BinOpData { op: 43, left: result, right: rhs }));

                } else {

                    cont = false;

                }

            }

            case Token::Minus => {

                advanceParser(p);

                if let right = parsePrimary(p) {

                    let rhs = parseMulDiv(p, right);

                    result = newExpr(p.pool, Expr::BinOp(BinOpData { op: 45, left: result, right: rhs }));

                } else {

                    cont = false;

                }

            }

            else => {

                cont = false;

            }

        }

    }

    return result;

}

/// Parse a full expression.

fn parseExpr(p: *mut Parser) -> ?u32 {

    let left = parsePrimary(p) else {

        return nil;

    };

    return parseAddSub(p, left);

}

/// Evaluate an expression tree.

fn eval(nodes: *[Expr], idx: u32) -> i32 {

    let node = nodes[idx];

    match node {

        case Expr::Num(n) => {

            return n;

        }

        case Expr::BinOp(data) => {

            let l = eval(nodes, data.left);

            let r = eval(nodes, data.right);

            if data.op == 43 {

                return l + r;

            } else if data.op == 45 {

                return l - r;

            } else if data.op == 42 {

                return l * r;

            } else if data.op == 47 {

                if r == 0 {

                    return 0;

                }

                return l / r;

            }

            return 0;

        }

        case Expr::Neg(child) => {

            return 0 - eval(nodes, child);

        }

    }

}

/// Count nodes in an expression tree.

fn countNodes(nodes: *[Expr], idx: u32) -> u32 {

    let node = nodes[idx];

    match node {

        case Expr::Num(_) => { return 1; }

        case Expr::BinOp(data) => {

            return 1 + countNodes(nodes, data.left) + countNodes(nodes, data.right);

        }

        case Expr::Neg(child) => {

            return 1 + countNodes(nodes, child);

        }

    }

}

/// Helper: tokenize, parse, and evaluate a string expression.

fn evaluate(

    source: *[u8],

    tokenBuf: *mut [Token],

    exprBuf: *mut [Expr]

) -> ?i32 {

    let mut lex = Lexer { source, pos: 0 };

    let mut list = TokenList { tokens: tokenBuf, count: 0 };

    if not tokenize(&mut lex, &mut list) {

        return nil;

    }

    let mut pool = ExprPool { nodes: exprBuf, count: 0 };

    let mut parser = Parser {

        tokens: tokenBuf,

        tokenCount: list.count,

        pos: 0,

        pool: &mut pool,

    };

    let root = parseExpr(&mut parser) else {

        return nil;

    };

    return eval(exprBuf, root);

}

/// Test simple number.

fn testNumber(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    let result = evaluate("42", tokenBuf, exprBuf) else {

        return 1;

    };

    assert result == 42;

    return 0;

}

/// Test addition.

fn testAdd(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    let result = evaluate("3 + 4", tokenBuf, exprBuf) else {

        return 1;

    };

    assert result == 7;

    return 0;

}

/// Test precedence: multiplication before addition.

fn testPrecedence(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    let result = evaluate("2 + 3 * 4", tokenBuf, exprBuf) else {

        return 1;

    };

    assert result == 14;

    return 0;

}

/// Test parenthesized expression.

fn testParens(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    let result = evaluate("(2 + 3) * 4", tokenBuf, exprBuf) else {

        return 1;

    };

    assert result == 20;

    return 0;

}

/// Test negation.

fn testNeg(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    let result = evaluate("-5 + 8", tokenBuf, exprBuf) else {

        return 1;

    };

    assert result == 3;

    return 0;

}

/// Test complex expression.

fn testComplex(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    // (10 - 3) * (2 + 1) = 7 * 3 = 21

    let result = evaluate("(10 - 3) * (2 + 1)", tokenBuf, exprBuf) else {

        return 1;

    };

    assert result == 21;

    return 0;

}

/// Test chained operations.

fn testChained(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    // 100 - 20 - 30 - 10 = 40

    let result = evaluate("100 - 20 - 30 - 10", tokenBuf, exprBuf) else {

        return 1;

    };

    assert result == 40;

    return 0;

}

/// Test token counting via for-in.

fn testTokenCount(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    let mut lex = Lexer { source: "1 + 2 * 3", pos: 0 };

    let mut list = TokenList { tokens: tokenBuf, count: 0 };

    assert tokenize(&mut lex, &mut list);

    // Should be: 1, +, 2, *, 3, Eof = 6 tokens

    assert list.count == 6;

    // Count number tokens using for-in.

    let mut numCount: u32 = 0;

    let slice = &tokenBuf[0..list.count];

    for tok in slice {

        match tok {

            case Token::Number(_) => {

                numCount += 1;

            }

            else => {}

        }

    }

    assert numCount == 3;

    return 0;

}

/// Test division and mixed operations.

fn testDivision(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    // 20 / 4 + 3 = 5 + 3 = 8

    let result = evaluate("20 / 4 + 3", tokenBuf, exprBuf) else {

        return 1;

    };

    assert result == 8;

    return 0;

}

/// Test deeply nested parentheses.

fn testDeepNesting(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    // ((((5)))) = 5

    let result = evaluate("((((5))))", tokenBuf, exprBuf) else {

        return 1;

    };

    assert result == 5;

    return 0;

}

/// Test node count of complex expression.

fn testNodeCount(tokenBuf: *mut [Token], exprBuf: *mut [Expr]) -> i32 {

    let mut lex = Lexer { source: "1 + 2 * 3", pos: 0 };

    let mut list = TokenList { tokens: tokenBuf, count: 0 };

    assert tokenize(&mut lex, &mut list);

    let mut pool = ExprPool { nodes: exprBuf, count: 0 };

    let mut parser = Parser {

        tokens: tokenBuf,

        tokenCount: list.count,

        pos: 0,

        pool: &mut pool,

    };

    let root = parseExpr(&mut parser) else {

        return 2;

    };

    // 1 + (2 * 3) = 5 nodes: Num(1), Num(2), Num(3), BinOp(*), BinOp(+)

    let count = countNodes(exprBuf, root);

    assert count == 5;

    return 0;

}

@default fn main() -> i32 {

    let mut tokenBuf: [Token; 128] = [Token::Eof; 128];

    let mut exprBuf: [Expr; 128] = [Expr::Num(0); 128];

    let r1 = testNumber(&mut tokenBuf[..], &mut exprBuf[..]);

    if r1 != 0 {

        return 10 + r1;

    }

    let r2 = testAdd(&mut tokenBuf[..], &mut exprBuf[..]);

    if r2 != 0 {

        return 20 + r2;

    }

    let r3 = testPrecedence(&mut tokenBuf[..], &mut exprBuf[..]);

    if r3 != 0 {

        return 30 + r3;

    }

    let r4 = testParens(&mut tokenBuf[..], &mut exprBuf[..]);

    if r4 != 0 {

        return 40 + r4;

    }

    let r5 = testNeg(&mut tokenBuf[..], &mut exprBuf[..]);

    if r5 != 0 {

        return 50 + r5;

    }

    let r6 = testComplex(&mut tokenBuf[..], &mut exprBuf[..]);

    if r6 != 0 {

        return 60 + r6;

    }

    let r7 = testChained(&mut tokenBuf[..], &mut exprBuf[..]);

    if r7 != 0 {

        return 70 + r7;

    }

    let r8 = testTokenCount(&mut tokenBuf[..], &mut exprBuf[..]);

    if r8 != 0 {

        return 80 + r8;

    }

    let r9 = testDivision(&mut tokenBuf[..], &mut exprBuf[..]);

    if r9 != 0 {

        return 90 + r9;

    }

    let r10 = testDeepNesting(&mut tokenBuf[..], &mut exprBuf[..]);

    if r10 != 0 {

        return 100 + r10;

    }

    let r11 = testNodeCount(&mut tokenBuf[..], &mut exprBuf[..]);

    if r11 != 0 {

        return 110 + r11;

    }

    return 0;

}