bidirectional type-checker in TypeScript

bidirectional-structural.ts

// in response to https://twitter.com/arntzenius/status/1022076441603829760

//  types  A,B ∷= base | A → B
//  terms  M ∷= E | λx.M
//  exprs  E ∷= x | E M | (M : A)
type Type = "Base" | {lhs: Type, tag: "->", rhs: Type};
type Var = string;
type Term = Expr | {tag: "Lam", varname: Var, body: Term};
type Expr = Var | {tag: "App", fn: Expr, arg: Term} | {term: Term, tag: ":", tp: Type};

function typesAreEqual(tp1: Type, tp2: Type): boolean {
  if (tp1 === "Base" && tp2 === "Base") {
    return true;
  } else if (typeof(tp1) === "object" && typeof(tp2) === "object") {
    return typesAreEqual(tp1.lhs, tp2.lhs) && typesAreEqual(tp1.rhs, tp2.rhs);
  } else {
    return false;
  }
}

// helpers to make it easier to construct Terms and Exprs
function arr(lhs: Type, rhs: Type): Type {
  return {lhs: lhs, tag: "->", rhs: rhs};
}
function lam(varname: Var, body: (varname:Var) => Term): Term {
  return {tag: "Lam", varname: varname, body: body(varname)};
}
function app(fn: Expr, arg: Term): Expr {
  return {tag: "App", fn: fn, arg: arg};
}
function ofType(term: Term, tp: Type): Expr {
  return {term: term, tag: ":", tp: tp};
}
function lamFn(varname: Var, body: (fn: (arg:Expr) => Expr) => Term): Term {
  return lam(varname, f => body(arg => app(f, arg)));
}

// example value: the church encoding of the number two
let two: Term = lamFn("succ", succ => lam("zero", zero => succ(succ(zero))));

//type ctx = (var * tp) list
type List<A> = "Nil" | {head: A, tag: "Cons", tail: List<A>};
type Ctx = List<[Var, Type]>;

function assoc(varname: Var, ctx: Ctx) : Type | "ScopeError" {
  if (typeof(ctx) === "object" && ctx.tag === "Cons") {
    let [k,v] = ctx.head;
    if (k === varname) {
      return v;
    } else {
      return assoc(varname, ctx.tail);
    }
  } else {
    return "ScopeError";
  }
}

//val infer: ctx -> expr -> tp
//val check: ctx -> tp -> term -> unit
function infer(ctx: Ctx, expr: Expr): Type | "ScopeError" | "TypeError" {
  //  | `Var x -> List.assoc x ctx
  if (typeof(expr) === "string") {
    let varname: Var = expr;
    return assoc(varname, ctx);
  }

  //  | `App (e,m) -> let Fn(a,b) = infer ctx e in check ctx a m; b
  else if (expr.tag === "App") {
    let tp = infer(ctx, expr.fn);
    if (typeof(tp) === "object" && tp.tag === "->") {
      let argResult = check(ctx, tp.lhs, expr.arg);
      if (argResult === "Typechecks") {
        return tp.rhs;
      } else {
        return argResult;
      }
    } else {
      return "TypeError";
    }
  }

  //  | `Asc (a,m) -> check ctx a m; a
  else {
    let termResult = check(ctx, expr.tp, expr.term);
    if (termResult === "Typechecks") {
      return expr.tp;
    } else {
      return termResult;
    }
  }
}
function check(ctx: Ctx, tp: Type, term: Term): "Typechecks" | "ScopeError" | "TypeError" {
  //  | `Lam(x,m) -> let Fn(a,b) = tp in check ((x,a)::ctx) b m
  if (typeof(term) === "object" && term.tag == "Lam") {
    if (typeof(tp) === "object" && tp.tag === "->") {
      let extendedCtx: Ctx = {head: [term.varname, tp.lhs], tag: "Cons", tail: ctx};
      return check(extendedCtx, tp.rhs, term.body);
    } else {
      return "TypeError";
    }
  }

  //  | e ->
  //     let inferred = infer ctx e in
  //     if inferred = tp then () else raise TypeError
  else {
    let expr: Expr = term;
    let inferred = infer(ctx, expr);
    if (inferred === "ScopeError" || inferred === "TypeError") {
      return inferred;
    } else if (typesAreEqual(inferred, tp)) {
      return "Typechecks";
    } else {
      return "TypeError";
    }
  }
}


// tests
var tp = arr("Base", "Base");
console.log(check("Nil", tp, two)); // TypeError
tp = arr(tp, tp);
console.log(check("Nil", tp, two)); // TypeChecks
tp = arr(tp, tp);
console.log(check("Nil", tp, two)); // TypeChecks