thm/main.ts

import isEqual from "lodash.isequal";

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 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: x2,
    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 {
  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 {
  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(
  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 {
  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 {
  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 {
  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) };
    }
  }

  const result = normalize(ctx, e);
  return result;
}

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));
}

// =============================================================================
// 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();