module type S = sig
  type ('term,'lit,'a) t
  (** An access key for theories which have per-class data ['a] *)

  val create :
    ?pp:'a Fmt.printer ->
    name:string ->
    eq:('a -> 'a -> bool) ->
    merge:('a -> 'a -> 'a) ->
    unit ->
    ('term,'lit,'a) t
  (** Generative creation of keys for the given theory data.

      @param eq : Equality. This is used to optimize backtracking info.

      @param merge :
        [merge d1 d2] is called when merging classes with data [d1] and [d2]
        respectively. The theory should already have checked that the merge
        is compatible, and this produces the combined data  for terms in the
        merged class.
      @param name name of the theory which owns this data
      @param pp a printer for the data
  *)

  val equal : ('t,'lit,_) t -> ('t,'lit,_) t -> bool
  (** Checks if two keys are equal (generatively) *)

  val pp : _ t Fmt.printer
  (** Prints the name of the key. *)
end


(** Custom keys for theory data.
    This imitates the classic tricks for heterogeneous maps
    https://blog.janestreet.com/a-universal-type/

    It needs to form a commutative monoid where values are persistent so
    they can be restored during backtracking.
*)
module Key = struct
  module type KEY_IMPL = sig
    type term
    type lit
    type t
    val id : int
    val name : string
    val pp : t Fmt.printer
    val equal : t -> t -> bool
    val merge : t -> t -> t
    exception Store of t
  end

  type ('term,'lit,'a) t =
    (module KEY_IMPL with type term = 'term and type lit = 'lit and type t = 'a)

  let n_ = ref 0

  let create (type term)(type lit)(type d)
      ?(pp=fun out _ -> Fmt.string out "<opaque>")
      ~name ~eq ~merge () : (term,lit,d) t =
    let module K = struct
      type nonrec term = term
      type nonrec lit = lit
      type t = d
      let id = !n_
      let name = name
      let pp = pp
      let merge = merge
      let equal = eq
      exception Store of d
    end in
    incr n_;
    (module K)

  let[@inline] id
    : type term lit a. (term,lit,a) t -> int
    = fun (module K) -> K.id

  let[@inline] equal
    : type term lit a b. (term,lit,a) t -> (term,lit,b) t -> bool
    = fun (module K1) (module K2) -> K1.id = K2.id

  let pp
    : type term lit a. (term,lit,a) t Fmt.printer
    = fun out (module K) -> Fmt.string out K.name
end



(*
  (** Map for theory data associated with representatives *)
  module K_map = struct
    type 'a key = (term,lit,'a) Key.t
    type pair = Pair : 'a key * exn -> pair

    type t = pair IM.t

    let empty = IM.empty

    let[@inline] mem k t = IM.mem (Key.id k) t

    let find (type a) (k : a key) (self:t) : a option =
      let (module K) = k in
      match IM.find K.id self with
      | Pair (_, K.Store v) -> Some v
      | _ -> None
      | exception Not_found -> None

    let add (type a) (k : a key) (v:a) (self:t) : t =
      let (module K) = k in
      IM.add K.id (Pair (k, K.Store v)) self
      
    let remove (type a) (k: a key) self : t =
      let (module K) = k in
      IM.remove K.id self

    let equal (m1:t) (m2:t) : bool =
      IM.equal
        (fun p1 p2 ->
           let Pair ((module K1), v1) = p1 in
           let Pair ((module K2), v2) = p2 in
           assert (K1.id = K2.id);
           match v1, v2 with K1.Store v1, K1.Store v2 -> K1.equal v1 v2 | _ -> false)
        m1 m2

    let merge ~f_both (m1:t) (m2:t) : t =
      IM.merge
        (fun _ p1 p2 ->
           match p1, p2 with
           | None, None -> None
           | Some v, None
           | None, Some v -> Some v
           | Some (Pair ((module K1) as key1, pair1)), Some (Pair (_, pair2)) ->
             match pair1, pair2 with
             | K1.Store v1, K1.Store v2 ->
               f_both K1.id pair1 pair2;  (* callback for checking compat *)
               let v12 = K1.merge v1 v2 in (* merge content *)
               Some (Pair (key1, K1.Store v12))
             | _ -> assert false
        )
        m1 m2
  end
   *)