zoo2b.ts

// Normalization by evaluation
//
// Inspired Andras Kovacs' by elaboration-zoo  typecheck-HOAS-names
//
// By inspired, I mean that I wrote this after studing Kovacs' stuff and then
// cheated and looked up some of the answers when I got to the infer/check bit.
// Any bugs are mine.
//
// In this version, instead of passing in each branch as an argument, we pass
// a single object with a named function for each constructor.  The case statements
// are better labelled this way.

const fail = (msg: string) => { throw new Error(msg); };

// Term is both Raw and Term, but let never gets quoted back.
type Term = <M>(_: {
  v(a: string): M;
  lam(nm: string, body: Term): M;
  app(t: Term, u: Term): M;
  tlet(nm: string, a: Term, b: Term, c: Term): M;
  tpi(nm: string, a: Term, b: Term): M;
  tu(): M;
}) => M;

// constructors
let tvar = (n: string): Term => (x) => x.v(n);
let tlam = (n: string, t: Term): Term => (x) => x.lam(n, t);
let tapp = (t: Term, u: Term): Term => (x) => x.app(t, u);
let tlet = (n: string, a: Term, t: Term, u: Term): Term => (x) => x.tlet(n, a, t, u);
let tpi = (n: string, ty: Term, sc: Term): Term => (x) => x.tpi(n, ty, sc);
let tu = (): Term => (x) => x.tu();

type Env = [string, Val][];
let lookup = (env: Env, name: string) => env.find((x) => x[0] == name)?.[1];

type Closure = (_: Val) => Val;

type Val = <C>(_: {
  nvar(_: string): C;
  napp(t: Val, u: Val): C;
  vlam(sc: Closure): C;
  vpi(a: Val, b: Closure): C;
  vu(): C;
}) => C;

// constructors
let nvar = (n: string): Val => (x) => x.nvar(n);
let napp = (t: Val, u: Val): Val => (x) => x.napp(t, u);
let vlam = (sc: Closure): Val => (x) => x.vlam(sc);
let vpi = (ty: Val, sc: Closure): Val => (x) => x.vpi(ty, sc);
let vu = (): Val => (x) => x.vu();

let vapp = (t: Val, u: Val): Val =>
  t({
    nvar: () => napp(t, u),
    napp: () => napp(t, u),
    vlam: (sc) => sc(u),
    vpi: (_, sc) => sc(u),
    vu: () => napp(t, u),
  });

let ev = (env: Env, tm: Term): Val =>
  tm<Val>({
    v: (name) => lookup(env, name) || nvar(name),
    lam: (n, sc) => vlam((v) => ev([[n, v], ...env], sc)),
    app: (t, u) => vapp(ev(env, t), ev(env, u)),
    tlet: (n, ty, t, sc) => ev([[n, ev(env, t)], ...env], sc),
    // Hold onto ty so we can quote it back
    tpi: (n, ty, sc) => vpi(ev(env, ty), (v) => ev([[n, v], ...env], sc)),
    tu: () => vu(),
  });

let quote = (l: number, v: Val): Term =>
  v({
    nvar: (n) => tvar(n),
    napp: (t, u) => tapp(quote(l, t), quote(l, u)),
    vlam: (sc) => tlam("_" + l, quote(l + 1, sc(nvar("_" + l)))),
    vpi: (ty, sc) =>
      tpi("_" + l, quote(l, ty), quote(l + 1, sc(nvar("_" + l)))),
    vu: () => tu(),
  });

let nf = (t: Term): Term => quote(0, ev([], t));

// same shape, but the Val is the type
type Ctx = Env;

