type ('k, 'v) t = ('k, 'v list ref) PMap.t ref

let create () : ('k, 'v) t = ref PMap.empty
let clear (m_ref : ('k, 'v) t) : unit = m_ref := PMap.empty

let add (m_ref : ('k, 'v) t) (x : 'k) (d : 'v) : unit =
  try
    let vlst_ref = PMap.find x !m_ref in
    vlst_ref := d :: !vlst_ref
  with Not_found -> m_ref := PMap.add x (ref [d]) !m_ref

let copy_ref r = ref !r
let copy (m_ref : ('k, 'v) t) : ('k, 'v) t = ref (PMap.map copy_ref !m_ref)
let find (m_ref : ('k, 'v) t) (x : 'k) : 'v = List.hd !(PMap.find x !m_ref)
let find_all m_ref x : 'v list = !(PMap.find x !m_ref)
let mem (m_ref : ('k, 'v) t) (x : 'k) : bool = PMap.mem x !m_ref

let remove (m_ref : ('k, 'v) t) (x : 'k) : unit =
  try
    let vlst_ref = PMap.find x !m_ref in
    match !vlst_ref with
    | [_] -> m_ref := PMap.remove x !m_ref
    | _ :: rest -> vlst_ref := rest
    | [] -> assert false
  with Not_found -> ()

let replace (m_ref : ('a, 'b) t) x d =
  try
    let vlst_ref = PMap.find x !m_ref in
    vlst_ref := d :: List.tl !vlst_ref
  with Not_found -> m_ref := PMap.add x (ref [d]) !m_ref

let iter (f : 'k -> 'v -> unit) (m_ref : ('k, 'v) t) : unit =
  let f' k vlst_ref = List.iter (f k) !vlst_ref in
  PMap.iter f' !m_ref

let fold (f : 'k -> 'v -> 'a -> 'a) (m_ref : ('k, 'v) t) (acc : 'a) : 'a =
  let f' k vlst_ref acc =
    List.fold_left (fun acc v -> f k v acc) acc !vlst_ref in
  PMap.fold f' !m_ref acc