(** Map of an id to a type *)
module IdentifierMap = Map.Make (Identifier)

(** Map type *)
type t = Type.t IdentifierMap.t

(* Empty substitution *)
let empty = IdentifierMap.empty

(** Create a substitution with one element *)
let singleton id ty = IdentifierMap.singleton id ty

(** Apply substitution to a type *)
let rec apply subst = function
  | Type.Int -> Type.Int
  | Type.Var id as t ->
    (* Look for a substitution in the map *)
    (match IdentifierMap.find_opt id subst with
     | Some ty' -> apply subst ty'
     | None -> t)
  | Type.Product (ty1, ty2) -> Type.Product (apply subst ty1, apply subst ty2)
  | Type.Arrow (ty1, ty2) -> Type.Arrow (apply subst ty1, apply subst ty2)
;;

(** Compose two substitutions *)
let compose s2 s1 =
  IdentifierMap.merge
    (fun _ ty1 ty2 ->
      match ty1, ty2 with
      (* If we have 2, we pick one of them *)
      | Some ty1', Some _ -> Some (apply s2 ty1')
      (* If we have 1, we pick the one we have *)
      | Some ty1', None -> Some (apply s2 ty1')
      | None, Some ty2' -> Some (apply s2 ty2')
      (* If we have 0, we return nothing *)
      | None, None -> None)
    s1
    s2
;;