preserves/implementations/javascript/packages/schema/src/compiler.ts

import { Pattern, NamedPattern, Schema, Input, Environment, Ref, lookup } from "./meta";
import * as M from './meta';
import { Annotated, Bytes, Dictionary, Fold, fold, KeyedSet, Position, preserves, Record, Set, Tuple, Value } from "@preserves/core";
import { Formatter, parens, seq, Item, opseq, block, commas, brackets, anglebrackets, braces } from "./block";
import { refPosition } from "./reader";
import { Alternative, Definition } from "gen/schema";

export interface CompilerOptions {
    preservesModule?: string;
    defaultPointer?: Ref;
    warn?(message: string, pos: Position | null): void;
};

function fnblock(... items: Item[]): Item {
    return seq('((() => ', block(... items), ')())');
}

class ModuleContext {
    readonly env: Environment;
    readonly schema: Schema;
    readonly options: CompilerOptions;

    readonly literals = new Dictionary<never, string>();
    readonly typedefs: Item[] = [];
    readonly functiondefs: Item[] = [];
    readonly imports = new KeyedSet<[string, string]>();

    constructor(env: Environment, schema: Schema, options: CompilerOptions) {
        this.env = env;
        this.schema = schema;
        this.options = options;
    }

    literal(v: Input): Item {
        let varname = this.literals.get(v);
        if (varname === void 0) {
            const s = v.asPreservesText()
                .replace('_', '__')
                .replace('*', '_STAR_');
            varname = M.isValidToken('_' + s, true) ? '$' + s : '__lit' + this.literals.size;
            this.literals.set(v, varname);
        }
        return varname;
    }

    derefPattern([_name, p]: [string, Alternative]): Definition {
        if (p.label === M.$ref) {
            return lookup(refPosition(p), p, this.env,
                          (p) => p,
                          (p) => p,
                          (_modId, _modPath, pp) => pp ?? p);
        } else {
            return p;
        }
    }

    defineType(f: Item): void {
        this.typedefs.push(f);
    }

    defineFunction(f: (ctx: FunctionContext) => Item): void {
        this.functiondefs.push(f(new FunctionContext(this)));
    }
}

class FunctionContext {
    readonly mod: ModuleContext;

    tempCounter = 0;
    temps: string[] = [];

    constructor(mod: ModuleContext) {
        this.mod = mod;
    }

    gentemp(): string {
        const varname = '_tmp' + this.tempCounter++;
        this.temps.push(varname);
        return varname;
    }

    gentemps(n: number): string[] {
        const temps = [];
        while (temps.length < n) temps.push(this.gentemp());
        return temps;
    }

    block(f: () => Item[]): Item {
        const oldTemps = this.temps;
        this.temps = [];
        const items = f();
        const ts = this.temps;
        this.temps = oldTemps;
        return block(
            ... ts.length > 0 ? [seq(`let `, commas(... ts), ': any')] : [],
            ... items);
    }
}

function unname(p: NamedPattern): Pattern {
    return (p.label === M.$named) ? p[1] : p;
}

function fieldName(np: NamedPattern, index: number): string {
    return (np.label === M.$named) ? np[0].description! : `_field${index}`;
}

function accumulateCompound(ctx: FunctionContext,
                            p: Pattern,
                            kFail: () => Item[],
                            kAcc: (temp: string) => Item[]): Item
{
    const t = ctx.gentemp();
    return seq(`while (!d.closeCompound()) `, ctx.block(() => [
        seq(`${t} = void 0`),
        ... decoderFor(ctx, p, t),
        seq(`if (${t} === void 0) `, block(... kFail(), seq(`break`))),
        ... kAcc(t)]));
}

