cg/src/compiler.ts

import type { AST, Pattern, Definition } from './ast';
import { _rt, store } from './runtime-js';

let matchCounter = 0;

export const definitions: Map<string, AST> = new Map();
export const dependencies: Map<string, Set<string>> = new Map();
export const dependents: Map<string, Set<string>> = new Map();

export function compile(ast: AST, useStore = true, bound = new Set<string>()): string {
    switch (ast.kind) {
        case 'literal':
            if (ast.value.kind === 'string')
                return JSON.stringify(ast.value.value);
            if (ast.value.kind === 'int' || ast.value.kind === 'float')
                return JSON.stringify(ast.value.value);
            throw new Error(`Cannot compile literal`); //  of kind ${ast.value.kind}`);

        case 'variable': {
            if (bound.has(ast.name)) {
            return sanitizeName(ast.name);
            }
            return sanitize(ast.name, useStore);
        }

        case 'lambda': {
            const newBound = new Set([...bound, ...ast.params]);
            const params = ast.params.map(sanitizeName).join(') => (');
            return `Object.assign((${params}) => ${compile(ast.body, useStore, newBound)}, { _ast: (${JSON.stringify(ast)}) })`;
        }

        case 'apply':
            // Constructor
            if (ast.func.kind === 'constructor') {
                const ctorName = ast.func.name;
                const arg = compile(ast.args[0], useStore, bound);
                return `((_a) => _a && typeof _a === 'object' && !Array.isArray(_a) && !_a._tag
                         ? { _tag: "${ctorName}", ..._a }
                         : { _tag: "${ctorName}", _0: _a })(${arg})`;
            }

            const args = ast.args.map(a => compile(a, useStore, bound)).join(')(');
            return `${compile(ast.func, useStore, bound)}(${args})`;

        case 'record': {
            const parts = ast.entries.map(entry =>
                                         entry.kind === 'spread'
                                         ? `...${compile(entry.expr, useStore, bound)}`
                                         : `${sanitizeName(entry.key)}: ${compile(entry.value, useStore, bound)}`
                                         )
            return `({${parts.join(', ')}})`;
        }

        case 'list': {
            const elements = ast.elements.map(e =>
                'spread' in e ? `...${compile(e.spread, useStore, bound)}` : compile(e, useStore, bound)
            );
            return `[${elements.join(', ')}]`;
        }

        case 'record-access':
            return `${compile(ast.record, useStore, bound)}.${sanitizeName(ast.field)}`;

        case 'record-update':
            const updates = Object.entries(ast.updates)
                .map(([k, v]) => `${sanitizeName(k)}: ${compile(v, useStore, bound)}`);
            return `({...${compile(ast.record, useStore, bound)}, ${updates.join(', ')}})`;

        case 'let':
            const newBound = new Set([...bound, ast.name]);
            return `((${sanitizeName(ast.name)}) =>
                     ${compile(ast.body, useStore, newBound)})(${compile(ast.value, useStore, bound)})`;

        case 'match':
            return compileMatch(ast, useStore, bound);

        case 'constructor':
            return `({ _tag: "${ast.name}" })`;
            /*
            return `((arg) => arg && typeof arg === 'object' && !arg._tag
                     ? { _tag: "${ast.name}", ...arg }
                     : { _tag: "${ast.name}", _0: arg })`;
                     */

        case 'rebind': {
            const rootName = getRootName(ast.target);
            const path = getPath(ast.target);
            const value = compile(ast.value, useStore, bound);

            if (!rootName) throw new Error('Rebind target must be a variable');

            if (path.length === 0) {
                return `({ _tag: "Rebind", _0: "${rootName}", _1: ${value} })`;
            } else {
                return `({ _tag: "Rebind", _0: "${rootName}", _1: ${JSON.stringify(path)}, _2: ${value} })`;
            }
        }

        default:
            throw new Error(`Cannot compile ${ast.kind}`);

    }
}

function sanitize(name: string, useStore = true): string {
    const ops: Record<string, string> = {
        'add': '_rt.add', 'sub': '_rt.sub', 'mul': '_rt.mul', 'div': '_rt.div',
        'mod': '_rt.mod', 'pow': '_rt.pow',
        'eq': '_rt.eq', 'gt': '_rt.gt', 'lt': '_rt.lt', 'gte': '_rt.gte', 'lte': '_rt.lte',
        'cat': '_rt.cat', 'max': '_rt.max', 'min': '_rt.min',
    };

    if (ops[name]) return ops[name];

    const natives = ['eval', 'measure', 'measureText', 'storeSearch', 'getSource', 'debug', 'nth', 'len', 'slice', 'join', 'chars', 'split', 'str', 'redefine', 'undefine', 'batch', 'noOp', 'focus', 'ui'];
    if (natives.includes(name)) return `_rt.${name}`;

    if (useStore) {
        return `store.${sanitizeName(name)}`;
    }
    return sanitizeName(name);
}

