module type S = Hashtbl_intf.S

module Make (H : sig
    include Hashable.S

    val to_dyn : t -> Dyn.t
  end) =
struct
  include MoreLabels.Hashtbl.Make (H)

  let[@ocaml.warning "-32"] add = `Use_set
  let find = find_opt

  let find_exn t key =
    match find_opt t key with
    | Some v -> v
    | None -> Code_error.raise "Hashtbl.find_exn" [ "key", H.to_dyn key ]
  ;;

  let set t key data = replace t ~key ~data

  let find_or_add t key ~f =
    match find t key with
    | Some x -> x
    | None ->
      let x = f key in
      set t key x;
      x
  ;;

  let foldi t ~init ~f = fold t ~init ~f:(fun ~key ~data acc -> f key data acc)
  let fold t ~init ~f = foldi t ~init ~f:(fun _ x -> f x)

  let of_list l =
    let h = create (List.length l) in
    let rec loop = function
      | [] -> Result.Ok h
      | (k, v) :: xs ->
        (match find h k with
         | None ->
           set h k v;
           loop xs
         | Some v' -> Error (k, v', v))
    in
    loop l
  ;;

  let of_list_exn l =
    match of_list l with
    | Result.Ok h -> h
    | Error (key, _, _) ->
      Code_error.raise "Hashtbl.of_list_exn duplicate keys" [ "key", H.to_dyn key ]
  ;;

  let add_exn t key data =
    match find t key with
    | None -> set t key data
    | Some _ ->
      Code_error.raise "Hashtbl.add_exn: key already exists" [ "key", H.to_dyn key ]
  ;;

  let add t key data =
    match find t key with
    | None ->
      set t key data;
      Result.Ok ()
    | Some p -> Result.Error p
  ;;

  let keys t = foldi t ~init:[] ~f:(fun key _ acc -> key :: acc)

  let to_dyn f t =
    Dyn.Map
      (foldi t ~init:[] ~f:(fun key data acc -> (H.to_dyn key, f data) :: acc)
       |> List.sort ~compare:(fun (k, _) (k', _) -> Dyn.compare k k'))
  ;;

  let to_list t = foldi t ~init:[] ~f:(fun key v acc -> (key, v) :: acc)

  let filteri_inplace t ~f =
    filter_map_inplace t ~f:(fun ~key ~data ->
      match f ~key ~data with
      | true -> Some data
      | false -> None)
  ;;

  let iter t ~f = iter t ~f:(fun ~key:_ ~data -> f data)
end