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

(** {1 Multimap} *)

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

module type S = sig
  type key
  type value
  type t

  val empty : t
  (** Empty multimap *)

  val is_empty : t -> bool
  (** Empty multimap? *)

  val add : t -> key -> value -> t
  (** Add a key/value binding *)

  val remove : t -> key -> value -> t
  (** Remove the binding *)

  val remove_all : t -> key -> t
  (** Remove the key from the map *)

  val mem : t -> key -> bool
  (** Is there a binding for this key? *)

  val find : t -> key -> value list
  (** List of values for this key *)

  val find_iter : t -> key -> value iter
  (** Iterate on bindings for this key *)

  val count : t -> key -> int
  (** Number of bindings for this key *)

  val iter : t -> (key -> value -> unit) -> unit
  (** Iterate on all key/value *)

  val fold : t -> 'a -> ('a -> key -> value -> 'a) -> 'a
  (** Fold on all key/value *)

  val size : t -> int
  (** Number of keys *)

  val union : t -> t -> t
  (** Union of multimaps *)

  val inter : t -> t -> t
  (** Intersection of multimaps *)

  val diff : t -> t -> t
  (** Difference of maps, ie bindings of the first that are not
      in the second *)

  val equal : t -> t -> bool
  (** Same multimap *)

  val compare : t -> t -> int
  (** Total order on multimaps *)

  val submap : t -> t -> bool
  (** [submap m1 m2] is true iff all bindings of [m1] are also in [m2] *)

  val to_iter : t -> (key * value) iter
  val of_iter : ?init:t -> (key * value) iter -> t
  val keys : t -> key iter

  val values : t -> value iter
  (** Some values may occur several times *)
end

module type OrderedType = sig
  type t

  val compare : t -> t -> int
end