let conv = (env: Env, ty1: Val, ty2: Val): boolean => {
  let push = (x: string): Env => [[x, nvar(x)], ...env];
  let no = () => false;
  // the u, VLam case
  let eta = (sc: Closure) =>
    fresh(env, (x) => conv(push(x), vapp(ty1, nvar(x)), sc(nvar(x))));
  return ty1({
    nvar: (n) =>
      ty2({
        nvar: (n2) => n == n2,
        napp: no,
        vlam: eta,
        vpi: no,
        vu: no,
      }),
    napp: (t, u) =>
      ty2({
        nvar: no,
        napp: (t2, u2) => conv(env, t, t2) && conv(env, u, u2),
        vlam: eta,
        vpi: no,
        vu: no,
      }),
    vlam: (sc) =>
      fresh(env, (x) => {
        let eta2 = () => conv(push(x), sc(nvar(x)), vapp(ty2, nvar(x)));
        return ty2({
          nvar: eta2,
          napp: eta2,
          vlam: (sc2) => conv(push(x), sc(nvar(x)), sc2(nvar(x))),
          vpi: eta2,
          vu: eta2,
        });
      }),
    vpi: (a, b) =>
      ty2({
        nvar: no,
        napp: no,
        vlam: eta,
        vpi: (a2, b2) =>
          conv(env, a, a2) &&
          fresh(env, (x) => conv(push(x), b(nvar(x)), b2(nvar(x)))),
        vu: no,
      }),
    vu: () => ty2({ nvar: no, napp: no, vlam: eta, vpi: no, vu: () => true }),
  });
};

let notpi = (ty: Val) => () => fail(`expected pi type, got ${showVal(ty)}`);

let fresh = <A>(e: Env, f: (_: string) => A): A => f(`__${e.length}`);

// v1 - we'll just check/infer.  Later, let's return an elaborated term
let check = (env: Env, ctx: Ctx, t: Term, ty: Val): unknown =>
  t<unknown>({
    v: () =>
      conv(env, infer(env, ctx, t), ty) ||
      fail(`conv ${showVal(infer(env, ctx, t))} ${showVal(ty)}`),
    lam: (n, sc) =>
      ty({
        nvar: notpi(ty), // ezoo does check/infer to throw
        napp: notpi(ty),
        vlam: notpi(ty),
        vpi: (a, b) =>
          fresh(env, (x) =>
            check([[n, nvar(x)], ...env], [[n, a], ...ctx], sc, b(nvar(x)))
          ),
        vu: notpi(ty)
      }), // lam
    app: () => conv(env, infer(env, ctx, t), ty), // app
    tlet: (n, a, t, u) => {
      check(env, ctx, a, vu());
      let va = ev(env, a);
      check(env, ctx, t, va);
      check([[n, ev(env, t)], ...env], [[n, va], ...ctx], u, ty);
    }, // let
    tpi: () => conv(env, infer(env, ctx, t), ty), // pi
    tu: () => conv(env, infer(env, ctx, t), ty)
  });

let infer = (env: Env, ctx: Ctx, t: Term): Val =>
  t({
    v: (n) => lookup(ctx, n) || fail(`undefined ${n}`),
    lam: (n, t) => fail(`can't infer lambda`),
    app: (t, u) => {
      let tty = infer(env, ctx, t);
      return tty({
        nvar: notpi(tty),
        napp: notpi(tty),
        vlam: notpi(tty),
        vpi: (a, b) => (check(env, ctx, u, a), b(ev(env, u))),
        vu: notpi(tty)
      });
    },
    tlet: (n, a, t, u) => {
      check(env, ctx, a, vu());
      let va = ev(env, a);
      check(env, ctx, t, va);
      return infer([[n, ev(env, t)], ...env], [[n, va], ...ctx], u);
    },
    tpi: (n, a, b) => {
      check(env, ctx, a, vu());
      check([[n, nvar(n)], ...env], [[n, ev(env, a)], ...ctx], b, vu());
      return vu();
    }, // pi
    tu: () => vu()
  });

let showTerm = (t: Term): string =>
  t({
    v: (n) => n,
    lam: (n, sc) => `(λ ${n}. ${showTerm(sc)})`,
    app: (t, u) => `(${showTerm(t)} ${showTerm(u)})`,
    tlet: (n, a, t, u) =>
      `(let ${n} : ${showTerm(a)} = ${showTerm(t)}; ${showTerm(u)})`,
    tpi: (n, a, b) => `(∏(${n} : ${showTerm(a)}). ${showTerm(b)})`,
    tu: () => "U"
  });
