cg/src/parser.ts

import type { Token } from './lexer'
import type { AST, MatchCase, Pattern } from './ast'
import { ParseError } from './error'

export class Parser {
    private tokens: Token[]
    private pos: number = 0
    private source: string

    constructor(tokens: Token[], source: string) {
        this.tokens = tokens;
        this.source = source;
    }

    private current(): Token {
        if (this.pos >= this.tokens.length) {
            return { kind: 'eof' } as Token;
        }

        return this.tokens[this.pos];
    }

    private advance(): Token {
        return this.tokens[this.pos++];
    }

    private peek(offset = 1): Token {
        const pos = this.pos + offset;
        if (pos >= this.tokens.length) {
            return { kind: 'eof' } as Token;
        }
        return this.tokens[pos];
    }

    private expect(kind: Token['kind']): Token {
        const token = this.current();

        if (token.kind !== kind) {
            throw ParseError(`Expected ${kind}, got ${token.kind}`, token.line, token.column, this.source);
        }

        return this.advance();
    }

    private getPos(token: Token) {
        return {
            line: token.line,
            column: token.column,
            start: token.start
        };
    }

    private isInfixOp(): boolean {
        const kind = this.current().kind;
        return kind === 'plus' || kind === 'minus' ||
               kind === 'star' || kind === 'slash' ||
               kind === 'percent' || kind === 'caret' ||
               kind === 'ampersand' ||
               kind === 'equals-equals' || kind === 'not-equals' ||
               kind === 'less-than' || kind === 'less-equals' ||
               kind === 'greater-than' || kind === 'greater-equals' ||
               kind === 'tilde';
    }

    private tokenToOpName(token: Token): string {
        switch (token.kind) {
            case 'ampersand': return 'cat';

            // Arithmetic
            case 'plus': return 'add';
            case 'minus': return 'sub';
            case 'star': return 'mul';
            case 'slash': return 'div';
            case 'percent': return 'mod';
            case 'caret': return 'pow';

            // Comparison
            case 'equals-equals': return 'eq';
            case 'not-equals': return 'neq';
            case 'greater-than': return 'gt';
            case 'greater-equals': return 'gte';
            case 'less-than': return 'lt';
            case 'less-equals': return 'lte';

            default: throw ParseError(`Not an operator: ${token.kind}`, token.line, token.column, this.source);
        }
    }

    private canStartPrimary(): boolean {
        const kind = this.current().kind;
        return kind === 'ident' || kind === 'type-ident' ||
            kind === 'int' || kind === 'float' || kind === 'string' ||
            kind === 'open-paren' || kind === 'open-bracket' ||
            kind === 'open-brace';
    }

    private isLambdaStart(): boolean {
        const kind = this.current().kind;

        if (kind === 'backslash') {
            return true;
        }

        if (kind !== 'ident') {
            return false;
        }

        let offset = 1;
        while (true) {
            const token = this.peek(offset);
            if (token.kind === 'backslash') return true;
            if (token.kind !== 'ident') return false;
            offset++;
        }
    }

    parse(): AST {
        return this.parseExpression();
    }

    private parseExpression(): AST {
        // Lambda
        if (this.isLambdaStart()) {
            return this.parseLambda();
        }

        // Let
        if ((this.current().kind === 'ident' || this.current().kind === 'underscore') && this.peek().kind === 'equals') {
            return this.parseLet();
        }

        let expr = this.parseInfix();

        // Match
        if (this.current().kind === 'pipe') {
            return this.parseMatch(expr);
        }

        return expr;
    }

    // Used in match, doesn't match another match
    private parseExpressionNoMatch(): AST {
        // Lambda
        if (this.isLambdaStart()) {
            return this.parseLambda();
        }

        // Let
        if (this.current().kind === 'ident' && this.peek().kind === 'equals') {
            return this.parseLet();
        }

        return this.parseInfix();
    }


