Concestor in OCaml: a purely functional Map backed by a hash table

concestor.ml

module Concestor(Key: Collections.HASH) : sig
  type 'v t

  val empty : unit -> 'v t
  val is_empty : 'v t -> bool
  val add : Key.t -> 'v -> 'v t -> 'v t
  val remove : Key.t -> 'v t -> 'v t
  val mem : Key.t -> 'v t -> bool
  val find : Key.t -> 'v t -> 'v
  val find_opt : Key.t -> 'v t -> 'v option
  val fold_left : ('a -> Key.t -> 'v -> 'a) -> 'a -> 'v t -> 'a
  val count : 'v t -> int
  val to_array : 'v t -> (Key.t * 'v) array
  val of_array : (Key.t * 'v) array -> 'v t
  val of_seq : (Key.t * 'v) Seq.t -> 'v t
  val pp : Format.formatter -> (Format.formatter -> 'v -> unit) -> 'v t -> unit
  val show : ('v -> string) -> 'v t -> string
end = struct
  module HashTable = Collections.HashTable(Key)

  type 'v ancestry =
    | Add of Key.t * 'v * 'v t
    | Remove of Key.t * 'v * 'v t
    | Replace of Key.t * 'v * 'v * 'v t
    | HashTable of 'v HashTable.t
  and 'v t = { mutable ancestry: 'v ancestry }

  let create d = { ancestry = HashTable d }

  let empty() = create(HashTable.make())

  let of_array s = create(HashTable.of_array s)

  let of_seq s = create(HashTable.of_seq s)

  let rec ossifyCPS concestor k =
    let recur relative f =
      ossifyCPS relative (fun d ->
        relative.ancestry <- f d;
        concestor.ancestry <- HashTable d;
        k d) in
    match concestor.ancestry with
    | HashTable s -> k s
    | Add(key, value, relative) ->
        recur relative (fun d ->
          HashTable.add d key value;
          Remove(key, value, relative))
    | Remove(key, value, relative) ->
        recur relative (fun d ->
          HashTable.remove d key;
          Add(key, value, relative))
    | Replace(key, oldValue, newValue, relative) ->
        recur relative (fun d ->
          HashTable.add d key newValue;
          Replace(key, newValue, oldValue, relative))

  let ossify concestor = ossifyCPS concestor (fun x -> x)

  let locked concestor f =
    (* Lock the concestor here for multicore code. *)
    f(ossify concestor)

  let add key value concestor =
    locked concestor (fun d ->
      let curr = { ancestry = HashTable d } in (* {concestor with ...} *)
      let ancestry =
        try
          let oldValue = HashTable.find d key in
          HashTable.add d key value;
          Replace(key, value, oldValue, curr)
        with Not_found ->
          HashTable.add d key value;
          Remove(key, value, curr) in
      concestor.ancestry <- ancestry;
      curr)

  let remove key concestor =
    locked concestor (fun d ->
      try
        let oldValue = HashTable.find d key in
        HashTable.remove d key;
        let curr = { ancestry = HashTable d } in
        concestor.ancestry <- Add(key, oldValue, curr);
        curr
      with Not_found -> concestor)

  let mem key concestor =
    locked concestor (fun d -> HashTable.mem d key)

  let find key concestor =
    locked concestor (fun d -> HashTable.find d key)

  let find_opt key concestor =
    locked concestor (fun d -> HashTable.find_opt d key)

  let count concestor =
    locked concestor (fun d -> HashTable.length d)

  let to_array concestor = locked concestor HashTable.to_array

  let fold_left f a concestor =
    to_array concestor
    |> Array.fold_left (fun a (k, v) -> f a k v) a

  let is_empty concestor =
    count concestor = 0

  let pp fmt pp_v xs = locked xs (HashTable.pp fmt pp_v)

  let show show xs = locked xs (HashTable.show show)
end