Skip to content

Instantly share code, notes, and snippets.

@dunhamsteve
Created November 20, 2023 04:02

Revisions

  1. dunhamsteve created this gist Nov 20, 2023.
    229 changes: 229 additions & 0 deletions zoo2.ts
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,229 @@
    // Normalization by evaluation
    //
    // I was playing around with mogensen-scott encoding of the STLC and got carried away.
    //
    // 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, though.

    // make function types easier to write
    type F0<A> = () => A;
    type F1<A, B> = (_: A) => B;
    type F2<A, B, C> = (a: A, b: B) => C;
    type F3<A, B, C, D> = (a: A, b: B, c: C) => D;
    type F4<A, B, C, D, E> = (a: A, b: B, c: C, d: D) => E;

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

    // Term is both Raw and Term, but let never gets quoted back.
    type Term = <M>(
    v: F1<string, M>,
    lam: F2<string, Term, M>,
    app: F2<Term, Term, M>,
    tlet: F4<string, Term, Term, Term, M>,
    tpi: F3<string, Term, Term, M>,
    tu: F0<M>
    ) => M;
    let tvar: F1<string, Term> = (n) => (fv, fl, fa) => fv(n);
    let tlam: F2<string, Term, Term> = (n, t) => (fv, fl, fa) => fl(n, t);
    let tapp: F2<Term, Term, Term> = (t, u) => (fv, fl, fa) => fa(t, u);
    let tlet: F4<string, Term, Term, Term, Term> = (n, a, t, u) => (fv, fl, fa, flet) => flet(n, a, t, u);
    let tpi: F3<string, Term, Term, Term> = (n, ty, sc) => (fv, fl, fa, flet, fpi) => fpi(n, ty, sc);
    let tu: Term = (fv, fl, fa, flet, fpi, fu) => fu();

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

    type Val = <C>(
    nvar: F1<string, C>,
    napp: F2<Val, Val, C>,
    vlam: F1<F1<Val, Val>, C>,
    vpi: F2<Val, F1<Val, Val>, C>,
    vu: F0<C>
    ) => C;
    let nvar: F1<string, Val> = (n) => (nv, na) => nv(n);
    let napp: F2<Val, Val, Val> = (t, u) => (nv, na) => na(t, u);
    let vlam: F1<F1<Val, Val>, Val> = (sc) => (nv, na, nl) => nl(sc);
    let vpi: F2<Val, F1<Val, Val>, Val> = (ty, sc) => (nv, na, nl, npi) => npi(ty, sc);
    let vu: Val = (nv, na, nl, npi, vu) => vu();

    let vapp = (t:Val, u: Val): Val =>
    t(
    () => napp(t, u),
    () => napp(t, u),
    (sc) => sc(u),
    (_, sc) => sc(u),
    () => napp(t, u)
    );

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

    let quote = (l: number, v: Val): Term =>
    v(
    (n) => tvar(n),
    (t, u) => tapp(quote(l, t), quote(l, u)),
    (sc) => tlam("_" + l, quote(l + 1, sc(nvar("_" + l)))),
    (ty, sc) => tpi("_" + l, quote(l, ty), quote(l + 1, sc(nvar("_" + l)))),
    () => 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: F1<Val, Val>) => fresh(env, (x) => conv(push(x), vapp(ty1, nvar(x)), sc(nvar(x))));
    return ty1(
    (n) => ty2((n2) => n == n2, no, eta, no, no),
    (t, u) => ty2(no, (t2, u2) => conv(env, t, t2) && conv(env, u, u2), eta, no, no),
    (sc) =>
    fresh(env, (x) => {
    let eta2 = () => conv(push(x), sc(nvar(x)), vapp(ty2, nvar(x)));
    return ty2(eta2, eta2, (sc2) => conv(push(x), sc(nvar(x)), sc2(nvar(x))), eta2, eta2);
    }),
    (a, b) => ty2(no, no, eta, (a2, b2) => conv(env, a, a2) && fresh(env, (x) => conv(push(x), b(nvar(x)), b2(nvar(x)))), no),
    () => ty2(no, no, eta, no, () => true)
    );
    };

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

    let fresh = <A>(e: Env, f: F1<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(
    () => conv(env, infer(env, ctx, t), ty) || fail(`conv ${showVal(infer(env, ctx, t))} ${showVal(ty)}`),
    (n, sc) =>
    ty(
    notpi(ty), // ezoo does check/infer to throw
    notpi(ty),
    notpi(ty),
    (a, b) => fresh(env, (x) => check([[n, nvar(x)], ...env], [[n, a], ...ctx], sc, b(nvar(x)))),
    notpi(ty)
    ), // lam
    () => conv(env, infer(env, ctx, t), ty), // app
    (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
    () => conv(env, infer(env, ctx, t), ty), // pi
    () => conv(env, infer(env, ctx, t), ty)
    );

    let infer = (env: Env, ctx: Ctx, t: Term): Val =>
    t(
    (n) => lookup(ctx, n) || fail(`undefined ${n}`),
    (n, t) => fail(`can't infer lambda`),
    (t, u) => {
    let tty = infer(env, ctx, t);
    return tty(notpi(tty), notpi(tty), notpi(tty), (a, b) => (check(env, ctx, u, a), b(ev(env, u))), notpi(tty));
    },
    (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);
    },
    (n, a, b) => {
    check(env, ctx, a, vu);
    check([[n, nvar(n)], ...env], [[n, ev(env, a)], ...ctx], b, vu);
    return vu;
    }, // pi
    () => vu
    );

    let showTerm = (t: Term): string =>
    t(
    (n) => n,
    (n, sc) => `(λ ${n}. ${showTerm(sc)})`,
    (t, u) => `(${showTerm(t)} ${showTerm(u)})`,
    (n, a, t, u) => `(let ${n} : ${showTerm(a)} = ${showTerm(t)}; ${showTerm(u)})`,
    (n, a, b) => `(∏(${n} : ${showTerm(a)}). ${showTerm(b)})`,
    () => "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
    `);