Source file set.ml

module type S = Set_intf.S

module Make (Key : Map_intf.Key) (M : Map_intf.S with type key = Key.t) = struct
  include Stdlib.MoreLabels.Set.Make (struct
    type t = Key.t

    let compare x y = Ordering.to_int (Key.compare x y)
  end)

  type 'a map = 'a M.t

  let to_list = elements

  let to_list_map t ~f =
    (* optimize if we get a right fold *)
    fold t ~init:[] ~f:(fun a acc -> f a :: acc) |> List.rev

  let mem t x = mem x t

  let add t x = add x t

  let remove t x = remove x t

  let compare a b = Ordering.of_int (compare a b)

  let is_subset t ~of_ = subset t of_

  let iter t ~f = iter t ~f

  let map t ~f = map t ~f

  let fold t ~init ~f = fold t ~init ~f

  let for_all t ~f = for_all t ~f

  let exists t ~f = exists t ~f

  let filter t ~f = filter t ~f

  let partition t ~f = partition t ~f

  let min_elt = min_elt_opt

  let max_elt = max_elt_opt

  let choose = choose_opt

  let split x t = split t x

  let union_map l ~f =
    List.fold_left ~init:empty l ~f:(fun acc x ->
        let s = f x in
        union acc s)

  let union_all l = union_map l ~f:(fun x -> x)

  exception Found of elt

  let find t ~f =
    match iter t ~f:(fun e -> if f e then raise_notrace (Found e) else ()) with
    | () -> None
    | exception Found e -> Some e

  let to_dyn t = Dyn.Set (to_list t |> List.map ~f:Key.to_dyn)

  let choose_exn t =
    match choose t with
    | Some e -> e
    | None -> Code_error.raise "Set.choose_exn" [ ("t", to_dyn t) ]

  let of_keys = M.foldi ~init:empty ~f:(fun k _ acc -> add acc k)

  let to_map t ~f = fold t ~init:M.empty ~f:(fun k acc -> M.set acc k (f k))

  let of_list_map xs ~f =
    (* We don't [fold_left] & [add] over [xs] because [of_list] has a
       specialized implementation *)
    List.map xs ~f |> of_list
end

module Of_map (Key : Map_intf.Key) (Map : Map_intf.S with type key = Key.t) =
struct
  type elt = Key.t

  type 'a map = 'a Map.t

  type t = unit Map.t

  let empty = Map.empty

  let is_empty = Map.is_empty

  let singleton x = Map.singleton x ()

  let cardinal = Map.cardinal

  let equal a b = Map.equal a b ~equal:(fun () () -> true)

  let compare a b = Map.compare a b ~compare:(fun () () -> Eq)

  let to_list = Map.keys

  let to_list_map t ~f =
    (* optimize if we get a right fold *)
    Map.to_list_map t ~f:(fun a () -> f a)

  let mem t x = Map.mem t x

  let add t x = Map.set t x ()

  let remove t x = Map.remove t x

  exception Not_a_subset

  let is_subset t ~of_ =
    match
      Map.merge t of_ ~f:(fun _key a b ->
          match (a, b) with
          | Some (), None -> raise_notrace Not_a_subset
          | _ -> None)
    with
    | (_ : t) -> true
    | exception Not_a_subset -> false

  let iter t ~f = Map.iteri t ~f:(fun k () -> f k)

  let fold t ~init ~f = Map.foldi t ~init ~f:(fun k () acc -> f k acc)

  let map t ~f = fold t ~init:empty ~f:(fun x acc -> add acc (f x))

  let for_all t ~f = Map.for_alli t ~f:(fun k () -> f k)

  let exists t ~f = Map.existsi t ~f:(fun k () -> f k)

  let filter t ~f = Map.filteri t ~f:(fun k () -> f k)

  let partition t ~f = Map.partitioni t ~f:(fun k () -> f k)

  let min_elt t = Map.min_binding t |> Option.map ~f:fst

  let max_elt t = Map.max_binding t |> Option.map ~f:fst

  let choose t = Map.choose t |> Option.map ~f:fst

  let split t x =
    let a, x, b = Map.split t x in
    (a, Option.is_some x, b)

  let union a b = Map.union a b ~f:(fun _k () () -> Some ())

  let diff a b =
    Map.merge a b ~f:(fun _k a b ->
        match b with
        | Some () -> None
        | None -> a)

  let inter a b =
    Map.merge a b ~f:(fun _k a b ->
        match a with
        | Some () -> b
        | None -> None)

  let union_map l ~f =
    List.fold_left ~init:empty l ~f:(fun acc x ->
        let s = f x in
        union acc s)

  let union_all l = union_map l ~f:(fun x -> x)

  exception Found of elt

  let find t ~f =
    match iter t ~f:(fun e -> if f e then raise_notrace (Found e) else ()) with
    | () -> None
    | exception Found e -> Some e

  let to_dyn t = Dyn.Set (to_list t |> List.map ~f:Key.to_dyn)

  let choose_exn t =
    match choose t with
    | Some e -> e
    | None -> Code_error.raise "Set.choose_exn" [ ("t", to_dyn t) ]

  let of_keys t = Map.map t ~f:ignore

  let to_map t ~f = Map.mapi t ~f:(fun k () -> f k)

  let of_list l = List.fold_left l ~init:empty ~f:add

  let of_list_map l ~f =
    List.fold_left l ~init:empty ~f:(fun acc x -> add acc (f x))

  let to_seq t = Map.to_seq t |> Seq.map ~f:fst
end