lstd/jit.ts

export const jit_format = Symbol.for(`re.lua.lstd.jit_format`);

export interface JitFragment {
  [jit_format](c: JitCompiler): void;
}

export function is_jit(x: unknown): x is JitFragment {
  return typeof x === "object" && x !== null && jit_format in x;
}

export class JitCompiler {
  #body = "";
  #args = new Map<unknown, string>();

  write(s: string | JitFragment) {
    if (is_jit(s)) s[jit_format](this);
    else this.#body += s;
  }

  enclose(x: unknown) {
    switch (typeof x) {
      case "boolean":
      case "number":
        return this.write(`${x}`);
      case "bigint":
        return this.write(`${x}n`);
      case "string":
        return this.write(JSON.stringify(x));
      case "object":
        if (x === null) return this.write(`null`);
        break;
      case "undefined":
        return this.write(`undefined`);
    }

    let name = this.#args.get(x);
    if (!name) this.#args.set(x, (name = `__x${this.#args.size}`));
    this.write(name);
  }

  format(x: unknown) {
    is_jit(x) ? this.write(x) : this.enclose(x);
  }

  compile() {
    return new Function(
      ...this.#args.values(),
      `"use strict"; return (${this.#body});`
    )(...this.#args.keys());
  }
}

export function jit(
  { raw: s }: TemplateStringsArray,
  ...xs: unknown[]
): JitFragment {
  return {
    [jit_format](c) {
      for (let i = 0, n = s.length; i < n; i++) {
        if (i !== 0) c.format(xs[i - 1]);
        c.write(s[i]);
      }
    },
  };
}

jit.compiled = compiled;
jit.raw = raw;
jit.fragment = fragment;
jit.map = map;
jit.if = condition;

export function compiled<T = unknown>(
  s: TemplateStringsArray,
  ...xs: unknown[]
): T {
  const c = new JitCompiler();
  return c.write(jit(s, ...xs)), c.compile();
}

export function raw(s: TemplateStringsArray, ...xs: unknown[]): JitFragment;
export function raw(s: string): JitFragment;
export function raw(
  s: TemplateStringsArray | string,
  ...xs: unknown[]
): JitFragment {
  s = typeof s === "string" ? s : String.raw(s, ...xs);
  return {
    [jit_format](c) {
      c.write(s);
    },
  };
}

export function fragment(
  sep: string | JitFragment,
  ...xs: unknown[]
): JitFragment {
  return {
    [jit_format](c) {
      for (let i = 0, n = xs.length; i < n; i++) {
        if (i !== 0) c.write(sep);
        c.format(xs[i]);
      }
    },
  };
}

export function map<T>(
  sep: string | JitFragment,
  xs: Iterable<T>,
  f: (value: T, index: number) => unknown
): JitFragment {
  return fragment(sep, ...Iterator.from(xs).map(f));
}

export function condition(
  test: unknown,
  consequent: JitFragment,
  alternate: JitFragment = jit``
): JitFragment {
  return test ? consequent : alternate;
}