module Make (K : OrderedType) (V : OrderedType) = struct
  type key = K.t
  type value = V.t

  module M = Map.Make (K)
  module S = Set.Make (V)

  type t = S.t M.t
  (** Map of sets *)

  let empty = M.empty
  let is_empty = M.is_empty

  let add m k v =
    let set = try M.find k m with Not_found -> S.empty in
    M.add k (S.add v set) m

  let remove m k v =
    try
      let set = M.find k m in
      let set' = S.remove v set in
      if S.is_empty set' then
        M.remove k m
      else
        M.add k set' m
    with Not_found -> m

  let remove_all m k = M.remove k m
  let mem m k = M.mem k m

  let find m k =
    try
      let set = M.find k m in
      S.elements set
    with Not_found -> []

  let find_iter m k f =
    try
      let set = M.find k m in
      S.iter f set
    with Not_found -> ()

  let count m k =
    try
      let set = M.find k m in
      S.cardinal set
    with Not_found -> 0

  let iter m f = M.iter (fun k set -> S.iter (fun v -> f k v) set) m

  let fold m acc f =
    M.fold (fun k set acc -> S.fold (fun v acc -> f acc k v) set acc) m acc

  let size m = M.cardinal m

  let union m1 m2 =
    M.merge
      (fun _k v1 v2 ->
        match v1, v2 with
        | None, None -> None
        | Some set1, Some set2 -> Some (S.union set1 set2)
        | Some set, None | None, Some set -> Some set)
      m1 m2

  let inter m1 m2 =
    M.merge
      (fun _k v1 v2 ->
        match v1, v2 with
        | None, _ | _, None -> None
        | Some set1, Some set2 ->
          let set = S.inter set1 set2 in
          if S.is_empty set then
            None
          else
            Some set)
      m1 m2

  let diff m1 m2 =
    M.merge
      (fun _k v1 v2 ->
        match v1, v2 with
        | None, _ -> None
        | Some set, None -> Some set
        | Some set1, Some set2 ->
          let set' = S.diff set1 set2 in
          if S.is_empty set' then
            None
          else
            Some set')
      m1 m2

  let equal m1 m2 = M.equal S.equal m1 m2
  let compare m1 m2 = M.compare S.compare m1 m2

  let submap m1 m2 =
    M.for_all
      (fun k set1 ->
        try
          let set2 = M.find k m2 in
          S.subset set1 set2
        with Not_found -> false)
      m1

  let to_iter m k = iter m (fun x y -> k (x, y))

  let of_iter ?(init = empty) seq =
    let m = ref init in
    seq (fun (k, v) -> m := add !m k v);
    !m

  let keys m k = M.iter (fun x _ -> k x) m
  let values m k = iter m (fun _ v -> k v)
end

module type BIDIR = sig
  type t
  type left
  type right

  val empty : t
  val is_empty : t -> bool

  val add : t -> left -> right -> t
  (** Add a binding (left,right) *)

  val remove : t -> left -> right -> t
  (** Remove a specific binding *)

  val cardinal_left : t -> int
  (** number of distinct left keys *)

  val cardinal_right : t -> int
  (** number of distinct right keys *)

  val remove_left : t -> left -> t
  (** Remove all bindings for the left key *)

  val remove_right : t -> right -> t
  (** Remove all bindings for the right key *)

  val mem_left : t -> left -> bool
  (** Is the left key present in at least one pair? *)

  val mem_right : t -> right -> bool
  (** Is the right key present in at least one pair? *)

  val find_left : t -> left -> right list
  (** List of values for this given left-key.
      This used to return an iter, but returns a list since 3.13. *)

  val find_left_iter : t -> left -> right iter
  (** Iterate on bindings for this given left-key
      @since 3.13 *)

  val find_right : t -> right -> left list
  (** List of values for this given right-key.
      This used to return an iter, but returns a list since 3.13. *)

  val find_right_iter : t -> right -> left iter
  (** Iterate on bindings for this given left-key
      @since 3.13 *)

  val find1_left : t -> left -> right option
  (** like {!find_left} but returns at most one value *)

  val find1_right : t -> right -> left option
  (** like {!find_right} but returns at most one value *)

  val fold : ('a -> left -> right -> 'a) -> 'a -> t -> 'a
  (** Fold on pairs *)

  val pairs : t -> (left * right) iter
  (** Iterate on pairs *)

  val add_pairs : t -> (left * right) iter -> t
  (** Add pairs *)

  val iter_left : t -> left iter
  val iter_right : t -> right iter
end

let _fold_iter f acc seq =
  let acc = ref acc in
  seq (fun x -> acc := f !acc x);
  !acc

let _head_iter seq =
  let r = ref None in
  (try
     seq (fun x ->
         r := Some x;
         raise Exit)
   with Exit -> ());
  !r

module MakeBidir (L : OrderedType) (R : OrderedType) = struct
  type left = L.t
  type right = R.t

  module MapL = Make (L) (R)
  module MapR = Make (R) (L)

  type t = {
    left: MapL.t;
    right: MapR.t;
  }

  let empty = { left = MapL.empty; right = MapR.empty }
  let is_empty m = MapL.is_empty m.left
  let add m a b = { left = MapL.add m.left a b; right = MapR.add m.right b a }

  let remove m a b =
    { left = MapL.remove m.left a b; right = MapR.remove m.right b a }

  let cardinal_left m = MapL.size m.left
  let cardinal_right m = MapR.size m.right
  let find_left m a = MapL.find m.left a
  let find_left_iter m a = MapL.find_iter m.left a
  let find_right m b = MapR.find m.right b
  let find_right_iter m b = MapR.find_iter m.right b

  let remove_left m a =
    _fold_iter (fun m b -> remove m a b) m (find_left_iter m a)

  let remove_right m b =
    _fold_iter (fun m a -> remove m a b) m (find_right_iter m b)

  let mem_left m a = MapL.mem m.left a
  let mem_right m b = MapR.mem m.right b
  let find1_left m a = _head_iter (find_left_iter m a)
  let find1_right m b = _head_iter (find_right_iter m b)
  let fold f acc m = MapL.fold m.left acc f
  let pairs m = MapL.to_iter m.left
  let add_pairs m seq = _fold_iter (fun m (a, b) -> add m a b) m seq
  let iter_left m = MapL.keys m.left
  let iter_right m = MapR.keys m.right
end
