Skip to content

Instantly share code, notes, and snippets.

@dunhamsteve

dunhamsteve/zoo2.ts

Created November 20, 2023 04:02
Show Gist options
  • Save dunhamsteve/1be0cbb346d75ee1be8f67d192e73234 to your computer and use it in GitHub Desktop.
Save dunhamsteve/1be0cbb346d75ee1be8f67d192e73234 to your computer and use it in GitHub Desktop.
Download ZIP
Normalization by evaluation in typescript
Raw
zoo2.ts
// 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
`);
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment