Multimap

multimap.ml

module MultiMap (T : Map.OrderedType) : sig
  type 'a t
  val of_list     : (T.t list * 'a) list -> 'a t
  val find        : T.t -> 'a t -> 'a option
  val find_update : T.t -> ('a -> 'a) -> 'a t -> ('a * 'a t) option
end =
struct

  module TMap = Map.Make (T)
  module IMap = Map.Make (struct type t = int let compare a b = compare a b end)

  type 'a t = int TMap.t * 'a IMap.t

  let of_list xxs =
    let rec go i m1 m2 = function
      | []           -> (m1, m2)
      | (ks, v)::xxs ->
          let m1 = List.fold_left (fun m x -> TMap.add x i m) m1 ks in
          let m2 = IMap.add i v m2 in
          go (succ i) m1 m2 xxs in
    go 0 TMap.empty IMap.empty xxs

  let find k (m1, m2) =
    try Some (IMap.find (TMap.find k m1) m2) with Not_found -> None

  let find_update k f (m1, m2) =
    try
      let i = TMap.find k m1 in
      let v = IMap.find i m2 in
      let m2 = IMap.add i (f v) m2 in
      Some (v, (m1, m2))
    with Not_found -> None
end