Source file sMap.ml

(*****************************************************************************)
(*                                                                           *)
(* MIT License                                                               *)
(* Copyright (c) 2022 Nomadic Labs <contact@nomadic-labs.com>                *)
(*                                                                           *)
(* Permission is hereby granted, free of charge, to any person obtaining a   *)
(* copy of this software and associated documentation files (the "Software"),*)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense,  *)
(* and/or sell copies of the Software, and to permit persons to whom the     *)
(* Software is furnished to do so, subject to the following conditions:      *)
(*                                                                           *)
(* The above copyright notice and this permission notice shall be included   *)
(* in all copies or substantial portions of the Software.                    *)
(*                                                                           *)
(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL   *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING   *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER       *)
(* DEALINGS IN THE SOFTWARE.                                                 *)
(*                                                                           *)
(*****************************************************************************)

module StringMap = struct
  module M = Map.Make (String)
  include M

  let of_list l = of_seq (List.to_seq l)

  let t (inner : 'a Repr.t) : 'a t Repr.t =
    let module M = Repr.Of_map (struct
      include M

      let key_t = Repr.string
    end) in
    M.t inner

  let to_bytes printer map =
    fold
      (fun key elt state ->
        Bytes.cat (Bytes.of_string key) (Bytes.cat (printer elt) state))
      map Bytes.empty

  let show (show_inner : 'a -> string) : 'a t -> string =
   fun m ->
    "{\n"
    ^ String.concat "\n"
        (List.map (fun (k, v) -> k ^ ": " ^ show_inner v) (bindings m))
    ^ "\n}"

  (* Return the union of two maps. The keys of the maps have to be disjoint unless
     specifically stated in common_keys. In this case both key's values
     are asserted to be equal, with a given equality function.
     If no equal function is given the polymorphic euqality is used.*)
  let union_disjoint ?(common_keys_equal_elt = ([], ( = ))) x y =
    let common_keys, equal_elt = common_keys_equal_elt in
    union
      (fun key elt_1 elt_2 ->
        if not (List.mem key common_keys) then
          raise
            (Invalid_argument
               (Printf.sprintf
                  "the key %s appears in both union arguments and does not \
                   belong\n\
                  \                                 to common_keys." key))
        else if not (equal_elt elt_1 elt_2) then
          raise
            (Invalid_argument
               (Printf.sprintf
                  "the key %s appears in both union argument with different \
                   values"
                  key))
        else Some elt_1)
      x y

  (* applies union_disjoint on a list of map*)
  let union_disjoint_list ?(common_keys_equal_elt = ([], ( = ))) map_list =
    List.fold_left (union_disjoint ~common_keys_equal_elt) empty map_list

  (* given a list of maps outputs a single map with the union of all keys and
     containing lists which consist of the concatenation of the data elements
     under the same key (order is preserved) *)
  let map_list_to_list_map map_list =
    let join _key x y = Some (x @ y) in
    List.fold_left
      (fun list_map m -> union join list_map (map (fun x -> [ x ]) m))
      empty map_list

  let sub_map sub_map map =
    let res = filter (fun name _ -> mem name sub_map) map in
    if cardinal res <> cardinal sub_map then
      failwith "sub_map : first argument is not contained in the second.";
    res

  module Aggregation = struct
    (* separator between prefixes & name ; must be only one character *)
    let sep = "~"

    let padded ~n i =
      let str = string_of_int i in
      let len = String.length (string_of_int (n - 1)) in
      String.(make (len - length str) '0') ^ str

    let add_prefix ?(n = 1) ?(i = 0) ?(shift = 0) prefix str =
      let out = if prefix = "" then str else prefix ^ sep ^ str in
      if n = 1 then out else padded ~n (i + shift) ^ sep ^ out

    (* adds prefix to each key of str_map *)
    let prefix_map ?n ?i ?shift prefix str_map =
      fold
        (fun k v acc -> add (add_prefix ?n ?i ?shift prefix k) v acc)
        str_map empty
  end
end

module type S = sig
  include Map.S with type key = string and type 'a t = 'a StringMap.t

  val t : 'a Repr.ty -> 'a t Repr.ty
  val of_list : (string * 'a) list -> 'a t

  val union_disjoint :
    ?common_keys_equal_elt:string list * ('a -> 'a -> bool) ->
    'a t ->
    'a t ->
    'a t

  val union_disjoint_list :
    ?common_keys_equal_elt:string list * ('a -> 'a -> bool) -> 'a t list -> 'a t

  val map_list_to_list_map : 'a t list -> 'a list t

  (* [sub_map m1 m2] returns m2 without the keys that do not appear in m1.
     Raises failure if some key of m1 is not in m2
  *)
  val sub_map : 'a t -> 'b t -> 'b t

  module Aggregation : sig
    val sep : string

    val add_prefix :
      ?n:int -> ?i:int -> ?shift:int -> string -> string -> string

    val prefix_map : ?n:int -> ?i:int -> ?shift:int -> string -> 'a t -> 'a t
  end
end

include (StringMap : S)