Source file utils.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 S = Csir.Scalar

(* Difference between the scalar order and the succeeding power of 2 *)
let alpha = Z.(shift_left one (Z.numbits S.order) - S.order)

let bitlist : ?le:bool -> bytes -> bool list =
 fun ?(le = false) b ->
  let l = Bytes.length b in
  let start = if le then 0 else l - 1 in
  let stop = if le then l else -1 in
  let next a = if le then a + 1 else a - 1 in
  let rec loop_byte acc n =
    if n = stop then acc
    else
      let byte = Bytes.get_uint8 b n in
      let rec loop_bit acc m =
        if m = 8 then acc
        else
          let mask = 1 lsl m in
          let bit = byte land mask in
          let bit = if bit = 0 then false else true in
          loop_bit (bit :: acc) (m + 1)
      in
      let acc = loop_bit acc 0 in
      loop_byte acc (next n)
  in
  List.rev @@ loop_byte [] start

let bytes_of_hex hs =
  let h = `Hex hs in
  Hex.to_bytes h

let bool_list_to_scalar : bool list -> S.t =
 fun b_list ->
  let res, _ =
    List.fold_left
      (fun (acc_res, acc_p) b ->
        let acc_res = if b then S.(acc_res + acc_p) else acc_res in
        let acc_p = S.double acc_p in
        (acc_res, acc_p))
      (S.zero, S.one) b_list
  in
  res

(* We use little endian notation (the lsb is on the head) *)
let bool_list_to_z : bool list -> Z.t =
 fun b_list ->
  let res, _ =
    List.fold_left
      (fun (acc_res, acc_p) b ->
        let acc_res = if b then Z.(acc_res + acc_p) else acc_res in
        let acc_p = Z.(acc_p + acc_p) in
        (acc_res, acc_p))
      (Z.zero, Z.one) b_list
  in
  res

(* We use little endian notation (the lsb is on the head) *)
let bool_list_of_z : ?nb_bits:int -> Z.t -> bool list =
 fun ?nb_bits z ->
  let two = Z.of_int 2 in
  let rec aux bits z = function
    | 0 -> List.rev bits
    | n ->
        let b = Z.(equal (z mod two) one) in
        aux (b :: bits) (Z.div z two) (n - 1)
  in
  aux [] z @@ Option.value ~default:(Z.numbits z) nb_bits

let rec transpose = function
  | [] | [] :: _ -> []
  | rows -> List.(map hd rows :: (transpose @@ map tl rows))

let tables_cs_encoding : (string list * Csir.CS.t) Data_encoding.t =
  Data_encoding.(obj2 (req "tables" (list string)) (req "cs" Csir.CS.encoding))

let save_cs_to_file path tables cs =
  let serialized_json =
    Data_encoding.Json.construct tables_cs_encoding (tables, cs)
    |> Data_encoding.Json.to_string
  in
  let outc = open_out path in
  Printf.fprintf outc "%s" serialized_json;
  close_out outc

let load_cs_from_file path =
  if not (Sys.file_exists path) then
    raise
    @@ Invalid_argument
         (Printf.sprintf "load_cs_from_file: %s does not exist." path);
  let inc = open_in path in
  let content = really_input_string inc (in_channel_length inc) in
  close_in inc;
  match Data_encoding.Json.from_string content with
  | Error _ -> failwith ("Error loading file: " ^ path)
  | Ok json -> Data_encoding.Json.destruct tables_cs_encoding json

let trace_info_encoding : Optimizer.trace_info Data_encoding.t =
  let open Optimizer in
  Data_encoding.(
    conv
      (fun { free_wires; assignments } -> (free_wires, assignments))
      (fun (free_wires, assignments) -> { free_wires; assignments })
      (obj2
         (req "free_vars" (list int31))
         (req "assignments"
            (list (tup2 int31 (list (tup2 Csir.Scalar.encoding int31)))))))

let cs_ti_encoding = Data_encoding.tup2 Csir.CS.encoding trace_info_encoding

let get_circuit_id cs =
  let serialized_bytes =
    Data_encoding.Binary.to_bytes_exn Csir.CS.encoding cs
  in
  Hacl_star.Hacl.Blake2b_32.hash serialized_bytes 32 |> Hex.of_bytes |> Hex.show

let circuit_dir = "/tmp/plompiler"

let circuit_path s =
  if not @@ Sys.file_exists circuit_dir then Sys.mkdir circuit_dir 0o755;
  circuit_dir ^ "/" ^ s

let dump_label_traces path (cs : Csir.CS.t) =
  let outc = open_out path in
  List.iter
    Csir.CS.(
      Array.iter (fun c ->
          Printf.fprintf outc "%s 1\n" @@ String.concat "; " (List.rev c.label)))
    cs;
  close_out outc