    private parseMatch(expr: AST): AST {
        const token = this.current();
        const cases: MatchCase[] = [];

        while(this.current().kind === 'pipe') {
            this.advance();
            const pattern = this.parsePattern();
            this.expect('backslash');
            const result = this.parseExpressionNoMatch();
            cases.push({ pattern, result, ...this.getPos(token) })
        }

        return { kind: 'match', expr, cases, ...this.getPos(token) };
    }

    private parsePattern(): Pattern {
        const token = this.current();

        // Wildcard
        if (token.kind === 'underscore') {
            this.advance();
            return { kind: 'wildcard' };
        }

        // Literal
        if (token.kind === 'int' || token.kind === 'float' || token.kind === 'string') {
            this.advance();
            return { kind: 'literal', value: token.value };
        }

        // Variable
        if (token.kind === 'ident') {
            this.advance();
            return { kind: 'var', name: token.value };
        }

        // Constructor
        if (token.kind === 'type-ident') {
            this.advance();
            const name = token.value;
            const args: Pattern[] = [];

            while (this.canStartPattern()) {
                args.push(this.parsePattern());
            }

            return { kind: 'constructor', name, args };
        }

        // List
        if (token.kind === 'open-bracket') {
            this.advance();
            const elements: Pattern[] = [];
            let first = true;
            let spreadName: string | null = null;

            while (this.current().kind !== 'close-bracket') {
                if (!first) this.expect('comma');
                first = false;

                // Spread
                if (this.current().kind === 'dot-dot-dot') {
                    this.advance();
                    const nameToken = this.expect('ident');
                    spreadName = (nameToken as { value: string }).value;
                    break;
                }

                elements.push(this.parsePattern());
            }

            this.expect('close-bracket');

            if (spreadName !== null) {
                return { kind: 'list-spread', head: elements, spread: spreadName, ...this.getPos(token) };
            }

            return { kind: 'list', elements };
        }

        // Record
        if (token.kind === 'open-brace') {
            this.advance();
            const fields: { [key: string]: Pattern } = {};
            let first = true;

            while (this.current().kind !== 'close-brace') {
                if (!first) this.expect('comma');
                first = false;

                const keyToken = this.expect('ident');
                const key = (keyToken as { value: string }).value;
                this.expect('equals');
                fields[key] = this.parsePattern();
            }

            this.expect('close-brace');
            return { kind: 'record', fields };
        }

        // Parens
        if (token.kind === 'open-paren') {
            this.advance();
            const pattern = this.parsePattern();
            this.expect('close-paren');
            return pattern;
        }

        throw ParseError(`Unexpected token in pattern: ${token.kind}`, token.line, token.column, this.source);
    }

    private canStartPattern(): boolean {
        const kind = this.current().kind;

        return kind === 'underscore' || kind === 'ident' ||
            kind === 'type-ident' || kind === 'int' ||
            kind === 'float' || kind === 'string' ||
            kind === 'open-paren';
    }

    private parseLambda(): AST {
        const token = this.current();
        const params: string[] = [];

        while (this.current().kind === 'ident') {
            const param = this.advance();
            params.push((param as { value: string }).value);
        }

        this.expect('backslash');
        const body = this.parseExpression();

        return { kind: 'lambda', params, body, ...this.getPos(token) };
    }

    private parseLet(): AST {
        const nameToken = this.current();
        let name: string;

        if (nameToken.kind === 'underscore') {
            name = '_';
            this.advance();
        } else if (nameToken.kind === 'ident') {
            name = (nameToken as { value: string }).value;
            this.advance();
        } else {
            throw ParseError(`Expected ident or underscore, got ${nameToken.kind}`, nameToken.line, nameToken.column, this.source);

        }

        this.expect('equals');
        const value = this.parseExpressionNoMatch();
        this.expect('semicolon');
        const body = this.parseExpressionNoMatch();

        return { kind: 'let', name, value, body, ...this.getPos(nameToken) };
    }

