thm/main.ts

226 lines
6.2 KiB
TypeScript
Raw Normal View History

2024-08-01 00:57:35 +00:00
import isEqual from "lodash.isequal";
import { match, P } from "ts-pattern";
export type expr =
| { kind: "var"; x: variable }
| { kind: "universe"; level: number }
| { kind: "pi"; abs: abstraction }
| { kind: "lambda"; abs: abstraction }
| { kind: "app"; left: expr; right: expr };
export interface abstraction {
x: variable;
type: expr;
body: expr;
}
// =============================================================================
// Substitution
export type variable =
| { kind: "string"; name: string }
| { kind: "gensym"; name: string; n: number }
| { kind: "dummy" };
export const dummy = (): variable => ({ kind: "dummy" });
export interface Substitution {
x: variable;
repl: expr;
}
let ctr = 0;
export function refresh(v: variable): variable {
const n = ctr + 1;
ctr += 1;
switch (v.kind) {
case "gensym":
case "string": {
return { kind: "gensym", name: v.name, n };
}
case "dummy":
return { kind: "gensym", name: "_", n };
}
}
export function subst(s: Substitution[], e: expr): expr {
switch (e.kind) {
case "var":
return s.find((sub) => isEqual(sub.x, e.x))?.repl ?? e;
case "universe":
return e;
case "pi":
return { kind: "pi", abs: substAbstraction(s, e.abs) };
case "lambda":
return { kind: "lambda", abs: substAbstraction(s, e.abs) };
case "app":
return { kind: "app", left: subst(s, e.left), right: subst(s, e.right) };
}
}
export function substAbstraction(
s: Substitution[],
{ x, type, body }: abstraction,
): abstraction {
const x2 = refresh(x);
return {
x: x,
type: subst(s, type),
body: subst([{ x, repl: { kind: "var", x: x2 } }, ...s], body),
};
}
// =============================================================================
// Type inference
export type context = ContextEntry[];
export interface ContextEntry {
x: variable;
type: expr;
value?: expr | undefined;
}
/** Returns the type of the definition associated with `x` inside the context `ctx` */
export function lookupTy(x: variable, ctx: context): expr | null {
return ctx.find((entry) => isEqual(entry.x, x))?.type ?? null;
}
/** Returns the value of the definition associated with `x` inside the context `ctx` */
export function lookupValue(x: variable, ctx: context): expr | null {
return ctx.find((entry) => isEqual(entry.x, x))?.value ?? null;
}
export function extend(
x: variable,
t: expr,
value: expr | undefined,
ctx: context,
): context {
return [...ctx, { x: x, type: t, value }];
}
/** Infer the type of the expression `e` in context `ctx` */
export function inferType(ctx: context, e: expr): expr {
switch (e.kind) {
case "var": {
const ty = lookupTy(e.x, ctx);
if (ty === null) throw new Error(`unknown identifier ${e.x}`);
return ty;
}
case "universe":
return { kind: "universe", level: e.level + 1 };
case "pi": {
const { x, type, body } = e.abs;
const k1 = inferUniverse(ctx, type);
const k2 = inferUniverse(extend(x, type, undefined, ctx), e.abs.body);
return { kind: "universe", level: Math.max(k1, k2) };
}
case "lambda": {
const { x, type, body } = e.abs;
const te = inferType(extend(x, type, undefined, ctx), body);
return { kind: "pi", abs: { x, type, body: te } };
}
case "app": {
const { x, type, body } = inferPi(ctx, e.left);
const te = inferType(ctx, e.right);
checkEqual(ctx, type, te);
return subst([{ x, repl: e.right }], body);
}
}
}
export function inferUniverse(ctx: context, t: expr): number {
const u = inferType(ctx, t);
const n = normalize(ctx, u);
switch (n.kind) {
case "universe":
return n.level;
default:
throw new Error("type expected.");
}
}
export function inferPi(ctx: context, e: expr): abstraction {
const t = inferType(ctx, e);
const n = normalize(ctx, t);
switch (n.kind) {
case "pi":
return n.abs;
default:
throw new Error("function expected");
}
}
export function checkEqual(ctx: context, e1: expr, e2: expr) {
if (!equal(ctx, e1, e2))
throw new Error(`expressions ${e1} and ${e2} not equal`);
}
// =============================================================================
// Normalization and equality
export function normalize(ctx: context, e: expr): expr {
switch (e.kind) {
case "var": {
const n = lookupTy(e.x, ctx);
if (!n) throw new Error(`unknown identifier ${e.x}`);
return n;
}
case "app": {
const e2 = normalize(ctx, e.right);
const e1 = normalize(ctx, e.left);
switch (e1.kind) {
case "lambda":
return normalize(
ctx,
subst([{ x: e1.abs.x, repl: e2 }], e1.abs.body),
);
default:
return { kind: "app", left: e1, right: e2 };
}
}
case "universe":
return e;
case "pi":
return { kind: "pi", abs: normalizeAbstraction(ctx, e.abs) };
case "lambda":
return { kind: "lambda", abs: normalizeAbstraction(ctx, e.abs) };
}
}
export function normalizeAbstraction(
ctx: context,
{ x, type, body }: abstraction,
): abstraction {
const t = normalize(ctx, type);
return { x, type: t, body: normalize(extend(x, t, undefined, ctx), body) };
}
export function equal(ctx: context, e1: expr, e2: expr): boolean {
function equalHelper(e1: expr, e2: expr): boolean {
if (e1.kind === "var" && e2.kind === "var") return isEqual(e1.x, e2.x);
if (e1.kind === "app" && e2.kind === "app")
return equalHelper(e1.left, e2.left) && equalHelper(e1.right, e2.right);
if (e1.kind === "universe" && e2.kind === "universe")
return e1.level === e2.level;
if (e1.kind === "pi" && e2.kind === "pi")
return equalAbstraction(e1.abs, e2.abs);
if (e1.kind === "lambda" && e2.kind === "lambda")
return equalAbstraction(e1.abs, e2.abs);
return false;
}
function equalAbstraction(
{ x, type: type1, body: body1 }: abstraction,
{ x: y, type: type2, body: body2 }: abstraction,
): boolean {
return (
equalHelper(type1, type2) &&
equalHelper(body1, subst([{ x: y, repl: { kind: "var", x } }], body2))
);
}
return equalHelper(normalize(ctx, e1), normalize(ctx, e2));
}