function decoderForTuple(ctx: FunctionContext,
                         tuplePattern: Pattern,
                         ps: Pattern[],
                         dest: string,
                         recordFields: boolean,
                         variablePattern: Pattern | undefined): Item[]
{
    const temps = ctx.gentemps(ps.length);

    function loop(i: number): Item[] {
        if (i < ps.length) {
            return [... decoderFor(ctx, ps[i], temps[i]),
                    seq(`if (${temps[i]} !== void 0) `, ctx.block(() => loop(i + 1)))];
        } else {
            if (variablePattern === void 0) {
                return [seq(`if (d.closeCompound()) ${dest} = `, brackets(... temps),
                            ` as `, typeFor(ctx.mod, tuplePattern))];
            } else {
                return [block(
                    seq(`let vN: `, typeFor(ctx.mod, tuplePattern),
                        ` | undefined = `, brackets(... temps)),
                    accumulateCompound(ctx,
                                       variablePattern,
                                       () => [`vN = void 0`],
                                       (t) => [`vN.push(${t})`]),
                    seq(`${dest} = vN`))];
            }
        }
    }

    return recordFields
        ? loop(0)
        : [seq(`if (d.openSequence()) `, ctx.block(() => loop(0)))];
}

function decoderFor(ctx: FunctionContext, p: Definition, dest: string, recordFields = false): Item[]
{
    switch (p.label) {
        case M.$atom:
            switch (p[0]) {
                case M.$Boolean: return [`${dest} = d.nextBoolean()`];
                case M.$Float: return [`${dest} = d.nextFloat()`];
                case M.$Double: return [`${dest} = d.nextDouble()`];
                case M.$SignedInteger: return [`${dest} = d.nextSignedInteger()`];
                case M.$String: return [`${dest} = d.nextString()`];
                case M.$ByteString: return [`${dest} = d.nextByteString()`];
                case M.$Symbol: return [`${dest} = d.nextSymbol()`];
            }
        case M.$lit: {
            let n: string;
            switch (typeof p[0]) {
                case 'boolean': n = `d.nextBoolean()`; break;
                case 'string': n = `d.nextString()`; break;
                case 'number': n = `d.nextSignedInteger()`; break;
                case 'symbol': n = `d.nextSymbol()`; break;
                default: n = `d.next()`; break;
            }
            return [`${dest} = _.asLiteral(${n}, ${ctx.mod.literal(p[0])})`];
        }
        case M.$ref:
            return lookup(refPosition(p), p, ctx.mod.env,
                          (_p) => [`${dest} = decode${p[1].description!}(d)`],
                          (p) => decoderFor(ctx, p, dest),
                          (modId, modPath,_p) => {
                              ctx.mod.imports.add([modId, modPath]);
                              return [`${dest} = ${modId}.decode${p[1].description!}(d)`];
                          });
        case M.$or: {
            const alts = p[0];
            const recs = alts.map(p => ctx.mod.derefPattern(p));
            if (recs.length > 1 && recs.every(pp => pp.label === M.$rec)) {
                // Hoist the record check up.
                // This is pretty hacky. If we lift the level of
                // discourse a little, we can do this
                // automatically and generically...
                return [seq(`if (d.openRecord()) `, ctx.block(() => {
                    const label = ctx.gentemp();
                    const mark = ctx.gentemp();
                    function loop(i: number): Item[] {
                        const alt = recs[i];
                        if (alt.label !== M.$rec) throw new Error("Internal error"); // avoid a cast
                        return [
                            seq(`if (`, predicateFor(ctx, label, unname(alt[0])), `) `, ctx.block(() => {
                                const fs = ctx.gentemp();
                                return [... decoderFor(ctx, unname(alt[1]), fs, true),
                                        seq(`if (${fs} !== void 0) ${dest} = _.Record`,
                                            anglebrackets(typeFor(ctx.mod, unname(alt[0])),
                                                          typeFor(ctx.mod, unname(alt[1]))),
                                            parens(seq(label, ` as any`),
                                                   seq(fs, ` as any`)))];
                            })),
                            ... (i < recs.length - 1)
                                ? [seq(`if (${dest} === void 0) `,
                                       ctx.block(() => [`d.restoreMark(${mark})`, ... loop(i + 1)]))]
                                : [],
                        ];
                    }
                    return [seq(`${label} = d.next()`),
                            seq(`${mark} = d.mark()`),
                            ... loop(0)];
                }))];
            } else {
                switch (alts.length) {
                    case 0: return []; // assume dest is already void 0
                    case 1: return decoderFor(ctx, alts[0][1], dest);
                    default: {
                        const mark = ctx.gentemp();
                        function loop(i: number): Item[] {
                            return [
                                ... decoderFor(ctx, alts[i][1], dest),
                                ... (i < alts.length - 1)
                                    ? [seq(`if (${dest} === void 0) `, ctx.block(() =>
                                        [`d.restoreMark(${mark})`, ... loop(i + 1)]))]
                                    : [],
                            ];
                        }
                        return [`${mark} = d.mark()`, ... loop(0)];
                    }
                }
            }
        }
        case M.$and:
            switch (p[0].length) {
                case 0: return [`${dest} = d.next()`];
                case 1: return decoderFor(ctx, unname(p[0][0]), dest);
                default: {
                    const [pp0, ... ppN] = p[0];
                    return [... decoderFor(ctx, unname(pp0), dest),
                            seq(`if (!`, opseq('true', ' && ',
                                               ... ppN.map(pp =>
                                                   predicateFor(ctx, dest, unname(pp)))),
                                `) ${dest} = void 0`)];
                }
            }
        case M.$pointer:
            return [`${dest} = _decodePtr(d)`];
        case M.$rec:
            // assume dest is already void 0
            return [seq(`if (d.openRecord()) `, ctx.block(() => {
                const label = ctx.gentemp();
                return [... decoderFor(ctx, unname(p[0]), label),
                        seq(`if (${label} !== void 0) `, ctx.block(() => {
                            const fs = ctx.gentemp();
                            return [... decoderFor(ctx, unname(p[1]), fs, true),
                                    seq(`if (${fs} !== void 0) ${dest} = _.Record`,
                                        anglebrackets(typeFor(ctx.mod, unname(p[0])),
                                                      typeFor(ctx.mod, unname(p[1]))),
                                        parens(seq(label, ` as any`),
                                               seq(fs, ` as any`)))];
                        }))];
            }))];
        case M.$tuple:
            // assume dest is already void 0
            return decoderForTuple(ctx, p, p[0].map(unname), dest, recordFields, void 0);
        case M.$tuple_STAR_:
            // assume dest is already void 0
            return decoderForTuple(ctx, p, p[0].map(unname), dest, recordFields, unname(p[1]));
        case M.$setof:
            // assume dest is already void 0
            return [seq(`if (d.openSet()) `, ctx.block(() => [
                seq(`let r: `, typeFor(ctx.mod, p), ` | undefined = new _.KeyedSet()`),
                accumulateCompound(ctx,
                                   p[0],
                                   () => [`r = void 0`],
                                   (t) => [`r.add(${t})`]),
                `${dest} = r`]))];
        case M.$dictof:
            // assume dest is already void 0
            return [seq(`if (d.openDictionary()) `, ctx.block(() => [
                seq(`let r: `, typeFor(ctx.mod, p), ` | undefined = new _.KeyedDictionary()`),
                seq(`while (!d.closeCompound()) `, ctx.block(() => [
                    seq(`let K: undefined | `, typeFor(ctx.mod, p[0]), ` = void 0`),
                    ... decoderFor(ctx, p[0], 'K'),
                    seq(`if (K === void 0) { r = void 0; break; }`),
                    seq(`let V: undefined | `, typeFor(ctx.mod, p[1]), ` = void 0`),
                    ... decoderFor(ctx, p[1], 'V'),
                    seq(`if (V === void 0) { r = void 0; break; }`),
                    seq(`r.set(K, V)`)])),
                seq(`${dest} = r`)]))];
        case M.$dict:
            return [seq(`${dest} = d.next()`),
                    seq(`if (${dest} !== void 0 && !(`, predicateFor(ctx, dest, p),
                        `)) ${dest} = void 0`)];
        default:
            ((_p: never) => {})(p);
            throw new Error("Unreachable");
    }
}

function typeFor(mod: ModuleContext, p: Pattern): Item {
    switch (p.label) {
        case M.$atom:
            switch (p[0]) {
                case M.$Boolean: return `boolean`;
                case M.$Float: return `_.SingleFloat`;
                case M.$Double: return `_.DoubleFloat`;
                case M.$SignedInteger: return `number`;
                case M.$String: return `string`;
                case M.$ByteString: return `_.Bytes`;
                case M.$Symbol: return `symbol`;
            }
        case M.$lit:
            return `(typeof ${mod.literal(p[0])})`;
        case M.$ref:
            return lookup(refPosition(p), p, mod.env,
                          (_p) => p[1].description!,
                          (p) => typeForAlternative(mod, p),
                          (modId, modPath,_p) => {
                              mod.imports.add([modId, modPath]);
                              return `${modId}.${p[1].description!}`;
                          });
        case M.$pointer:
            return `_ptr`;
        case M.$rec:
            return seq('_.Record', anglebrackets(typeFor(mod, unname(p[0])),
                                                 typeFor(mod, unname(p[1])), '_ptr'));
        case M.$tuple:
            return brackets(... p[0].map(pp => typeFor(mod, unname(pp))));
        case M.$tuple_STAR_:
            if (p[0].length === 0) {
                return seq('Array<', typeFor(mod, unname(p[1])), '>');
            } else {
                return brackets(... p[0].map(pp => typeFor(mod, unname(pp))),
                                seq('... Array<', typeFor(mod, unname(p[1])), '>'));
            }
        case M.$setof:
            return seq('_.KeyedSet', anglebrackets(typeFor(mod, p[0]), '_ptr'));
        case M.$dictof:
            return seq('_.KeyedDictionary', anglebrackets(
                typeFor(mod, p[0]),
                typeFor(mod, p[1]),
                '_ptr'));
        case M.$dict:
            return parens(seq(
                block(
                    ... Array.from(p[0]).map(([k, vp]) =>
                        seq(`get(k: typeof ${mod.literal(k)}): `, typeFor(mod, unname(vp)))),
                    ... Array.from(p[0]).map(([k, _vp]) =>
                        seq(`has(k: typeof ${mod.literal(k)}): true`))),
                ' & _.Dictionary<_ptr>'));
        default:
            ((_p: never) => {})(p);
            throw new Error("Unreachable");
    }
}

