type expr =
  | EUnit (* () *)
  | EVar of string (* x *)
  | EAbs of string * expr (* λx.e *)
  | EApp of expr * expr (* e.e *)
  | EAnnot of expr * ty (* e:A *)

and ty =
  | TUnit (* () *)
  | TFun of ty * ty (* A->A *)

type err =
  | UnboundName of string
  | CantApplyNonFunc
  | NeedsAnnotation of string

type ctx = (string * ty) list

let rec infer (ctx: ctx) (expr: expr): (ty, err) result =
  match expr with
  | EUnit -> Ok TUnit
  | EVar name ->
      List.find_opt (fun (n, _) -> n = name) ctx
      |> Option.to_result ~none:(UnboundName name)
      |> Result.map snd
  | EAnnot (_, ty) -> Ok ty
  | EApp (func, arg) ->
      begin match infer ctx func with
      | Ok (TFun (a, b)) -> check ctx arg a |> Result.map (Fun.const b)
      | Ok _ -> Error CantApplyNonFunc
      | Error e -> Error e
      end
  | EAbs (name, _) -> Error (NeedsAnnotation name)

and check (ctx: ctx) (expr: expr) (ty: ty): (unit, err) result =
  match (expr, ty) with
  | (EAbs (name, body), TFun (argt, bodyt)) -> check ((name, argt) :: ctx) body bodyt
  | _ -> infer ctx expr |> Result.map (Fun.const ())

(* vim: set sw=2 : *)