let showVal = (v: Val): string => showTerm(quote(0, v));
// using Π a : A . B for telescopes to keep the parser simple
let parse = (x: string): Term => {
  let toks = x.match(/\w+|[&|=]+|\S/g)!;
  let next = () => toks.shift()!;
  let must = (s: string) =>
    toks[0] == s ? next() : fail(`Expected ${s} got ${next()}`);
  let stop = ") ; λ . =".split(" ");
  return (function expr(prec: number): Term {
    let left: Term;
    let t = next();
    // C doesn't guarantee those arguments are evaluated in order, I hope JS does.
    if (t == "let")
      left = ((n, _1, t, _2, a, _3, sc) => tlet(n, t, a, sc))(
        next(),
        must(":"),
        expr(0),
        must("="),
        expr(0),
        must(";"),
        expr(0)
      );
    else if (t == "λ")
      left = ((n, _, sc) => tlam(n, sc))(next(), must("."), expr(0));
    else if (t == "U") left = tu();
    // the greek keymap and Opt-Sh-P give me two different characters
    else if (t == "Π" || t == "∏")
      left = ((n, _1, a, _2, b) => tpi(n, a, b))(
        next(),
        must(":"),
        expr(0),
        must("."),
        expr(0)
      );
    else if (t == "(") left = ((t, _) => t)(expr(0), must(")"));
    else left = tvar(t); // should check t is ident
    while (prec == 0 && toks[0] && !stop.includes(toks[0]))
      left = tapp(left, expr(1));
    return left;
  })(0);
};

function testnf(s: string) {
  let tm = parse(s);
  console.log(showTerm(tm));
  console.log(showTerm(nf(tm)));
}
testnf("λx.x y z");
testnf("(λx.y x) z");
testnf("(λz.λx.x z) x");

function testInfer(s: string) {
  console.log("parsing", s);
  let tm = parse(s);
  console.log("parsed", showTerm(tm));
  let ty = infer([], [], tm);
  console.log(showTerm(nf(tm)), ":", showVal(ty));
}

parse("Π Α : U . A");

// I'll probably want some sugar on the ∏ and maybe the λ
testnf(`
let id : ∏  A : U . ∏ x : A . A = λ A . λ x . x ;
let const : ∏ A : U . ∏ B : U . ∏ x : A . ∏ y : B. A = λ A . λ B . λ x . λ y . x;
let Nat : U = ∏ N : U . ∏ f : (∏ _ : N . N) . ∏ _ : N . N;
let five : Nat = λ N . λ s . λ z . s (s (s (s (s z))));
let add : ∏ _ : Nat . ∏ _ : Nat . Nat = λa.λb.λN.λs.λz. a N s (b N s z);
let mul : ∏ _ : Nat . ∏ _ : Nat . Nat = λa.λb.λN.λs.λz. a N (b N s) z;
let ten : Nat = add five five;
let hundred : Nat = mul ten ten;
let thousand : Nat = mul ten hundred;
thousand
`);

testInfer(`
let id : ∏  A : U . ∏ x : A . A = λ A . λ x . x ;
let const : ∏ A : U . ∏ B : U . ∏ x : A . ∏ y : B. A = λ A . λ B . λ x . λ y . x;
let Nat : U = ∏ N : U . ∏ f : (∏ _ : N . N) . ∏ _ : N . N;
let five : Nat = λ N . λ s . λ z . s (s (s (s (s z))));
let add : ∏ _ : Nat . ∏ _ : Nat . Nat = λa.λb.λN.λs.λz. a N s (b N s z);
let mul : ∏ _ : Nat . ∏ _ : Nat . Nat = λa.λb.λN.λs.λz. a N (b N s) z;
let ten : Nat = add five five;
let hundred : Nat = mul ten ten;
let thousand : Nat = mul ten hundred;
U
`);

export default 1;

Previous version with mogensen-scott encoding is at https://gist.github.com/dunhamsteve/1be0cbb346d75ee1be8f67d192e73234