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

inf3.ml

type ty =
| TVar of tvar ref
| TArrow of ty * ty
and tvar =
| Unbound of int
| Link of ty
| Generic of int

type exp =
| EVar of string
| EAbs of string * exp
| EApp of exp * exp
| ELet of string * 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
| TVar{contents=Generic _} -> false
| 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
| _ -> failwith "cannot unify types"

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
| TVar {contents=Generic _} -> failwith "Generic type encountered, expecting arrow or instantiated variable"

let rec generalize env_tys ty = match ty with
| TVar{contents=Unbound id} when not (List.exists (occursin id) env_tys) -> TVar (ref (Generic id))
| TVar{contents=Unbound _} -> ty
| TVar{contents=Link ty} -> generalize env_tys ty
| TVar{contents=Generic _} -> ty
| TArrow(tparam, tret) -> TArrow(generalize env_tys tparam, generalize env_tys tret)

let instantiate ty =
  let id_var_hash = Hashtbl.create 10 in
  let rec f ty = match ty with
  | TVar{contents=Generic id} ->
    (try Hashtbl.find id_var_hash id
     with Not_found ->
       let var = new_tvar () in
        Hashtbl.add id_var_hash id var;
        var)
  | TVar{contents=Unbound _} -> ty
  | TVar{contents=Link ty} -> f ty
  | TArrow(tparam, tret) -> TArrow(f tparam, f tret)
  in f ty

let rec typeof env = function
| EVar s -> instantiate (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
| ELet(svar, e, ebody) ->
  let tvar = typeof env e in
  let tgen = generalize (List.map snd env) tvar in
  typeof ((svar,tgen)::env) ebody