From 102b5efc465fbc34211b3888b8df22eba9057305 Mon Sep 17 00:00:00 2001
From: Michael Zhang <mail@mzhang.io>
Date: Wed, 31 Jul 2024 20:48:49 -0500
Subject: [PATCH] it works

---
 main.ts | 274 +++++++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 220 insertions(+), 54 deletions(-)

diff --git a/main.ts b/main.ts
index 74a0369..ef5ef47 100644
--- a/main.ts
+++ b/main.ts
@@ -1,5 +1,4 @@
 import isEqual from "lodash.isequal";
-import { match, P } from "ts-pattern";
 
 export type expr =
   | { kind: "var"; x: variable }
@@ -22,8 +21,6 @@ export type variable =
   | { kind: "gensym"; name: string; n: number }
   | { kind: "dummy" };
 
-export const dummy = (): variable => ({ kind: "dummy" });
-
 export interface Substitution {
   x: variable;
   repl: expr;
@@ -64,7 +61,7 @@ export function substAbstraction(
 ): abstraction {
   const x2 = refresh(x);
   return {
-    x: x,
+    x: x2,
     type: subst(s, type),
     body: subst([{ x, repl: { kind: "var", x: x2 } }, ...s], body),
   };
@@ -82,13 +79,23 @@ export interface ContextEntry {
 }
 
 /** 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;
+export function lookupTy(x: variable, ctx: context): expr {
+  const res = ctx.find((entry) => isEqual(entry.x, x));
+  if (res === undefined)
+    throw new Error(
+      `unknown identifier ${ppVariable(x)} (ctx: ${ppContext(ctx)})`,
+    );
+  return res.type;
 }
 
 /** 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;
+  const res = ctx.find((entry) => isEqual(entry.x, x));
+  if (res === undefined)
+    throw new Error(
+      `unknown identifier ${ppVariable(x)} (ctx: ${ppContext(ctx)})`,
+    );
+  return res.value ?? null;
 }
 
 export function extend(
@@ -102,32 +109,37 @@ export function extend(
 
 /** 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);
+  function inferType(ctx: context, e: expr): expr {
+    switch (e.kind) {
+      case "var": {
+        const ty = lookupTy(e.x, ctx);
+        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), body);
+        return { kind: "universe", level: Math.max(k1, k2) };
+      }
+      case "lambda": {
+        const { x, type, body } = e.abs;
+        inferUniverse(ctx, type);
+        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);
+      }
     }
   }
+
+  const result = inferType(ctx, e);
+  return result;
 }
 
 export function inferUniverse(ctx: context, t: expr): number {
@@ -161,32 +173,35 @@ export function checkEqual(ctx: context, e1: expr, e2: expr) {
 // 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 };
+  function normalize(ctx: context, e: expr): expr {
+    switch (e.kind) {
+      case "var": {
+        const n = lookupValue(e.x, ctx);
+        return n ? normalize(ctx, n) : e;
       }
+      case "app": {
+        const e1 = normalize(ctx, e.left);
+        const e2 = normalize(ctx, e.right);
+        switch (e1.kind) {
+          case "lambda": {
+            const { x, body } = e1.abs;
+            return normalize(ctx, subst([{ x, repl: e2 }], body));
+          }
+          default:
+            return { kind: "app", left: e1, right: e2 };
+        }
+      }
+      case "universe":
+        return { kind: "universe", level: e.level };
+      case "pi":
+        return { kind: "pi", abs: normalizeAbstraction(ctx, e.abs) };
+      case "lambda":
+        return { kind: "lambda", abs: normalizeAbstraction(ctx, e.abs) };
     }
-    case "universe":
-      return e;
-    case "pi":
-      return { kind: "pi", abs: normalizeAbstraction(ctx, e.abs) };
-    case "lambda":
-      return { kind: "lambda", abs: normalizeAbstraction(ctx, e.abs) };
   }
+
+  const result = normalize(ctx, e);
+  return result;
 }
 
 export function normalizeAbstraction(
@@ -223,3 +238,154 @@ export function equal(ctx: context, e1: expr, e2: expr): boolean {
 
   return equalHelper(normalize(ctx, e1), normalize(ctx, e2));
 }
+
+// =============================================================================
+// Pretty printing
+
+export function ppVariable(v: variable): string {
+  switch (v.kind) {
+    case "string":
+      return `\`${v.name}`;
+    case "gensym":
+      return `\`${v.name}$${v.n}`;
+    case "dummy":
+      return "`_";
+  }
+}
+
+function occursIn(x: variable, e: expr): boolean {
+  switch (e.kind) {
+    case "var":
+      return isEqual(e.x, x);
+    case "universe":
+      return false;
+  }
+}
+
+export function ppExpr(e: expr): string {
+  switch (e.kind) {
+    case "var":
+      return ppVariable(e.x);
+    case "universe":
+      if (e.level === 0) return "U";
+      return `U(${e.level})`;
+    case "pi":
+      if (occursIn(e.abs.x, e.abs.body)) {
+        return `(∏ (${ppVariable(e.abs.x)} : ${ppExpr(e.abs.type)}), ${ppExpr(
+          e.abs.type,
+        )})`;
+      }
+      return `(${ppExpr(e.abs.type)} → ${ppExpr(e.abs.body)})`;
+    case "lambda":
+      return `(λ (${ppVariable(e.abs.x)} : ${ppExpr(e.abs.type)}), ${ppExpr(
+        e.abs.body,
+      )})`;
+    case "app":
+      return `(${ppExpr(e.left)} ${ppExpr(e.right)})`;
+  }
+}
+
+export function ppContext(ctx: context): string {
+  const ss = ctx.map(
+    ({ x, type, value }) =>
+      `${ppVariable(x)} : ${ppExpr(type)}${value ? ` = ${ppExpr(value)}` : ""}`,
+  );
+
+  return `[${ss.join(", ")}]`;
+}
+
+// =============================================================================
+// Testing
+
+function test() {
+  const ctx: context = [];
+
+  const dummy: variable = { kind: "dummy" };
+  const varify = (name: string): expr => ({
+    kind: "var",
+    x: { kind: "string", name },
+  });
+  const define = (name: string, e: expr) => {
+    const type = inferType(ctx, e);
+    ctx.push({ x: { kind: "string", name }, type, value: e });
+  };
+
+  // N : Set lzero
+  ctx.push({
+    x: { kind: "string", name: "N" },
+    type: { kind: "universe", level: 0 },
+  });
+  const N = varify("N");
+
+  // z : N
+  ctx.push({ x: { kind: "string", name: "z" }, type: N });
+  const z = varify("z");
+
+  // s : N -> N
+  const NtoN: expr = { kind: "pi", abs: { x: dummy, type: N, body: N } };
+  ctx.push({
+    x: { kind: "string", name: "s" },
+    type: NtoN,
+  });
+  const s = varify("s");
+
+  const f = varify("f");
+  const x = varify("x");
+  const threeExpr: expr = {
+    kind: "lambda",
+    abs: {
+      x: { kind: "string", name: "f" },
+      type: NtoN,
+      body: {
+        kind: "lambda",
+        abs: {
+          x: { kind: "string", name: "x" },
+          type: N,
+          body: {
+            kind: "app",
+            left: f,
+            right: {
+              kind: "app",
+              left: f,
+              right: { kind: "app", left: f, right: x },
+            },
+          },
+        },
+      },
+    },
+  };
+  const three = varify("three");
+  define("three", threeExpr);
+
+  const nineExpr: expr = {
+    kind: "app",
+    left: three,
+    right: {
+      kind: "app",
+      left: three,
+      right: s,
+    },
+  };
+  const nine = varify("nine");
+  define("nine", nineExpr);
+
+  const evaled: expr = {
+    kind: "app",
+    left: nine,
+    right: z,
+  };
+  define("result", normalize(ctx, evaled));
+
+  console.log("ctx =");
+  // biome-ignore lint/complexity/noForEach: <explanation>
+  ctx.forEach(({ x, type, value }, idx) => {
+    console.log(
+      `  ${idx === 0 ? "[" : ","} ${ppVariable(x)} : ${ppExpr(type)}${
+        value ? ` = ${ppExpr(value)}` : ""
+      }`,
+    );
+  });
+  console.log("  ]");
+}
+
+test();