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 keys t = bindings t |> List.map fst

  let values t = bindings t |> List.map snd

  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}"

  let to_pair m = (map fst m, map snd m)

  let add_unique k v m =
    if mem k m then
      raise
        (Invalid_argument (Printf.sprintf "key %s already present in map." k))
    else add k v m

  (* 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

  let update_keys f map = fold (fun k v acc -> add (f k) v acc) map empty

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

    let update_key_name f str =
      match String.rindex_from_opt str (String.length str - 1) sep.[0] with
      | None -> f str
      | Some i ->
          String.sub str 0 (i + 1)
          ^ f (String.sub str (i + 1) (String.length str - i - 1))

    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 ?(no_sep = false) ?(n = 0) ?(i = 0) ?(shift = 0) prefix str =
      let prefix = if prefix = "" || no_sep then prefix else prefix ^ sep in
      if n = 0 then prefix ^ str else prefix ^ padded ~n (i + shift) ^ sep ^ str

    let build_all_names prefix n name =
      List.init n (fun i -> add_prefix ~n ~i prefix name)

    let prefix_map ?n ?i ?shift prefix str_map =
      fold (fun k -> add (add_prefix ?n ?i ?shift prefix k)) str_map empty

    let of_list ?n ?shift prefix name l =
      of_list
      @@ List.mapi (fun i x -> (add_prefix ?n ~i ?shift prefix name, x)) l

    (* This function will merge the maps of the list, by prefixing each key with it’s index in the list, optionnally with a shift, with the index prefix prefixed with zero to we able to handle n elements with the same prefix size (with n either given by [shift] or by the length of [list_map]) ; if a [prefix] is given, it will be put before the index.
       *)
    let map_of_list_map ?(prefix = "") ?shift list_map =
      let shift, n = Option.value ~default:(0, List.length list_map) shift in
      List.mapi (fun i m -> prefix_map ~n ~i ~shift prefix m) list_map
      |> union_disjoint_list

    let smap_of_smap_smap mapmap =
      fold (fun prefix map res -> prefix_map prefix map :: res) mapmap []
      |> union_disjoint_list

    let gather_maps ?(shifts_map = empty) map_list_map =
      mapi
        (fun name list_map ->
          map_of_list_map ?shift:(find_opt name shifts_map) list_map)
        map_list_map
      |> smap_of_smap_smap

    let filter_by_circuit_name circuit_name =
      let sep_char =
        assert (String.length sep = 1) ;
        String.get sep 0
      in
      filter (fun key _ ->
          let name_parts = String.split_on_char sep_char key in
          circuit_name = ""
          || List.exists
               (String.equal circuit_name)
               (* we exclude the last element in [name_parts] *)
               (List.rev name_parts |> List.tl))

    let select_answers_by_circuit circuit_name =
      map (filter_by_circuit_name circuit_name)

    let add_map_list_map m1 m2 =
      mapi
        (fun k l1 ->
          match find_opt k m2 with
          | Some l2 -> List.map2 union_disjoint l1 l2
          | None -> l1)
        m1
  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 keys : 'a t -> string list

  val values : 'a t -> 'a list

  (* Splits a map of pairs into a pair of maps *)
  val to_pair : ('a * 'b) t -> 'a t * 'b t

  (* [add_unique k v map] adds [k -> v] to [map] & throw an error if [k] is
     already in [map]
  *)
  val add_unique : key -> 'a -> 'a t -> '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

  (* USE WITH CAUTION : be sure your update function won’t create duplications *)
  val update_keys : (key -> key) -> 'a t -> 'a t

  module Aggregation : sig
    (* Separator for prefixing *)
    val sep : string

    (* applies the input function on the last part of the input string (parts
       are delimited by sep).
       [update_key_name f ("hello" ^ sep ^ "world" ^ sep ^ "!!")]
       returns ["hello" ^ sep ^ "world" ^ sep ^ (f "!!")]
    *)
    val update_key_name : (key -> key) -> key -> key

    (* [add_prefix ~n ~i ~shift prefix str] return idx^prefix^sep^str
       idx = [i] + [shift] as a string, eventually padded with '0' before to
       allow a numbering until [n] with the same number of caracters
       for instance, [prefix ~n:11 ~i:5 ~shift:1 "hello" "world"] will return
       "06~hello~world"
       [n] is zero by default, this means if no n is specified, no idx will be
        added
       [no_sep] is false by default ; if set to true, the separator before the
        string to prefix will be ommitted :
       [prefix ~no_sep:true ~n:11 ~i:5 ~shift:1 "hello" "world"] will return
       "06~helloworld"
    *)
    val add_prefix :
      ?no_sep:bool ->
      ?n:int ->
      ?i:int ->
      ?shift:int ->
      string ->
      string ->
      string

    (* [build_all_names prefix n k] build the list of all prefixed [k] with
       n proofs : [build_all_names "hello" 11 "world"] will return
       ["hello~00~world" ; "hello~01~world" ; … ; "hello~10~world"] *)
    val build_all_names : key -> int -> key -> key list

    (* adds prefix to each key of str_map ; [i] will be added as a string
       before the prefix
       For instance [prefix_map ~n:3000 ~i:5 ~shift:1 "hello" map] will prefix
       all the keys of [map] with "0006~hello~"
    *)
    val prefix_map : ?n:int -> ?i:int -> ?shift:int -> string -> 'a t -> 'a t

    (* [Aggregation.of_list ~n ~shift s name l] is the same as
       [of_list @@ List.mapi (fun i x -> Aggregation.add_prefix ~n ~i ~shift s name, x) l] *)
    val of_list : ?n:int -> ?shift:int -> string -> string -> 'a list -> 'a t

    (* "c1" -> {"a" ; "b"} ; "c2" -> {"a" ; "c"} becomes
       {"c1~a" ; "c1~b" ; "c2~a" ; "c2~c"} with the same values *)
    val smap_of_smap_smap : 'a t t -> 'a t

    (* Converts a map of list of map in a map, by merging each list of map in
       one map, prefixing all keys with their proof index, and then merging all
       the new maps into one prefixing the keys with the outside map’s keys.
       shifts_maps map outside key to pairs of integers.
       'key1' -> (7, 20) means that 20 proofs will be produced for key1 in
       total and we should start from the 8th one, assuming 7 of them were done
       independently. (Note that we may not even finish the whole 20, this
       depends on the map_list length).
       For example, on input:
       { 'circuit_foo' -> [ {'a' -> fa0; 'b' -> fb0; 'c' -> fc0};
                            {'a' -> fa1; 'b' -> fb1; 'c' -> fc1} ];
         'circuit_bar' -> [ {'a' -> ga0; 'b' -> gb0; 'c' -> gc0} ]; }
       outputs
       { 'circuit_foo~0~a' -> fa0;
         'circuit_foo~0~b' -> fb0
         'circuit_foo~0~c' -> fc0
         'circuit_foo~1~a' -> fa1
         'circuit_foo~1~b' -> fb1
         'circuit_foo~1~c' -> fc1
         'circuit_bar~0~a' -> ga0
         'circuit_bar~0~b' -> gb0
         'circuit_bar~0~c' -> gc0
       }
    *)
    val gather_maps : ?shifts_map:(int * int) t -> 'a t list t -> 'a t

    (* Filter a map keeping the elements whose key corresponds to the given
       circuit name *)
    val filter_by_circuit_name : string -> 'a t -> 'a t

    (* [select_answers_by_circuit circuit_name s_map_map] takes a [circuit_name]
       and a map with the structure:
       { 'x' -> { 'circuit_foo~0~a' -> [scalar] ;
                  'circuit_foo~0~b' -> [scalar] ;
                   ...
                }
       }
       and filters the keys of the inner map, keeping the elements whose key
       corresponds to the given circuit name. *)
    val select_answers_by_circuit : string -> 'a t t -> 'a t t

    (* [add_map_list_map m1 m2] will merge [m1] & [m2] ; the resulting map will contain the same keys as [m1] ; [m1] & [m2] can be disjoint, if a key is not found in [m2], the resulting map contains the same binding as [m1] for this key  *)
    val add_map_list_map : 'a t list t -> 'a t list t -> 'a t list t
  end
end

include (StringMap : S)