function sanitizeName(name: string): string {
    const reserved = [
        'default','class','function','return','const','let','var',
        'if','else','switch','case','for','while','do','break',
        'continue','new','delete','typeof','in','this','super',
        'import','export','extends','static','yield','await','async',
        'try','catch','finally','throw','null','true','false',
        'eval','Function','window','document','console'
    ];
    if (reserved.includes(name)) return `_${name}`;

    return name.replace(/-/g, '_');
}

function compileMatch(ast: AST & { kind: 'match'}, useStore = true, bound = new Set<string>()): string {
    const expr = compile(ast.expr, useStore, bound);
    const tmpVar = `_m${matchCounter++}`;

    let code = `((${tmpVar}) => { `;

    for (const c of ast.cases) {
        const { condition, bindings } = compilePattern(c.pattern, tmpVar);
        const patternBound = patternVars(c.pattern);
        const newBound = new Set([...bound, ...patternBound]);
        code += `if (${condition}) { `;
        if (bindings.length > 0) {
            code += `const ${bindings.join(', ')}; `;
        }
        code += `return ${compile(c.result, useStore, newBound)}; }`;
    }

    code += `console.error("No match for:", ${tmpVar}); throw new Error("No match"); })(${expr})`;
    return code;
}

function compilePattern(pattern: Pattern, expr: string): { condition: string, bindings: string[] } {
    switch (pattern.kind) {
        case 'wildcard':
            return { condition: 'true', bindings: [] };

        case 'var':
            return { condition: 'true', bindings: [`${sanitizeName(pattern.name)} = ${expr}`] };

        case 'literal':
            return { condition: `${expr} === ${JSON.stringify(pattern.value)}`, bindings: [] };

        case 'constructor': {
            let condition = `${expr}?._tag === "${pattern.name}"`;
            const bindings: string[] = [];
            pattern.args.forEach((argPattern, i) => {
                const sub = compilePattern(argPattern, `${expr}._${i}`);
                condition += ` && ${sub.condition}`;
                bindings.push(...sub.bindings);
            });

            return { condition, bindings };
        }

        case 'list': {
            let condition = `Array.isArray(${expr}) && ${expr}.length === ${pattern.elements.length}`;
            const bindings: string[] = [];
            pattern.elements.forEach((elemPattern, i) => {
                const sub = compilePattern(elemPattern, `${expr}[${i}]`);
                if (sub.condition !== 'true') condition += ` && ${sub.condition}`;
                bindings.push(...sub.bindings);
            });

            return { condition, bindings };
        }

        case 'list-spread': {
            let condition = `Array.isArray(${expr}) && ${expr}.length >= ${pattern.head.length}`;
            const bindings: string[] = [];
            pattern.head.forEach((elemPattern, i) => {
                const sub = compilePattern(elemPattern, `${expr}[${i}]`);
                if (sub.condition !== 'true') condition += ` && ${sub.condition}`;
                bindings.push(...sub.bindings);
            });
            bindings.push(`${sanitizeName(pattern.spread)} = ${expr}.slice(${pattern.head.length})`);

            return { condition, bindings };
        }

        case 'record': {
            let condition = 'true';
            const bindings: string[] = [];
            for (const [field, fieldPattern] of Object.entries(pattern.fields)) {
                const sub = compilePattern(fieldPattern, `${expr}.${sanitizeName(field)}`);
                if (sub.condition !== 'true') condition += ` && ${sub.condition}`;
                bindings.push(...sub.bindings);
            }

            return { condition, bindings };
        }

        default:
            return { condition: 'true', bindings: [] };
    }
}

export function compileAndRun(defs: Definition[]) {
    const compiledDefs: string[] = [];

    const topLevel = new Set(defs.map(d => d.name));

    for (const def of defs) {
        definitions.set(def.name, def.body);
        const free = freeVars(def.body);
        const deps = new Set([...free].filter(v => topLevel.has(v)));
        dependencies.set(def.name, deps);
    }

    dependents.clear();
    for (const name of topLevel) {
        dependents.set(name, new Set());
    }

    for (const [name, deps] of dependencies) {
        for (const dep of deps) {
            dependents.get(dep)?.add(name);
        }
    }

    for (const def of defs) {
        const compiled = `const ${sanitizeName(def.name)} = ${compile(def.body, false)};`;
        compiledDefs.push(compiled);

        try {
            new Function('_rt', compiled);
        } catch (e) {
            console.error(`=== BROKEN: ${def.name} ===`);
            console.error(compiled);
            throw e;
        }
    }

    const lastName = defs[defs.length - 1].name;
    const defNames = defs.map(d => sanitizeName(d.name)).join(', ');

    const code = `${compiledDefs.join('\n')}
return { ${defNames}, __result: ${sanitizeName(lastName)} };`;

    const fn = new Function('_rt', 'store', code);
    const allDefs = fn(_rt, store);

    // Populate store
    for (const [name, value] of Object.entries(allDefs)) {
        if (name !== '__result') {
            store[name] = value;
        }
    }

    return allDefs.__result;
}