function typeForDefinition(mod: ModuleContext, _name: symbol, d: Definition): Item {
    if (d.label === M.$or) {
        return opseq('never', ' | ', ... d[0].map(a => typeForAlternative(mod, a[1])));
    } else {
        return typeForAlternative(mod, d);
    }
}

function typeForAlternative(mod: ModuleContext, a: Alternative): Item {
    if (a.label === M.$and) {
        return opseq('_val', ' & ', ... a[0].map(p => typeFor(mod, unname(p))));
    } else {
        return typeFor(mod, a);
    }
}

function predicateFor(ctx: FunctionContext, v: string, p: Definition, recordOkAsTuple = false): Item
{
    switch (p.label) {
        case M.$atom:
            switch (p[0]) {
                case M.$Boolean: return `typeof ${v} === 'boolean'`;
                case M.$Float: return `_.Float.isSingle(${v})`;
                case M.$Double: return `_.Float.isDouble(${v})`;
                case M.$SignedInteger: return `typeof ${v} === 'number'`;
                case M.$String: return `typeof ${v} === 'string'`;
                case M.$ByteString: return `_.Bytes.isBytes(${v})`;
                case M.$Symbol: return `typeof ${v} === 'symbol'`;
            }
        case M.$lit:
            return `_.is(${v}, ${ctx.mod.literal(p[0])})`;
        case M.$ref:
            return lookup(refPosition(p), p, ctx.mod.env,
                          (_p) => `is${Ref._.name(p).description!}(${v})`,
                          (pp) => predicateFor(ctx, v, pp),
                          (modId, modPath, _p) => {
                              ctx.mod.imports.add([modId, modPath]);
                              return `${modId}.is${Ref._.name(p).description!}(${v})`;
                          });
        case M.$or: {
            const alts = p[0];
            const recs = alts.map(p => ctx.mod.derefPattern(p));
            if (recs.length > 1 && recs.every(pp => pp.label === M.$rec)) {
                return seq(
                    `_.Record.isRecord<_val, _.Tuple<_val>, _ptr>(${v}) && `,
                    parens(opseq('false', ' || ',
                                 ... recs.map(r =>
                                     (r.label !== M.$rec) ? '' : parens(seq(
                                         predicateFor(ctx, `${v}.label`, unname(r[0])),
                                         ' && ',
                                         predicateFor(ctx, v, unname(r[1]), true)))))));
            } else {
                return opseq('false', ' || ', ... p[0].map(pp => predicateFor(ctx, v, pp[1])));
            }
        }
        case M.$and:
            return opseq('true', ' && ', ... p[0].map(pp => predicateFor(ctx, v, unname(pp))));
        case M.$pointer:
            return `_.isPointer(${v})`;
        case M.$rec:
            return opseq('true', ' && ',
                         `_.Record.isRecord<_val, _.Tuple<_val>, _ptr>(${v})`,
                         predicateFor(ctx, `${v}.label`, unname(p[0])),
                         predicateFor(ctx, v, unname(p[1]), true));
        case M.$tuple:
            return opseq('true', ' && ',
                         ... (recordOkAsTuple ? []
                             : [`_.Array.isArray(${v})`, `!_.Record.isRecord<_val, _.Tuple<_val>, _ptr>(${v})`]),
                         `(${v}.length === ${p[0].length})`,
                         ... p[0].map((pp, i) => predicateFor(ctx, `${v}[${i}]`, unname(pp))));
        case M.$tuple_STAR_:
            return opseq('true', ' && ',
                         ... (recordOkAsTuple ? []
                             : [`_.Array.isArray(${v})`, `!_.Record.isRecord<_val, _.Tuple<_val>, _ptr>(${v})`]),
                         `(${v}.length >= ${p[0].length})`,
                         seq(p[0].length > 0 ? `${v}.slice(${p[0].length})` : v,
                             `.every(v => `,
                             parens(predicateFor(ctx, 'v', unname(p[1]))),
                             `)`),
                         ... p[0].map((pp, i) => predicateFor(ctx, `${v}[${i}]`, unname(pp))));
        case M.$setof:
            return opseq('true', ' && ',
                         `_.Set.isSet<_val>(${v})`,
                         fnblock(
                             seq(`for (const vv of ${v}) `, block(
                                 seq('if (!(', predicateFor(ctx, 'vv', p[0]), ')) return false'))),
                             seq('return true')));
        case M.$dictof:
            return opseq('true', ' && ',
                         `_.Dictionary.isDictionary<_ptr>(${v})`,
                         fnblock(
                             seq(`for (const e of ${v}) `, block(
                                 seq('if (!(', predicateFor(ctx, 'e[0]', p[0]), ')) return false'),
                                 seq('if (!(', predicateFor(ctx, 'e[1]', p[1]), ')) return false'))),
                             seq('return true')));
        case M.$dict:
            return opseq('true', ' && ',
                         `_.Dictionary.isDictionary<_ptr>(${v})`,
                         ... Array.from(p[0]).map(([k, vp]) => {
                             const tmp = ctx.gentemp();
                             return parens(seq(
                                 `(${tmp} = ${v}.get(${ctx.mod.literal(k)})) !== void 0 && `,
                                 predicateFor(ctx, tmp, unname(vp))));
                         }));
        default:
            ((_p: never) => {})(p);
            throw new Error("Unreachable");
    }
}

