(** This modules describes complete (and potentially finite) mappings. *)


(** The signature of the functor's parameter module. *)

module type T = sig
  type key
  type 'a t
  (** [keys] needs to be given for finite mappings, and needs to be [None] for
      non finite mappings. *)
  val keys  : key list option
  val mem   : key -> 'a t -> bool
  val find  : key -> 'a t -> 'a
  val add   : key -> 'a -> 'a t -> 'a t
  val empty : 'a t
  val map   : ('a -> 'b) -> 'a t -> 'b t
  val fold  : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
end

(** The result signature of the functor. *)

module type S = sig
  type key
  type 'a t
  val empty     : 'a -> 'a t
  val find      : key -> 'a t -> 'a
  val upd       : key -> 'a -> 'a t -> 'a t
  val map       : ('a -> 'b) -> 'a t -> 'b t
  (** Keys in case of finite mappings. *)
  val keys      : key list option
  (** Only works on [keys]. Raises [Failure "CompleteMap.fold: not finite
      domain"] is no keys. *)
  val fold      : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
  val merge     : ('a -> 'b -> 'c) -> 'a t -> 'b t -> 'c t
  val cmp       : ('a -> 'b -> bool -> bool) -> 'a t -> 'b t -> bool -> bool
  val to_string : (key -> string) -> ('a -> string) -> 'a t -> string
end

module Make (T : T) : S with type key = T.key = struct

  type key = T.key

  type 'a t = ('a * 'a T.t)

  let empty a = (a, T.empty)

  let find x (a, map) = if T.mem x map then T.find x map else a

  let (!!) = find

  let upd x b (a, map) = (a, T.add x b map)

  let map f (a, map) = (f a, T.map f map)

  let keys = T.keys

  let fold f cmap a = match T.keys with
    | None -> raise (Failure "CompleteMap.fold: not finite domain")
    | Some keys ->
      List.fold_left (fun res key -> f key (find key cmap) res) a keys

  let f_merge f_default f_map ((a1, map1) as cmap1) ((a2, map2) as cmap2) a =
    let f1 x v1 a = f_map x v1 (find x cmap2) a in
    let f2 x v2 a = f_map x (find x cmap1) v2 a in
    let a = T.fold f1 map1 a in
    f_default a1 a2 (T.fold f2 map2 a)

  let merge f cmap1 cmap2 =
    let f_default a1 a2 map = (f a1 a2, map) in
    let f_map x v1 v2 map = T.add x (f v1 v2) map in
    f_merge f_default f_map cmap1 cmap2 T.empty

  let cmp f cmap1 cmap2 a =
    let f_default a1 a2 a = f a1 a2 a in
    let f_map _ v1 v2 a = f v1 v2 a in
    f_merge f_default f_map cmap1 cmap2 a

  let to_string string_of_key string_of_a (_, map) =
    let f key a res =
      res ^ "--KEY\n" ^ (string_of_key key) ^ "\n--VAL\n" ^ (string_of_a a) in
    T.fold f map ""

end