export function freeVars(ast: AST, bound: Set<string> = new Set()): Set<string> {
    switch (ast.kind) {
        case 'literal':
        case 'constructor':
            return new Set();

        case 'variable':
            return bound.has(ast.name) ? new Set() : new Set([ast.name]);

        case 'lambda': {
            const newBound = new Set([...bound, ...ast.params]);
            return freeVars(ast.body, newBound);
        }

        case 'let': {
            const valueVars = freeVars(ast.value, bound);
            const bodyVars = freeVars(ast.body, new Set([...bound, ast.name]));
            return new Set([...valueVars, ...bodyVars]);
        }

        case 'apply': {
            const funcVars = freeVars(ast.func, bound);
            const argVars = ast.args.flatMap(a => [...freeVars(a, bound)])
            return new Set([...funcVars, ...argVars]);
        }

        case 'record': {
            const allVars = ast.entries.flatMap(entry =>
                                               entry.kind === 'spread'
                                               ? [...freeVars(entry.expr, bound)]
                                               : [...freeVars(entry.value, bound)]
                                               );
            return new Set(allVars);
        }

        case 'list': {
            const allVars = ast.elements.flatMap(e =>
                                                 'spread' in e ? [...freeVars(e.spread, bound)]
                                                 : [...freeVars(e, bound)]);
            return new Set(allVars);
        }

        case 'record-access':
            return freeVars(ast.record, bound);

        case 'record-update': {
            const recordVars = freeVars(ast.record, bound);
            const updateVars = Object.values(ast.updates).flatMap(v => [...freeVars(v, bound)]);
            return new Set([...recordVars, ...updateVars]);
        }

        case 'match': {
            const exprVars = freeVars(ast.expr, bound);
            const caseVars = ast.cases.flatMap(c => {
                const patternBindings = patternVars(c.pattern);
                const newBound = new Set([...bound, ...patternBindings]);
                return [...freeVars(c.result, newBound)];
            })
            return new Set([...exprVars, ...caseVars]);
        }

        case 'rebind':
            return freeVars(ast.value, bound);

        default:
            return new Set();
    }
}

function patternVars(pattern: Pattern): string[] {
    switch (pattern.kind) {
        case 'var':
            return [pattern.name];
        case 'constructor':
            return pattern.args.flatMap(patternVars);
        case 'list':
            return pattern.elements.flatMap(patternVars);
        case 'list-spread':
            return [...pattern.head.flatMap(patternVars), pattern.spread];
        case 'record':
            return Object.values(pattern.fields).flatMap(patternVars);
        default:
            return [];
    }

}

export function recompile(name: string, newAst: AST) {
    definitions.set(name, newAst);

    const topLevel = new Set(definitions.keys());
    const free = freeVars(newAst);
    const newDeps = new Set([...free].filter(v => topLevel.has(v)));

    const oldDeps = dependencies.get(name) || new Set();
    for (const dep of oldDeps) {
        dependents.get(dep)?.delete(name);
    }

    for (const dep of newDeps) {
        dependents.get(dep)?.add(name);
    }
    dependencies.set(name, newDeps);

    const toRecompile: string[] = [];
    const visited = new Set<string>();

    function collectDependents(n: string) {
        if (visited.has(n)) return;
        visited.add(n);
        toRecompile.push(n);
        for (const dep of dependents.get(n) || []) {
            collectDependents(dep);
        }
    }
    collectDependents(name);

    for (const defName of toRecompile) {
        const ast = definitions.get(defName)!;
        const compiled = compile(ast);

        const fn = new Function('_rt', 'store', `return ${compiled}`);
        store[defName] = fn(_rt, store);
    }
}

function getRootName(ast: AST): string | null {
    if (ast.kind === 'variable') return ast.name;
    if (ast.kind === 'record-access') return getRootName(ast.record);
    return null;
}

function getPath(ast: AST): string[] {
    if (ast.kind === 'variable') return [];
    if (ast.kind === 'record-access') {
        return [...getPath(ast.record), ast.field];
    }
    return [];
}