export function compile(env: Environment, schema: Schema, options: CompilerOptions = {}): string {
    const mod = new ModuleContext(env, schema, options);

    const pointerName = Schema._._field0(schema).get(M.$pointer);
    mod.defineType(seq(`export type _ptr = `,
                       pointerName === false ? 'never' : typeFor(mod, pointerName),
                       `;`));
    mod.defineType(`export type _val = _.Value<_ptr>;`);

    mod.defineFunction(ctx =>
        seq(`export const _decodePtr = `,
            (pointerName === false
                ? '() => { throw new _.DecodeError("Pointers forbidden"); }'
                : seq(`(d: _.TypedDecoder<_ptr>) => `, ctx.block(() => [
                    seq(`let result`),
                    ... decoderFor(ctx, pointerName, 'result'),
                    seq(`return result`)]))),
            `;`));

    for (const [name0, def] of Schema._._field0(schema).get(M.$definitions)) {
        const name = name0 as symbol;

        if (def.label === M.$rec &&
            def[0].label === M.$lit &&
            def[1].label === M.$tuple)
        {
            function fieldEntry(np: NamedPattern, index: number): Item {
                return seq(JSON.stringify(fieldName(np, index)), ': ', typeFor(mod, unname(np)));
            }
            mod.defineType(
                seq(`export const ${name.description!} = _.Record.makeConstructor<`,
                    braces(... def[1][0].map(fieldEntry)),
                    `, _ptr>()(${mod.literal(def[0][0])}, `,
                    JSON.stringify(def[1][0].map(fieldName)), `);`));
        }

        mod.defineType(
            seq(`export type ${name.description!} = `, typeForDefinition(mod, name, def), `;`));
    }

    for (const [name0, def] of Schema._._field0(schema).get(M.$definitions)) {
        const name = name0 as symbol;

        mod.defineFunction(ctx =>
            seq(`export function is${name.description!}`,
                '(v: any): v is ', name.description!, ' ',
                ctx.block(() => [seq('return ', predicateFor(ctx, 'v', def))])));

        mod.defineFunction(ctx =>
            seq(`export function as${name.description!}`,
                '(v: any): ', name.description!, ' ',
                ctx.block(() => [
                    seq(`if (!is${name.description!}(v)) `,
                        block(`throw new TypeError(\`Invalid ${name.description!}: \${_.stringify(v)}\`)`),
                        ' else ',
                        block(`return v`))])));

        mod.defineFunction(ctx =>
            seq(`export function decode${name.description!}`,
                `(d: _.TypedDecoder<_ptr>): ${name.description!} | undefined `,
                ctx.block(() => [seq(`let result`),
                                 ... decoderFor(ctx, def, 'result'),
                                 seq(`return result`)])));
    }

    const f = new Formatter();
    f.write(`import * as _ from ${JSON.stringify(options.preservesModule ?? '@preserves/core')};\n`);
    mod.imports.forEach(([identifier, path]) => {
        f.write(`import * as ${identifier} from ${JSON.stringify(path)};\n`);
    });
    f.newline();

    const sortedLiterals = Array.from(mod.literals);
    sortedLiterals.sort((a, b) => a[1] < b[1] ? -1 : a[1] === b[1] ? 0 : 1);
    for (const [lit, varname] of sortedLiterals) {
        f.write(seq(`export const ${varname} = `, sourceCodeFor(lit), `;\n`));
    }
    f.newline();

    mod.typedefs.forEach(t => {
        f.write(t);
        f.newline();
        f.newline();
    });
    f.newline();

    mod.functiondefs.forEach(p => {
        f.write(p);
        f.newline();
        f.newline();
    });

    return f.toString();
}

