Source file circuit.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.                                                 *)
(*                                                                           *)
(*****************************************************************************)

open Plompiler.Csir

module Gates =
  Custom_gates.Make ((* TODO this should be more generic *) Polynomial_protocol)

(* We assert here that all modules/selectors have been used.
   The "+2" is to take into account lookup selectors which are not defined in
   plonk/gates/custom_gates.ml
   The "-1" is to remove the Public_gate that is not used *)
let () = assert (List.length CS.all_selectors = Gates.nb_custom_gates - 1 + 2)

let gates_to_string m =
  SMap.fold
    (fun k v s ->
      s ^ k ^ " " ^ String.concat "," (List.map Scalar.to_string v) ^ "\n")
    m ""

module Circuit : sig
  type t = private {
    wires : int array SMap.t;
    gates : Scalar.t array SMap.t;
    tables : Scalar.t array list list;
    public_input_size : int;
    circuit_size : int;
    nb_wires : int;
    table_size : int;
    nb_lookups : int;
    ultra : bool;
  }

  val make_wires :
    a:int list ->
    b:int list ->
    c:int list ->
    ?d:int list ->
    ?e:int list ->
    ?f:int list ->
    ?g:int list ->
    ?h:int list ->
    unit ->
    int list SMap.t

  val make_gates :
    ?qc:Scalar.t list ->
    ?ql:Scalar.t list ->
    ?qr:Scalar.t list ->
    ?qo:Scalar.t list ->
    ?qlg:Scalar.t list ->
    ?qrg:Scalar.t list ->
    ?qog:Scalar.t list ->
    ?qm:Scalar.t list ->
    ?qx2b:Scalar.t list ->
    ?qx5a:Scalar.t list ->
    ?qx5c:Scalar.t list ->
    ?qecc_ws_add:Scalar.t list ->
    ?qecc_ed_add:Scalar.t list ->
    ?q_plookup:Scalar.t list ->
    ?q_table:Scalar.t list ->
    unit ->
    Scalar.t list SMap.t

  val make :
    wires:int list SMap.t ->
    gates:Scalar.t list SMap.t ->
    ?tables:Scalar.t array list list ->
    public_input_size:int ->
    unit ->
    t

  val get_nb_of_constraints : t -> int

  (* /////////////////////////////////////////////////////////////////////// *)

  val get_selectors : t -> string list
  val sat : CS.t -> Table.t list -> Scalar.t array -> bool
  val to_plonk : public_input_size:int -> ?tables:Table.t list -> CS.t -> t
end = struct
  type t = {
    wires : int array SMap.t;
    gates : Scalar.t array SMap.t;
    tables : Scalar.t array list list;
    public_input_size : int;
    circuit_size : int;
    nb_wires : int;
    table_size : int;
    nb_lookups : int;
    ultra : bool;
  }

  let get_selectors circuit = List.map fst (SMap.bindings circuit.gates)

  let make_wires ~a ~b ~c ?(d = []) ?(e = []) ?(f = []) ?(g = []) ?(h = []) () =
    (* Filtering and mapping selectors with labels. *)
    let wire_map = SMap.of_list [ ("a", a); ("b", b); ("c", c) ] in
    let add_map map (label, l) = if l = [] then map else SMap.add label l map in
    List.fold_left add_map wire_map
      [ ("d", d); ("e", e); ("f", f); ("g", g); ("h", h) ]

  let make_gates ?(qc = []) ?(ql = []) ?(qr = []) ?(qo = []) ?(qlg = [])
      ?(qrg = []) ?(qog = []) ?(qm = []) ?(qx2b = []) ?(qx5a = []) ?(qx5c = [])
      ?(qecc_ws_add = []) ?(qecc_ed_add = []) ?(q_plookup = []) ?(q_table = [])
      () =
    (* Filtering and mapping selectors with labels. *)
    let gate_list =
      CS.q_list ~qc ~ql ~qr ~qo ~qlg ~qrg ~qog ~qm ~qx2b ~qx5a ~qx5c
        ~qecc_ws_add ~qecc_ed_add ~q_plookup ()
    in

    let add_map map (label, q) =
      match q with
      | [] -> map
      | l -> if List.for_all Scalar.is_zero l then map else SMap.add label q map
    in
    let base =
      if q_table = [] then SMap.empty else SMap.singleton "q_table" q_table
    in
    List.fold_left add_map base gate_list

  let verify_name name i =
    assert (i < 26);
    let alphabet = "abcdefghijklmnopqrstuvwxyz" in
    let letter_i = alphabet.[i] in
    let msg =
      Printf.sprintf "%d-th wire must be named '%c' (current name is '%s')." i
        letter_i name
    in
    if String.length name <> 1 then raise (Invalid_argument msg)
    else if Char.equal name.[0] letter_i then ()
    else raise (Invalid_argument msg)

  (* For efficiency reason this function does not check for all-zero selectors,
     it's responsability of the caller to filter them out.
     If public_input_size is greater than 0, selector ql will be added if not
     already present.
     Wires and gates cannot be empty and must all have the same length.
  *)
  let make ~wires ~gates ?(tables = []) ~public_input_size () =
    if SMap.is_empty wires then
      raise @@ Invalid_argument "Make Circuit: empty wires.";
    if SMap.is_empty gates then
      raise @@ Invalid_argument "Make Circuit: empty gates.";
    let circuit_size = List.length (snd (SMap.choose wires)) in
    if Int.equal circuit_size 0 then
      raise (Invalid_argument "Make Circuit: empty circuit.");

    (* Check that all wires have same size, and each wire i is named
       as the i-th alphabet’s letter. *)
    let nb_wires = SMap.cardinal wires in
    let () =
      List.iteri
        (fun i (name, l) ->
          verify_name name i;
          if List.compare_length_with l circuit_size = 0 then ()
          else raise (Invalid_argument "Make Circuit: different length wires."))
        (SMap.bindings wires)
    in
    (* Check that all selectors have the same size, and that they are available. *)
    let () =
      SMap.iter
        (fun label q ->
          if List.compare_length_with q circuit_size = 0 then ()
          else raise (Invalid_argument "Make Circuit: different length gates.");
          if List.mem label (List.map fst CS.all_selectors) then ()
          else raise (Invalid_argument "Make Circuit: unknown gates."))
        gates
    in
    (* Check all tables' columns have the same size. *)
    let () =
      List.iter
        (fun l ->
          let sub_table_size = Array.length (List.hd l) in
          if List.length l > nb_wires then
            raise
              (Invalid_argument "Make Circuit: table(s) with too many columns.");
          List.iter
            (fun t ->
              if Array.length t != sub_table_size then
                raise
                  (Invalid_argument
                     "Make Circuit: table(s) with columns of different length.")
              else ())
            l)
        tables
    in
    let table_size =
      if tables = [] then 0
      else List.fold_left (fun acc t -> acc + Array.length (List.hd t)) 0 tables
    in
    (* Determining if UltraPlonk or TurboPlonk needs to be used. *)
    let ultra = SMap.mem "q_plookup" gates in
    let nb_lookups =
      if not ultra then 0
      else
        let q_plookup = SMap.find "q_plookup" gates in
        List.fold_left
          (fun acc qi -> if Scalar.is_zero qi then acc else acc + 1)
          0 q_plookup
    in
    if ultra && not (SMap.mem "q_table" gates) then
      raise (Invalid_argument "Make Circuit: expected table selector.");
    if ultra && tables = [] then
      raise (Invalid_argument "Make Circuit: tables empty.");
    if (not ultra) && (tables != [] || SMap.mem "q_table" gates) then
      raise (Invalid_argument "Make Circuit: table(s) given with no lookups.");
    let wires = Tables.map Array.of_list wires in
    let gates = Tables.map Array.of_list gates in
    let gates =
      (* Define ql if undefined as it is the gate taking the public input in. *)
      if public_input_size > 0 && (not @@ SMap.mem "ql" gates) then
        Tables.add "ql" (Array.init circuit_size (fun _ -> Scalar.zero)) gates
      else gates
    in
    {
      circuit_size;
      wires;
      gates;
      tables;
      public_input_size;
      nb_wires;
      table_size;
      nb_lookups;
      ultra;
    }

  let get_nb_of_constraints cs = Array.length (snd (SMap.choose cs.wires))
  (* ////////////////////////////////////////////////////////// *)

  let sat_gate identities gate trace tables =
    let open CS in
    let nb_cs = Array.length gate in
    let identities =
      (* For each constraint *)
      List.init nb_cs (fun i ->
          (* Retrieving its values as well as the next constraint's values *)
          let j = (i + 1) mod nb_cs in
          let ci, cj = (gate.(i), gate.(j)) in
          let a, b, c = (trace.(ci.a), trace.(ci.b), trace.(ci.c)) in
          let ag, bg, cg = (trace.(cj.a), trace.(cj.b), trace.(cj.c)) in
          (* Folding on selectors *)
          List.fold_left
            (fun id_map (s_name, q) ->
              match s_name with
              | x when x = "q_plookup" -> id_map
              | x when x = "q_table" ->
                  (* We assume there can be only one lookup per gate *)
                  let entry : Table.entry = Table.{ a; b; c } in
                  let sub_table = List.nth tables (Scalar.to_z q |> Z.to_int) in
                  let b = Table.mem entry sub_table in
                  let id = [| (if b then Scalar.zero else Scalar.one) |] in
                  SMap.add "q_table" id id_map
              | _ ->
                  (* Retrieving the selector's identity name and equations *)
                  let s_id_name, _ = Gates.get_ids s_name in
                  let s_ids = SMap.find s_id_name id_map in
                  (* Updating the identities with the equations' output *)
                  List.iteri
                    (fun i s -> s_ids.(i) <- Scalar.(s_ids.(i) + s))
                    ((Gates.get_eqs s_name) ~q ~a ~b ~c ~ag ~bg ~cg ());
                  SMap.add s_id_name s_ids id_map)
            identities ci.sels)
    in
    (* Checking all identities are verified, i.e. the map contains only 0s *)
    List.for_all
      (SMap.for_all (fun _id_name id ->
           let b = Array.for_all Scalar.is_zero id in
           (* if Bool.not b then Printf.printf "\nIdentity '%s' not satisfied" id_name
              else () ; *)
           b))
      identities

  let sat cs tables trace =
    (* We initialise a map with all ids used in the circuit *)
    let identities =
      List.fold_left
        (fun map m ->
          let id, nb_ids = Gates.get_ids m in
          if SMap.mem id map then map
          else SMap.add id (Array.init nb_ids (fun _ -> Scalar.zero)) map)
        (SMap.singleton "q_table" [| Scalar.zero |])
        Gates.gates_list
    in
    let exception Constraint_not_satisfied of string in
    try
      (* We check in each gate, constraint by constraint, that all ids are satisfied *)
      List.iteri
        (fun i gate ->
          (* Printf.printf "\n\nGate %i: %s" i (CS.to_string_gate gate); *)
          let b = sat_gate identities gate trace tables in
          if b then ()
          else
            (* just to exit the iter *)
            raise
              (Constraint_not_satisfied
                 (Printf.sprintf "\nGate #%i not satisfied." i)))
        cs;
      true
    with Constraint_not_satisfied _ -> false

  let to_plonk ~public_input_size ?(tables = []) cs =
    let open CS in
    let cs = List.rev Array.(to_list @@ concat cs) in
    assert (cs <> []);
    let add_wires a b c wires =
      let aa = SMap.find "a" wires in
      let bb = SMap.find "b" wires in
      let cc = SMap.find "c" wires in
      let wires = SMap.add "a" (a :: aa) wires in
      let wires = SMap.add "b" (b :: bb) wires in
      let wires = SMap.add "c" (c :: cc) wires in
      wires
    in
    let add_selectors sels map pad =
      (* Add to the map all new selectors with the coresponding padding
         (array of [pad] zeroes). *)
      let map =
        List.fold_left
          (fun map (k, _) ->
            if SMap.mem k map then map
            else
              let zeros = List.init pad (fun _ -> Scalar.zero) in
              SMap.add k zeros map)
          map sels
      in

      (* Extend every binding in the map by either add the coefficient
         or pad with a zero. *)
      SMap.fold
        (fun label qq map ->
          let q =
            match List.find_opt (fun (s, _) -> s = label) sels with
            | None -> Scalar.zero
            | Some (_, coeff) -> coeff
          in
          SMap.add label (q :: qq) map)
        map map
    in

    let wires_map = SMap.of_list [ ("a", []); ("b", []); ("c", []) ] in
    let selectors_map = SMap.empty in

    List.fold_left
      (fun (wires_map, selectors_map, pad) { a; b; c; sels; _ } ->
        let wires_map = add_wires a b c wires_map in
        let selectors_map = add_selectors sels selectors_map pad in
        (wires_map, selectors_map, pad + 1))
      (wires_map, selectors_map, 0)
      cs
    |> fun (wires, selectors, _) ->
    let tables = List.map Table.to_list tables in
    make ~wires ~gates:selectors ~public_input_size ~tables ()
end

include Circuit