(** Unify 2 types and, if exists, returns the substitution *)
let rec unify ty1 ty2 =
  match ty1, ty2 with
  (* Same types *)
  | Type.Int, Type.Int -> Some TypeSubstitution.empty
  | Type.Var id1, Type.Var id2 when id1 = id2 -> Some TypeSubstitution.empty
  (* Different types *)
  | Type.Var id, ty | ty, Type.Var id -> Some (TypeSubstitution.singleton id ty)
  | Type.Product (l1, r1), Type.Product (l2, r2) | Type.Arrow (l1, r1), Type.Arrow (l2, r2)
    ->
    (match unify l1 l2 with
     | Some s1 ->
       (* Apply the substitution *)
       (match unify (TypeSubstitution.apply s1 r1) (TypeSubstitution.apply s1 r2) with
        | Some s2 -> Some (TypeSubstitution.compose s2 s1)
        | None -> None)
     | None -> None)
  | _, _ -> None
;;