Source file csir.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 Scalar = struct
  include Bls12_381.Fr

  type scalar = t

  let mone = negate one

  let string_of_scalar x =
    if eq x (of_string "-1") then "-1"
    else if eq x (of_string "-2") then "-2"
    else
      let s = to_string x in
      if String.length s > 3 then "h" ^ string_of_int (Z.hash (to_z x)) else s

  let equal a b = Bytes.equal (to_bytes a) (to_bytes b)

  let encoding =
    Data_encoding.(conv to_bytes of_bytes_exn (Fixed.bytes size_in_bytes))
end

(* If multiple tables are used, they all need to have the same number of wires,
   so any smaller one will be padded. *)
module Table : sig
  type t

  val empty : t
  val size : t -> int

  type entry = { a : Scalar.t; b : Scalar.t; c : Scalar.t }

  type partial_entry = {
    a : Scalar.t option;
    b : Scalar.t option;
    c : Scalar.t option;
  }

  val mem : entry -> t -> bool
  val find : partial_entry -> t -> entry option
  val to_list : t -> Scalar.t array list
  val of_list : Scalar.t array list -> t
end = struct
  (* Rows are variables, columns are entries in the table.
     If the table is full it would be |domain|^#variables e.g. 2^3=8
     Example OR gate:
     [
       [|0; 0; 1; 1|] ;
       [|0; 1; 0; 1|] ;
       [|0; 1; 1; 1|] ;
     ]
  *)

  type entry = { a : Scalar.t; b : Scalar.t; c : Scalar.t }

  type partial_entry = {
    a : Scalar.t option;
    b : Scalar.t option;
    c : Scalar.t option;
  }

  type t = Scalar.t array array

  let empty = [||]
  let size table = Array.length table.(0)

  (* Function returning the first table corresponding to the input partial entry.
     A partial entry is found on the table at row i if it coincides
     with the table values in all specified (i.e., not None) columns *)
  let find_entry_i : partial_entry -> t -> int -> entry option =
   fun pe table i ->
    let match_partial_entry o s =
      Option.(value ~default:true @@ map (Scalar.eq s) o)
    in
    if
      match_partial_entry pe.a table.(0).(i)
      && match_partial_entry pe.b table.(1).(i)
      && match_partial_entry pe.c table.(2).(i)
    then Some { a = table.(0).(i); b = table.(1).(i); c = table.(2).(i) }
    else None

  let find pe table =
    (* TODO make it a binary search *)
    let sz = size table in
    let rec aux i =
      match i with
      | 0 -> find_entry_i pe table 0
      | _ ->
          let o = find_entry_i pe table i in
          if Option.is_some o then o else aux (i - 1)
    in
    aux (sz - 1)

  let mem : entry -> t -> bool =
   fun e table ->
    match find { a = Some e.a; b = Some e.b; c = Some e.c } table with
    | Some _ -> true
    | None -> false

  let to_list table =
    Format.printf "\n%i %i\n" (Array.length table) (Array.length table.(0));
    Array.to_list table

  let of_list table = Array.of_list table
end

let table_or =
  Table.of_list
    Scalar.
      [
        [| zero; zero; one; one |];
        [| zero; one; zero; one |];
        [| zero; one; one; one |];
      ]

module Tables = Map.Make (String)

let table_registry = Tables.add "or" table_or Tables.empty

