Monomorphic Type Inference with Mutable Type Variable Generation from Types and Programming Languages, simplified based on https://github.com/hsk/simpletapl

inf1.ml

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

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

let new_tvar =
  let i = ref 0 in
  let f () = incr i; TVar !i in
  f

let rec infer env = function
| EVar s -> (List.assoc s env, [])
| EAbs(sparam, fbody) ->
  let tparam = new_tvar () in
  let (tret, constr) = infer ((sparam,tparam)::env) fbody in
  (TArrow(tparam,tret), constr)
| EApp(func, arg) ->
  let (tfunc,constr1) = infer env func in
  let (targ,constr2) = infer env arg in
  let tret = new_tvar () in
  let newconstr = [tfunc,TArrow(targ,tret)] in
  (tret, 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 constr e =
  let (t, constr1) = infer env e in
  let constr = unify (constr @ constr1) in
  let t = applysubst constr t in
  (t, constr)