Purely Functional Monomorphic Type Inference from Types and Programming Languages, simplified based on https://github.com/hsk/simpletapl

inf0.ml

type ty =
| TVar of int
| TArrow of ty * ty

type exp =
| EVar of string
| EAbs of string * exp
| EApp of exp * exp

type new_tvar = NewTVar of ty * (unit -> new_tvar)

let new_tvar =
  let rec f n () =
    NewTVar (TVar n, f (n+1))
  in f 0

let rec infer env new_tvar = function
| EVar s -> (List.assoc s env, new_tvar, [])
| EAbs(sparam, fbody) ->
  let NewTVar(tparam,new_tvar) = new_tvar () in
  let (tret, new_tvar, constr) = infer ((sparam,tparam)::env) new_tvar fbody in
  (TArrow(tparam,tret), new_tvar, constr)
| EApp(func, arg) ->
  let (tfunc,new_tvar,constr1) = infer env new_tvar func in
  let (targ,new_tvar,constr2) = infer env new_tvar arg in
  let NewTVar(tret,new_tvar) = new_tvar () in
  let newconstr = [tfunc,TArrow(targ,tret)] in
  (tret, new_tvar, List.concat [newconstr; constr1; constr2])

let rec substinty s t t1 = match t1 with
| TVar(s1) -> if s = s1 then t else t1
| TArrow(tparam1,tret1) -> TArrow(substinty s t tparam1, substinty s t tret1)

let applysubst constr t =
  let f t = function
  | (TVar(s),t1) -> substinty s t1 t
  | _ -> assert false
in List.fold_left f t (List.rev constr)

let substinconstr s t constr =
  List.map (fun (t1,t2) -> (substinty s t t1, substinty s t t2)) constr

let rec occursin s = function
| TVar s2 -> s = s2
| TArrow(tparam,tret) -> occursin s tparam || occursin s tret

let rec unify = function
| [] -> []
| (TVar s1, TVar s2)::rest when s1 = s2 -> unify rest
| (t, TVar s)::rest | (TVar s, t)::rest ->
  if occursin s t then failwith "Unification failed due to occurs check"
  else List.append (unify (substinconstr s t rest)) [(TVar s, t)]
| (TArrow(tparam1,tret1),TArrow(tparam2,tret2))::rest ->
  unify ((tparam1,tparam2)::(tret1,tret2)::rest)

let typeof env new_tvar constr e =
  let (t, new_tvar, constr1) = infer env new_tvar e in
  let constr = unify (constr @ constr1) in
  let t = applysubst constr t in
  (t, new_tvar, constr)