Monomorphic Type Inference with Mutable Type Variable Generation and Unification from TaPL, https://github.com/hsk/simpletapl and https://github.com/tomprimozic/type-systems

inf2.ml

type ty =
| TVar of tvar ref
| TArrow of ty * ty
and tvar =
| Unbound of int
| Link of 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(ref @@ Unbound !i) in
  f

let rec occursin id = function
| TVar{contents=Unbound id1} -> id = id1
| TVar{contents=Link t} -> occursin id t
| TArrow(tparam, tret) -> occursin id tparam || occursin id tret

let rec unify t1 t2 = match t1, t2 with
| _, _ when t1 = t2 -> ()
| TArrow(tparam1, tret1), TArrow(tparam2, tret2) ->
  unify tparam1 tparam2;
  unify tret1 tret2
| TVar{contents=Link t1}, t2 | t1, TVar{contents=Link t2} -> unify t1 t2
| TVar({contents=Unbound id} as tvar), ty | ty, TVar({contents=Unbound id} as tvar) ->
  if occursin id ty then failwith "Unification failed due to occurs check"
  else tvar := Link ty

let rec match_fun_ty tfunc targ = match tfunc with
| TArrow(tparam,tret) -> unify tparam targ; tret
| TVar {contents=Link ty} -> match_fun_ty ty targ
| TVar ({contents=Unbound _} as tvar) ->
  let tparam = new_tvar () in
  let tret = new_tvar () in
  tvar := Link(TArrow(tparam,tret));
  unify tparam targ;
  tret

let rec typeof env = function
| EVar s -> List.assoc s env
| EAbs(sparam, fbody) ->
  let tparam = new_tvar () in
  let tret = typeof ((sparam,tparam)::env) fbody in
  TArrow(tparam,tret)
| EApp(func, arg) ->
  let tfunc = typeof env func in
  let targ = typeof env arg in
  match_fun_ty tfunc targ