    private parseInfix(): AST {
        const token = this.current();
        let left = this.parseApplication();

        while (this.isInfixOp()) {
            const opToken = this.advance();

            if (opToken.kind === 'tilde') {
                // function application operator
                const right = this.parseApplication();

                left = {
                    kind: 'apply',
                    func: right,
                    args: [left],
                    ...this.getPos(token)
                };
            } else {
                // operators desugar to function calls
                const opName = this.tokenToOpName(opToken);
                const right = this.parseApplication();

                left = {
                    kind: 'apply',
                    func: { kind: 'variable', name: opName, ...this.getPos(token) },
                    args: [left, right],
                    ...this.getPos(token)
                }
            }
        }

        return left;
    }

    private parseApplication(): AST {
        const token = this.current();
        let func = this.parsePostfix();

        while (this.canStartPrimary()) {
            const arg = this.parsePostfix();
            func = { kind: 'apply', func, args: [arg], ...this.getPos(token) };
        }

        return func;
    }

    private parsePostfix(): AST {
        const token = this.current();
        let expr = this.parsePrimary();

        while (true) {
            if (this.current().kind === 'dot') {
                this.advance();

                if (this.current().kind === 'open-brace') {
                    // Record update
                    this.advance();
                    const updates: { [key: string]: AST } = {};
                    let first = true;

                    while (this.current().kind !== 'close-brace') {
                        if (!first) {
                            this.expect('comma');
                        }
                        first = false;

                        const keyToken = this.expect('ident');
                        const key = (keyToken as { value: string }).value;
                        this.expect('equals');
                        updates[key] = this.parseExpression();
                    }

                    this.expect('close-brace');
                    expr = { kind: 'record-update', record: expr, updates, ...this.getPos(token) }

                } else {
                    // Record access
                    const fieldToken = this.expect('ident');
                    const field = (fieldToken as { value: string }).value;
                    expr = { kind: 'record-access', record: expr, field, ...this.getPos(token) };
                }
            } else {
                break;
            }
        }

        return expr;
    }

    private parsePrimary(): AST {
        const token = this.current();

        if (token.kind === 'open-paren') {
            this.advance();
            const expr = this.parseExpression();
            this.expect('close-paren');
            return expr;
        }

        if (token.kind === 'open-bracket') {
            this.advance();

            const items: AST[] = [];
            let first = true;

            while (this.current().kind !== 'close-bracket') {
                if (!first) {
                    this.expect('comma');
                }
                first = false;

                // Spread
                if (this.current().kind === 'dot-dot-dot') {
                    const spreadToken = this.current();
                    this.advance();
                    const expr = this.parseExpression();
                    items.push({ kind: 'list-spread', spread: expr, ...this.getPos(spreadToken) })
                } else {
                    items.push(this.parseExpression());
                }
            }

            this.expect('close-bracket');
            return { kind: 'list', elements: items, ...this.getPos(token) };
        }

        if (token.kind === 'open-brace') {
            this.advance();

            const fields: { [key: string]: AST } = {};
            let first = true;

            while (this.current().kind !== 'close-brace') {
                if (!first) {
                    this.expect('comma');
                }
                first = false;

                const keyToken = this.expect('ident');
                const key = (keyToken as { value: string }).value;
                this.expect('equals');
                fields[key] = this.parseExpression();
            }

            this.expect('close-brace');
            return { kind: 'record', fields, ...this.getPos(token) };
        }

        if (token.kind === 'int') {
            this.advance();
            return { kind: 'literal', value: { kind: 'int', value: token.value }, ...this.getPos(token) };
        }

        if (token.kind === 'float') {
            this.advance();
            return { kind: 'literal', value: { kind: 'float', value: token.value }, ...this.getPos(token) };
        }

        if (token.kind === 'string') {
            this.advance();
            return { kind: 'literal', value: { kind: 'string', value: token.value }, ...this.getPos(token) };
        }

        if (token.kind === 'ident') {
            this.advance();
            return { kind: 'variable', name: token.value, ...this.getPos(token) };
        }

        if (token.kind === 'type-ident') {
            this.advance();
            return { kind: 'constructor', name: token.value, ...this.getPos(token) };
        }

        throw ParseError(`Unexpected token: ${token.kind}`, token.line, token.column, this.source);
    }
}