export function stringSource(s: string) {
    return JSON.stringify(s);
}

export function sourceCodeFor(v: Value<any>): Item {
    return fold(v, {
        boolean(b: boolean): Item { return b.toString(); },
        single(f: number): Item { return f.toString(); },
        double(f: number): Item { return f.toString(); },
        integer(i: number): Item { return i.toString(); },
        string(s: string): Item { return stringSource(s); },
        bytes(b: Bytes): Item {
            return seq(`Uint8Array.from(`, brackets(... Array.from(b).map(b => b.toString())), `)`);
        },
        symbol(s: symbol): Item { return `Symbol.for(${stringSource(s.description!)})`; },

        record(r: Record<Value<any>, Tuple<Value<any>>, any>, k: Fold<any, Item>): Item {
            return seq(`_.Record<_val, _.Tuple<_val>, _ptr>`, parens(k(r.label), brackets(... r.map(k))));
        },
        array(a: Array<Value<any>>, k: Fold<any, Item>): Item {
            return brackets(... a.map(k));
        },
        set(s: Set<any>, k: Fold<any, Item>): Item {
            return seq('new _.Set<_val>', parens(brackets(... Array.from(s).map(k))));
        },
        dictionary(d: Dictionary<any>, k: Fold<any, Item>): Item {
            return seq('new _.Dictionary<_ptr>', parens(brackets(... Array.from(d).map(([kk,vv]) =>
                brackets(k(kk), k(vv))))));
        },

        annotated(a: Annotated<any>, k: Fold<any, Item>): Item {
            return seq('_.annotate<_ptr>', parens(k(a.item), ... a.annotations.map(k)));
        },

        pointer(t: any, _k: Fold<any, Item>): Item {
            throw new Error(preserves`Cannot emit source code for construction of pointer ${t}`);
        },
    });
}