sidekick/src/base/Het.ml
2022-08-05 21:56:17 -04:00

129 lines
2.9 KiB
OCaml

(* This file is free software, part of containers. See file "license" for more details. *)
(** {1 Associative containers with Heterogenerous Values} *)
(*$R
let k1 : int Key.t = Key.create() in
let k2 : int Key.t = Key.create() in
let k3 : string Key.t = Key.create() in
let k4 : float Key.t = Key.create() in
let tbl = Tbl.create () in
Tbl.add tbl k1 1;
Tbl.add tbl k2 2;
Tbl.add tbl k3 "k3";
assert_equal (Some 1) (Tbl.find tbl k1);
assert_equal (Some 2) (Tbl.find tbl k2);
assert_equal (Some "k3") (Tbl.find tbl k3);
assert_equal None (Tbl.find tbl k4);
assert_equal 3 (Tbl.length tbl);
Tbl.add tbl k1 10;
assert_equal (Some 10) (Tbl.find tbl k1);
assert_equal 3 (Tbl.length tbl);
assert_equal None (Tbl.find tbl k4);
Tbl.add tbl k4 0.0;
assert_equal (Some 0.0) (Tbl.find tbl k4);
()
*)
type 'a iter = ('a -> unit) -> unit
type 'a gen = unit -> 'a option
module type KEY_IMPL = sig
type t
exception Store of t
val id : int
end
module Key = struct
type 'a t = (module KEY_IMPL with type t = 'a)
let _n = ref 0
let create (type k) () =
incr _n;
let id = !_n in
let module K = struct
type t = k
let id = id
exception Store of k
end in
(module K : KEY_IMPL with type t = k)
let id (type k) (module K : KEY_IMPL with type t = k) = K.id
let equal : type a b. a t -> b t -> bool =
fun (module K1) (module K2) -> K1.id = K2.id
end
type pair = Pair : 'a Key.t * 'a -> pair
type exn_pair = E_pair : 'a Key.t * exn -> exn_pair
let pair_of_e_pair (E_pair (k, e)) =
let module K = (val k) in
match e with
| K.Store v -> Pair (k, v)
| _ -> assert false
module Map = struct
module M = Map.Make (struct
type t = int
let compare (i : int) j = CCInt.compare i j
end)
type t = exn_pair M.t
let empty = M.empty
let mem k t = M.mem (Key.id k) t
let find_exn (type a) (k : a Key.t) t : a =
let module K = (val k) in
let (E_pair (_, e)) = M.find K.id t in
match e with
| K.Store v -> v
| _ -> assert false
let find k t = try Some (find_exn k t) with Not_found -> None
let add_e_pair_ p t =
let (E_pair ((module K), _)) = p in
M.add K.id p t
let add_pair_ p t =
let (Pair (((module K) as k), v)) = p in
let p = E_pair (k, K.Store v) in
M.add K.id p t
let add (type a) (k : a Key.t) v t =
let module K = (val k) in
add_e_pair_ (E_pair (k, K.Store v)) t
let remove (type a) (k : a Key.t) t =
let module K = (val k) in
M.remove K.id t
let cardinal t = M.cardinal t
let length = cardinal
let iter f t = M.iter (fun _ p -> f (pair_of_e_pair p)) t
let to_iter t yield = iter yield t
let to_list t = M.fold (fun _ p l -> pair_of_e_pair p :: l) t []
let add_list t l = List.fold_right add_pair_ l t
let add_iter t seq =
let t = ref t in
seq (fun pair -> t := add_pair_ pair !t);
!t
let of_list l = add_list empty l
let of_iter seq = add_iter empty seq
end