module CS = struct
  let q_list ?q_table ~qc ~ql ~qr ~qo ~qlg ~qrg ~qog ~qm ~qx5 ~qecc_ws_add
      ~qecc_ed_add ~q_plookup () =
    let base =
      [
        ("qc", qc);
        ("ql", ql);
        ("qr", qr);
        ("qo", qo);
        ("qlg", qlg);
        ("qrg", qrg);
        ("qog", qog);
        ("qm", qm);
        ("qx5", qx5);
        ("qecc_ws_add", qecc_ws_add);
        ("qecc_ed_add", qecc_ed_add);
        ("q_plookup", q_plookup);
      ]
    in
    Option.(map (fun q -> ("q_table", q)) q_table |> to_list) @ base

  type selector_tag = Linear | ThisConstr | NextConstr | WireA | WireB | WireC

  let all_selectors =
    q_list ~qc:[ ThisConstr ]
      ~ql:[ ThisConstr; Linear; WireA ]
      ~qr:[ ThisConstr; Linear; WireB ]
      ~qo:[ ThisConstr; Linear; WireC ]
      ~qlg:[ NextConstr; Linear; WireA ]
      ~qrg:[ NextConstr; Linear; WireB ]
      ~qog:[ NextConstr; Linear; WireC ]
      ~qm:[ ThisConstr; WireA; WireB ]
      ~qx5:[ ThisConstr; WireA ]
      ~qecc_ws_add:[ ThisConstr; NextConstr; WireA; WireB; WireC ]
      ~qecc_ed_add:[ ThisConstr; NextConstr; WireA; WireB; WireC ]
      ~q_plookup:[ ThisConstr; WireA; WireB; WireC ]
      ~q_table:[ ThisConstr; WireA; WireB; WireC ]
      ()

  let selectors_with_tags tags =
    List.filter
      (fun (_, sel_tags) -> List.for_all (fun t -> List.mem t sel_tags) tags)
      all_selectors
    |> List.map fst

  let this_constr_selectors = selectors_with_tags [ ThisConstr ]
  let next_constr_selectors = selectors_with_tags [ NextConstr ]
  let this_constr_linear_selectors = selectors_with_tags [ ThisConstr; Linear ]
  let next_constr_linear_selectors = selectors_with_tags [ NextConstr; Linear ]

  type raw_constraint = {
    a : int;
    b : int;
    c : int;
    sels : (string * Scalar.t) list;
    label : string list;
  }

  type gate = raw_constraint array
  type t = gate list

  let selectors_encoding = Data_encoding.(list (tup2 string Scalar.encoding))

  let raw_constraint_encoding : raw_constraint Data_encoding.t =
    Data_encoding.(
      conv
        (fun { a; b; c; sels; label } -> (a, b, c, sels, label))
        (fun (a, b, c, sels, label) -> { a; b; c; sels; label })
        (obj5 (req "a" int31) (req "b" int31) (req "c" int31)
           (req "sels" selectors_encoding)
           (req "label" (list string))))

  let gate_encoding = Data_encoding.array raw_constraint_encoding
  let encoding = Data_encoding.list gate_encoding
  let cs_pub_size_encoding = Data_encoding.(tup2 encoding int31)

  let q_list ?q_table ~qc ~ql ~qr ~qo ~qlg ~qrg ~qog ~qm ~qx5 ~qecc_ws_add
      ~qecc_ed_add ~q_plookup () =
    let base =
      [
        ("qc", qc);
        ("ql", ql);
        ("qr", qr);
        ("qo", qo);
        ("qlg", qlg);
        ("qrg", qrg);
        ("qog", qog);
        ("qm", qm);
        ("qx5", qx5);
        ("qecc_ws_add", qecc_ws_add);
        ("qecc_ed_add", qecc_ed_add);
        ("q_plookup", q_plookup);
      ]
    in
    Option.(map (fun q -> ("q_table", q)) q_table |> to_list) @ base

  let new_constraint ~a ~b ~c ?qc ?ql ?qr ?qo ?qlg ?qrg ?qog ?qm ?qx5
      ?qecc_ws_add ?qecc_ed_add ?q_plookup ?q_table ?(labels = []) label =
    let sels =
      List.filter_map
        (fun (l, x) -> Option.bind x (fun c -> Some (l, c)))
        (q_list ~qc ~ql ~qr ~qo ~qlg ~qrg ~qog ~qm ~qx5 ~qecc_ws_add
           ~qecc_ed_add ~q_plookup ~q_table ())
    in
    { a; b; c; sels; label = label :: labels }

  let get_sel sels s =
    match List.find_opt (fun (x, _) -> s = x) sels with
    | None -> Scalar.zero
    | Some (_, c) -> c

  let to_string_raw_constraint { a; b; c; label; sels } : string =
    let selectors =
      String.concat " "
        (List.map (fun (s, c) -> s ^ ":" ^ Scalar.string_of_scalar c) sels)
    in
    Format.sprintf "a:%i b:%i c:%i %s %s" a b c selectors
      (String.concat ";" label)

  let to_string_gate g =
    String.concat "\n" @@ Array.to_list @@ Array.map to_string_raw_constraint g

  let to_string cs =
    List.fold_left (fun acc con -> acc ^ to_string_gate con ^ "\n") "" cs

  let is_linear_raw_constr constr =
    let linear_selectors =
      ("qc" :: this_constr_linear_selectors) @ next_constr_linear_selectors
    in
    let is_linear_sel (s, _q) = List.mem s linear_selectors in
    List.for_all is_linear_sel constr.sels

  let used_selectors gate i =
    let this_sels = gate.(i).sels in
    let prev_sels = if i = 0 then [] else gate.(i - 1).sels in
    List.filter (fun (s, _) -> List.mem s this_constr_selectors) this_sels
    @ List.filter (fun (s, _) -> List.mem s next_constr_selectors) prev_sels

  let wires_of_constr_i gate i =
    let a_selectors = selectors_with_tags [ WireA ] in
    let b_selectors = selectors_with_tags [ WireB ] in
    let c_selectors = selectors_with_tags [ WireC ] in
    let intersect names = List.exists (fun (s, _q) -> List.mem s names) in
    let sels = used_selectors gate i in
    List.map2
      (fun wsels w -> if intersect wsels sels then w else -1)
      [ a_selectors; b_selectors; c_selectors ]
      [ gate.(i).a; gate.(i).b; gate.(i).c ]

  let gate_wires gate =
    List.init (Array.length gate) (wires_of_constr_i gate)
    |> List.concat |> List.sort_uniq Int.compare
    |> List.filter (fun x -> x >= 0)

  (* the relationship of this function wrt is_linear_raw_constr is a bit weird *)
  let linear_terms constr =
    if not @@ is_linear_raw_constr constr then
      raise @@ Invalid_argument "constraint is non-linear"
    else
      List.map
        (fun (sel_name, coeff) ->
          match sel_name with
          | "qc" -> (coeff, -1)
          | "ql" -> (coeff, constr.a)
          | "qr" -> (coeff, constr.b)
          | "qo" -> (coeff, constr.c)
          | _ -> assert false)
        constr.sels
      |> List.filter (fun (q, _) -> not @@ Scalar.is_zero q)

  let mk_linear_constr (wires, sels) =
    match wires with
    | [ a; b; c ] -> { a; b; c; sels; label = [ "linear" ] }
    | _ -> assert false

  let raw_constraint_equal c1 c2 =
    c1.a = c2.a && c1.b = c2.b && c1.c = c2.c && c1.label = c2.label
    && List.for_all2
         (fun (name, coeff) (name', coeff') ->
           name = name' && Scalar.eq coeff coeff')
         c1.sels c2.sels
end