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

(** {1 Mutable Set} *)

type 'a iter = ('a -> unit) -> unit
type 'a printer = Format.formatter -> 'a -> unit

module type S = sig
  type t
  type elt

  val create : int -> t
  (** [create n] makes a new set with the given capacity [n] *)

  val singleton : elt -> t
  (** [singleton x] is the singleton [{x}] *)

  val clear : t -> unit
  (** [clear s] removes all elements from [s] *)

  val copy : t -> t
  (** Fresh copy *)

  val copy_into : into:t -> t -> unit
  (** [copy_into ~into s] copies all elements of [s] into [into] *)

  val insert : t -> elt -> unit
  (** [insert s x] adds [x] into [s] *)

  val remove : t -> elt -> unit
  (** Remove the element, if it were in there *)

  val cardinal : t -> int
  (** [cardinal s] returns the number of elements in [s] *)

  val mem : t -> elt -> bool
  (** [mem s x] returns [true] iff [x] is in [s] *)

  val find_exn : t -> elt -> elt
  (** [find_exn s x] returns [y] if [x] and [y] are equal, and [mem s y].
      @raise Not_found if [x] not in [s] *)

  val find : t -> elt -> elt option
  (** Safe version of {!find_exn} *)

  val inter : t -> t -> t
  (** [inter a b] returns [a ∩ b] *)

  val inter_mut : into:t -> t -> unit
  (** [inter_mut ~into a] changes [into] into [a ∩ into] *)

  val union : t -> t -> t
  (** [union a b] returns [a ∪ b] *)

  val union_mut : into:t -> t -> unit
  (** [union_mut ~into a] changes [into] into [a ∪ into] *)

  val diff : t -> t -> t
  (** [diff a b] returns [a - b] *)

  val subset : t -> t -> bool
  (** [subset a b] returns [true] if all elements of [a] are in [b] *)

  val equal : t -> t -> bool
  (** [equal a b] is extensional equality ([a] and [b] have the same elements) *)

  val for_all : (elt -> bool) -> t -> bool
  val exists : (elt -> bool) -> t -> bool

  val iter : (elt -> unit) -> t -> unit
  (** Iterate on values *)

  val fold : ('a -> elt -> 'a) -> 'a -> t -> 'a
  (** Fold on values *)

  val elements : t -> elt list
  (** List of elements *)

  val of_list : elt list -> t
  val to_iter : t -> elt iter
  val of_iter : elt iter -> t
  val add_iter : t -> elt iter -> unit

  val pp : ?pp_sep:unit printer -> elt printer -> t printer
  (** [pp pp_elt] returns a set printer, given a printer for
      individual elements *)
end

module type ELEMENT = sig
  type t

  val equal : t -> t -> bool

  val hash : t -> int
  (** Positive value *)
end

module Make (E : ELEMENT) : S with type elt = E.t = struct
  module Tbl = Hashtbl.Make (E)

  type elt = E.t
  type t = elt Tbl.t (* map [x -> x], for find *)

  let create = Tbl.create

  let singleton x =
    let s = create 8 in
    Tbl.replace s x x;
    s

  let clear = Tbl.clear
  let copy = Tbl.copy
  let copy_into ~into s = Tbl.iter (fun x _ -> Tbl.replace into x x) s
  let insert s x = Tbl.replace s x x
  let remove = Tbl.remove
  let cardinal = Tbl.length
  let mem = Tbl.mem
  let find_exn = Tbl.find
  let find s x = try Some (Tbl.find s x) with Not_found -> None
  let iter f s = Tbl.iter (fun x _ -> f x) s
  let fold f acc s = Tbl.fold (fun x _ acc -> f acc x) s acc

  let inter a b =
    let res = create (min (cardinal a) (cardinal b)) in
    iter (fun x -> if mem a x then insert res x) b;
    res

  let inter_mut ~into a =
    iter (fun x -> if not (mem a x) then remove into x) into

  let union a b =
    let res = copy a in
    copy_into ~into:res b;
    res

  let union_mut ~into a = copy_into ~into a

  let diff a b =
    let res = copy a in
    iter (fun x -> remove res x) b;
    res

  exception FastExit

  let for_all p s =
    try
      Tbl.iter (fun x _ -> if not (p x) then raise FastExit) s;
      true
    with FastExit -> false

  let exists p s =
    try
      Tbl.iter (fun x _ -> if p x then raise FastExit) s;
      false
    with FastExit -> true

  let subset a b = for_all (fun x -> mem b x) a
  let equal a b = subset a b && subset b a
  let elements s = Tbl.fold (fun x _ acc -> x :: acc) s []

  let of_list l =
    let res = create (List.length l) in
    List.iter (insert res) l;
    res

  let to_iter s yield = iter yield s
  let add_iter s seq = seq (insert s)

  let of_iter seq =
    let s = create 32 in
    seq (insert s);
    s

  let pp ?(pp_sep = fun out () -> Format.fprintf out ",@ ") pp_x out s =
    Format.pp_print_string out "{";
    let first = ref true in
    Tbl.iter
      (fun x _ ->
        if !first then
          first := false
        else
          pp_sep out ();
        pp_x out x)
      s;
    Format.pp_print_string out "}"
end
