ohmap.ml

module T0 = struct
  module type S = sig type t end
end

(** Finite-map; open-to-extension; heterogenous values (type determined by the key).
 *
 *  See the bottom of this file for an example use.
 *
 *  New key types can be generated on the fly and used to modify or query the
 *  contents of the map.  *)
module OpenHetMap = struct

  (** Basic signature (more helper methods added in signature 'S').  *)
  module type S0 = sig
    type t
    type 'a key
    val empty      : t
    val create_key : (module T0.S with type t = 'a) -> 'a key
    val find_opt   : 'a key -> t -> 'a option
    val update     : 'a key -> ('a option -> 'a option) -> t -> t
  end

  (** Basic generator (use Make() instead...).  *)
  module Make0() : S0 = struct

    (** Underlying key representation, called 'tags'.  *)
    module Tag = struct
      include Int
      let next = ref 0
      let gen () = next := !next + 1; !next
    end

    (** Underlying map representation.  *)
    module Map = Map.Make(Tag)
    type entry = ..
    type t = entry Map.t
    let empty = Map.empty

    (** "Keys" tell us how to wrap/unwrap entries and which tag to use.  *)
    module type Key_S = sig
      type a
      val tag    : Tag.t
      val wrap   : a -> entry
      val unwrap : entry option -> a option
    end
    type 'a key = (module Key_S with type a = 'a)

    let create_key (type a) (module A : T0.S with type t = a) =
      (module struct
         type nonrec a = a
         type entry += Inj of a
         let tag = Tag.gen ()
         let wrap x = Inj x
         let unwrap = function
           | Some (Inj x) -> Some x
           | _ -> None
       end : Key_S with type a = a)

    (* Call into underlying map operations but use tag/wrap/unwrap appropriately *)
    let find_opt (type a) (module K : Key_S with type a = a) m =
      K.unwrap (Map.find_opt K.tag m)
    let update (type a) (module K : Key_S with type a = a) f m =
      Map.update K.tag (fun x -> Option.map K.wrap (f (K.unwrap x))) m
  end

  (** Signature for OHM's with some additional operations. Note that a lot of the
   *  operations usually present on map's require higher polymorphic signatures now,
   *  so they aren't present...  *)
  module type S = sig
    include S0
    val find   : 'a key -> t -> 'a
    val add    : 'a key -> 'a -> t -> t
    val remove : 'a key -> t -> t
  end

  (** Generate a new map type, initially without any keys.  *)
  module Make() : S = struct
    include Make0()
    let find k m = match find_opt k m with
      | Some x -> x
      | None -> raise Not_found
    let add k v m = update k (fun _ -> Some v) m
    let remove k m = update k (fun _ -> None) m
  end
end

include OpenHetMap

(****************************************************************************************
 * Example usage: key/values for configuration data
 ****************************************************************************************)

module Config = OpenHetMap.Make()

let name    : string Config.key = Config.create_key (module String)
let fedi    : string Config.key = Config.create_key (module String)
let hidden  : bool Config.key   = Config.create_key (module Bool)
let timeout : int Config.key    = Config.create_key (module Int)

let example_map =
  Config.empty
  |> Config.add name "milo"
  |> Config.add fedi "iitalics@cybre.space"
  |> Config.add timeout 20

let () =
  ( Printf.printf "%S" (Config.find name example_map)
  ; Printf.printf "%d" (Config.find timeout example_map)
  )