(* This file is free software, part of containers. See file "license" for more details. *)

(** {1 Associative containers with Heterogenerous Values} *)

(*$R
  let k1 : int Key.t = Key.create() in
  let k2 : int Key.t = Key.create() in
  let k3 : string Key.t = Key.create() in
  let k4 : float Key.t = Key.create() in

  let tbl = Tbl.create () in

  Tbl.add tbl k1 1;
  Tbl.add tbl k2 2;
  Tbl.add tbl k3 "k3";

  assert_equal (Some 1) (Tbl.find tbl k1);
  assert_equal (Some 2) (Tbl.find tbl k2);
  assert_equal (Some "k3") (Tbl.find tbl k3);
  assert_equal None (Tbl.find tbl k4);
  assert_equal 3 (Tbl.length tbl);

  Tbl.add tbl k1 10;
  assert_equal (Some 10) (Tbl.find tbl k1);
  assert_equal 3 (Tbl.length tbl);
  assert_equal None (Tbl.find tbl k4);

  Tbl.add tbl k4 0.0;
  assert_equal (Some 0.0) (Tbl.find tbl k4);

  ()
*)

type 'a iter = ('a -> unit) -> unit
type 'a gen = unit -> 'a option

module type KEY_IMPL = sig
  type t

  exception Store of t

  val id : int
end

module Key = struct
  type 'a t = (module KEY_IMPL with type t = 'a)

  let _n = ref 0

  let create (type k) () =
    incr _n;
    let id = !_n in
    let module K = struct
      type t = k

      let id = id

      exception Store of k
    end in
    (module K : KEY_IMPL with type t = k)

  let id (type k) (module K : KEY_IMPL with type t = k) = K.id

  let equal : type a b. a t -> b t -> bool =
   fun (module K1) (module K2) -> K1.id = K2.id
end

type pair = Pair : 'a Key.t * 'a -> pair
type exn_pair = E_pair : 'a Key.t * exn -> exn_pair

let pair_of_e_pair (E_pair (k, e)) =
  let module K = (val k) in
  match e with
  | K.Store v -> Pair (k, v)
  | _ -> assert false

module Map = struct
  module M = Map.Make (struct
    type t = int

    let compare (i : int) j = CCInt.compare i j
  end)

  type t = exn_pair M.t

  let empty = M.empty
  let mem k t = M.mem (Key.id k) t

  let find_exn (type a) (k : a Key.t) t : a =
    let module K = (val k) in
    let (E_pair (_, e)) = M.find K.id t in
    match e with
    | K.Store v -> v
    | _ -> assert false

  let find k t = try Some (find_exn k t) with Not_found -> None

  let add_e_pair_ p t =
    let (E_pair ((module K), _)) = p in
    M.add K.id p t

  let add_pair_ p t =
    let (Pair (((module K) as k), v)) = p in
    let p = E_pair (k, K.Store v) in
    M.add K.id p t

  let add (type a) (k : a Key.t) v t =
    let module K = (val k) in
    add_e_pair_ (E_pair (k, K.Store v)) t

  let remove (type a) (k : a Key.t) t =
    let module K = (val k) in
    M.remove K.id t

  let cardinal t = M.cardinal t
  let length = cardinal
  let iter f t = M.iter (fun _ p -> f (pair_of_e_pair p)) t
  let to_iter t yield = iter yield t
  let to_list t = M.fold (fun _ p l -> pair_of_e_pair p :: l) t []
  let add_list t l = List.fold_right add_pair_ l t

  let add_iter t seq =
    let t = ref t in
    seq (fun pair -> t := add_pair_ pair !t);
    !t

  let of_list l = add_list empty l
  let of_iter seq = add_iter empty seq
end