Source file permutation_gate.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 L = Plompiler.LibCircuit

module Permutation_gate_impl
    (PP : Polynomial_protocol.S with type PC.Scalar.t = Plompiler.S.t) =
struct
  module PP = PP
  module Domain = PP.PC.Polynomial.Domain
  module Poly = PP.PC.Polynomial.Polynomial
  module Scalar = PP.PC.Scalar
  module Commitment = PP.PC.Commitment
  module Fr_generation = PP.PC.Fr_generation
  module Evaluations = PP.Evaluations

  let z_name = "Z"
  let zg_name = "Zg"

  (* element preprocessed and known by both prover and verifier *)
  type public_parameters = {
    g_map_perm_PP : Poly.t SMap.t;
    cm_g_map_perm_PP : Commitment.t SMap.t;
    s_poly_map : Poly.t SMap.t;
    cm_s_poly_map : Commitment.t SMap.t;
    permutation : int array;
  }

  let srs_size ~zero_knowledge ~n = if zero_knowledge then n + 9 else n
  let one = Scalar.one
  let zero = Scalar.zero
  let mone = Scalar.negate one
  let quadratic_non_residues = Fr_generation.build_quadratic_non_residues 8

  let get_k k =
    if k < 8 then quadratic_non_residues.(k)
    else raise (Invalid_argument "Permutation.get_k : k must be lower than 8.")

  module Preprocessing = struct
    (* returns l1 polynomial such that l1(generator) = 1 & l1(a) = 0 for all a != generator in H *)
    let l1 domain =
      let size_domain = Domain.length domain in
      let scalar_list =
        Array.append [| zero; one |]
          Array.(init (size_domain - 2) (fun _ -> zero))
      in
      Evaluations.interpolation_fft2 domain scalar_list

    (* returns [sid_0, …, sid_k] *)
    let sid_list_non_quadratic_residues size =
      if size > 8 then
        raise (Failure "sid_list_non_quadratic_residues: sid list too long")
      else List.init size (fun i -> Poly.of_coefficients [ (get_k i, 1) ])

    let sid_map_non_quadratic_residues_prover size =
      if size > 8 then
        raise (Failure "sid_map_non_quadratic_residues: sid map too long")
      else
        SMap.of_list
          (List.init size (fun i ->
               let k = get_k i in
               ("Si" ^ string_of_int (i + 1), Poly.of_coefficients [ (k, 1) ])))

    let evaluations_sid size domain_evals =
      SMap.of_list
        (List.init size (fun i ->
             let k = get_k i in
             let evals =
               Evaluations.mul_by_scalar k (Evaluations.of_domain domain_evals)
             in
             ("Si" ^ string_of_int (i + 1), evals)))

    let ssigma_map_non_quadratic_residues ~prefix permutation domain size =
      let n = Domain.length domain in
      let ssigma_map =
        SMap.of_list
          (List.init size (fun i ->
               let offset = i * n in
               let coeff_list =
                 Array.init n (fun j ->
                     let s_ij = permutation.(offset + j) in
                     let coeff = get_k (s_ij / n) in
                     let index = s_ij mod n in
                     Scalar.mul coeff (Domain.get domain index))
               in
               ( prefix ^ "Ss" ^ string_of_int (i + 1),
                 Evaluations.interpolation_fft2 domain coeff_list )))
      in
      ssigma_map
  end

  module Permutation_poly = struct
    (* compute f' & g' = (f + β×Sid + γ) & (g + β×Sσ + γ) products with Z *)
    (* compute_prime computes the following
       z_name * (w_1 + beta * s_1 + gamma) * ... * (w_n + beta * s_n + gamma)
       - z_name could be either "Z" or "Zg"
       - evaluations contains "Z" but not "Zg"
       - if z_name = "Zg", we compute "Zg" as composition_gx of "Z" with 1 *)
    let compute_prime ~prefix res_evaluation tmp_evaluation tmp2_evaluation beta
        gamma evaluations wire_names s_names this_z_name n =
      let z_evaluation =
        Evaluations.find_evaluation evaluations (prefix z_name)
      in

      let _i, res_evaluation =
        let f_fold (i, acc_evaluation) wire_name s_name =
          let comp = if i = 0 && this_z_name = zg_name then 1 else 0 in
          let res_evaluation =
            (* tmp_evaluation <- wire_name + beta * s_name + gamma *)
            let evaluation_linear_i =
              Evaluations.linear ~res:tmp_evaluation ~evaluations
                ~poly_names:[ wire_name; s_name ] ~linear_coeffs:[ one; beta ]
                ~add_constant:gamma ()
            in
            (* tmp2_evaluation <- acc_evaluation * evaluation_linear_i *)
            let acc_evaluation_new =
              Evaluations.mul_c ~res:tmp2_evaluation
                ~evaluations:[ evaluation_linear_i; acc_evaluation ]
                ~composition_gx:([ 0; comp ], n)
                ()
            in
            Evaluations.copy ~res:res_evaluation acc_evaluation_new
          in
          (i + 1, res_evaluation)
        in
        List.fold_left2 f_fold (0, z_evaluation) wire_names s_names
      in
      res_evaluation

    (* evaluations must contain z’s evaluation *)
    let precompute_perm_identity_poly ~prefix wire_names beta gamma n
        evaluations =
      let z_evaluation =
        Evaluations.find_evaluation evaluations (prefix z_name)
      in
      let z_evaluation_len = Evaluations.length z_evaluation in
      let tmp_evaluation = Evaluations.create z_evaluation_len in
      let tmp2_evaluation = Evaluations.create z_evaluation_len in
      let id1_evaluation = Evaluations.create z_evaluation_len in
      let id2_evaluation = Evaluations.create z_evaluation_len in

      let wire_names = List.map prefix wire_names in

      let identity_zfg =
        let nb_wires = List.length wire_names in

        (* changes f (resp g) array to f'(resp g') array, and multiply them together
            and with z (resp zg) *)
        let f_evaluation =
          let sid_names =
            List.init nb_wires (fun i -> "Si" ^ string_of_int (i + 1))
          in
          compute_prime ~prefix tmp_evaluation id2_evaluation tmp2_evaluation
            beta gamma evaluations wire_names sid_names z_name n
        in
        let g_evaluation =
          let ss_names =
            List.init nb_wires (fun i -> prefix @@ "Ss" ^ string_of_int (i + 1))
          in
          compute_prime ~prefix id2_evaluation id1_evaluation tmp2_evaluation
            beta gamma evaluations wire_names ss_names zg_name n
        in
        Evaluations.linear_c ~res:id1_evaluation
          ~evaluations:[ f_evaluation; g_evaluation ]
          ~linear_coeffs:[ one; mone ] ()
      in
      let identity_l1_z =
        let l1_evaluation = Evaluations.find_evaluation evaluations "L1" in
        let z_mone_evaluation =
          Evaluations.linear_c ~res:tmp_evaluation ~evaluations:[ z_evaluation ]
            ~add_constant:mone ()
        in

        Evaluations.mul_c ~res:id2_evaluation
          ~evaluations:[ l1_evaluation; z_mone_evaluation ]
          ()
      in
      SMap.of_list
        [ (prefix "Perm.a", identity_l1_z); (prefix "Perm.b", identity_zfg) ]

    (* compute_Z performs the following steps in the two loops.
       ----------------------
       | f_11 f_21 ... f_k1 | -> f_prod_1 (no need to compute as Z(g) is always one)
       | f_12 f_22 ... f_k2 | -> f_prod_2 = f_12 * f_22 * ... * f_k2
       |     ...........    | -> ...
       | f_1n f_2n ... f_kn | -> f_prod_n = f_1n * f_2n * ... * f_kn
        --------------------
       1. compute f_res = [ f_prod_1; f_prod_2; ...; f_prod_n ]
       2. compute g_res = [ g_prod_1; g_prod_2; ...; g_prod_n ]
       3. compute f_over_g = [ f_prod_1 / g_prod_1; ...; f_prod_n / g_prod_n ]
       4. apply fold_mul_array to f_over_g:
          [f_over_g_1; f_over_g_1 * f_over_g_2; ..; f_over_g_1 * f_over_g_2 * .. * f_over_n ]
       5. as step 4 computes [Z(g); Z(g^2); ..; Z(g^n)], we need to do a rotate right by 1
          (i.e., composition_gx with n - 1): [Z(g^n); Z(g); Z(g^2); ..; Z(g^{n-1})] *)
    let compute_Z s domain beta gamma values indices =
      let size_domain = Domain.length domain in
      let scalar_array_Z =
        let indices = Array.of_list (List.map snd (SMap.bindings indices)) in
        let size_res = Array.length indices.(0) in
        assert (size_res = size_domain);
        let g_res = Array.init size_res (fun _ -> Scalar.zero) in
        let f_prev = ref Scalar.one in
        let f_res = ref Scalar.one in
        let tmp = Scalar.(copy one) in
        (* the first element of scalar_array_Z is always one *)
        for i = 1 to size_res - 1 do
          for j = 0 to Array.length indices - 1 do
            let indices_list_j_i = indices.(j).(i) in
            let v_gamma =
              Scalar.add gamma @@ Evaluations.get values indices_list_j_i
            in
            let f_coeff =
              let gi = Domain.get domain i in
              Scalar.(
                mul_inplace tmp gi (get_k j);
                mul_inplace gi tmp beta;
                add_inplace gi gi v_gamma;
                gi)
            in
            let g_coeff =
              let sj = s.((j * size_domain) + i) in
              let gj = Domain.get domain (sj mod size_domain) in
              Scalar.(
                mul_inplace tmp gj (get_k (Int.div sj size_domain));
                mul_inplace gj tmp beta;
                add_inplace gj gj v_gamma;
                gj)
            in
            if j = 0 then (
              f_res := f_coeff;
              g_res.(i) <- g_coeff)
            else
              Scalar.(
                mul_inplace !f_res !f_res f_coeff;
                mul_inplace g_res.(i) g_res.(i) g_coeff)
          done;
          let f_over_g = Scalar.div_exn !f_res g_res.(i) in
          Scalar.(
            mul_inplace f_over_g f_over_g !f_prev;
            g_res.(i) <- !f_prev;
            f_prev := f_over_g)
        done;

        g_res.(0) <- !f_prev;
        g_res
      in
      Evaluations.interpolation_fft2 domain scalar_array_Z
  end

  (* max degree needed is the degree of Perm.b, which is sum of wire’s degree plus z degree *)
  let polynomials_degree ~nb_wires = nb_wires + 1

  (* d = polynomials’ max degree
     n = generator’s order
     Returns SRS of decent size, preprocessed polynomials for permutation and
     their commitments (g_map_perm, cm_g_map_perm (="L1" -> L₁, preprocessed
     polynomial for verify perm’s identity), s_poly_map, cm_s_poly_map) & those
     for PP (g_map_PP, cm_g_map_PP)
     permutation for ssigma_list computation is deducted of cycles
     Details for SRS size :
       max size needed is deg(T)+1
       v polynomials all have degree 1
       according to identities_list_perm[0], t has max degree of Z×fL×fR×fO ;
       interpolation makes polynomials of degree n-1, so Z has degree of X²×Zh =
       X²×(X^n - 1) which is n+2, and each f has degree of X×Zh so n+1
       As a consequence, deg(T)-deg(Zs) = (n+2)+3(n+1) - n = 3n+5
       (for gates’ identity verification, max degree is degree of qM×fL×fR which
       is (n-1)+(n+1)+(n+1) < 3n+5) *)
  let preprocessing ?(circuit_name = "") ~domain ~permutation ~nb_wires () =
    let prefix = SMap.Aggregation.add_prefix circuit_name "" in
    Preprocessing.ssigma_map_non_quadratic_residues ~prefix permutation domain
      nb_wires

  let common_preprocessing ~compute_l1 ~domain ~nb_wires ~domain_evals =
    let sid_evals = Preprocessing.evaluations_sid nb_wires domain_evals in
    let g_map_perm_PP =
      if not compute_l1 then SMap.empty
      else SMap.singleton "L1" (Preprocessing.l1 domain)
    in
    (g_map_perm_PP, sid_evals)

  let prover_identities ?(circuit_name = "") ~wire_names ~generator:_ ~beta
      ~gamma ~n () : PP.prover_identities =
    let prefix = SMap.Aggregation.add_prefix circuit_name in
    fun evaluations ->
      Permutation_poly.precompute_perm_identity_poly ~prefix wire_names beta
        gamma n evaluations

  let verifier_identities ?(circuit_name = "") ~nb_proofs ~generator ~n
      ~wire_names ~beta ~gamma ~delta () : PP.verifier_identities =
    let prefix = SMap.Aggregation.add_prefix circuit_name in
    let prefix_j j =
      SMap.Aggregation.add_prefix ~n:nb_proofs ~i:j circuit_name
    in
    fun x answers ->
      let get_ss i =
        PP.(get_answer answers X (prefix @@ "Ss" ^ string_of_int (i + 1)))
      in
      (* compute the delta-aggregated wire evaluations at x for each wire name *)
      let batched =
        let wire_j w j = PP.(get_answer answers X @@ prefix_j j w) in
        List.map
          (fun w -> Fr_generation.batch delta (List.init nb_proofs (wire_j w)))
          wire_names
      in
      let z = PP.(get_answer answers X (prefix z_name)) in
      let zg = PP.(get_answer answers GX (prefix z_name)) in
      (* compute the first identity: (Z(x) - 1) * L1(x) *)
      let res1 =
        let n = Z.of_int n in
        let l1_num = Scalar.(generator * sub (pow x n) one) in
        let l1_den = Scalar.(of_z n * sub x generator) in
        Scalar.(sub z one * div_exn l1_num l1_den)
      in
      (* compute the second identity *)
      let res2 =
        let z_factors =
          List.mapi Scalar.(fun i w -> w + (beta * get_k i * x) + gamma) batched
        in
        let zg_factors =
          List.mapi Scalar.(fun i w -> w + (beta * get_ss i) + gamma) batched
        in
        let multiply l = List.fold_left Scalar.mul (List.hd l) (List.tl l) in
        Scalar.sub
          (multiply @@ (z :: z_factors))
          (multiply @@ (zg :: zg_factors))
      in
      SMap.of_list [ (prefix "Perm.a", res1); (prefix "Perm.b", res2) ]

  let f_map_contribution ~permutation ~values ~indices ~beta ~gamma ~domain =
    let z_poly =
      Permutation_poly.compute_Z permutation domain beta gamma values indices
    in
    SMap.of_list [ (z_name, z_poly) ]

  let cs ~sum_alpha_i ~l1 ~ss1 ~ss2 ~ss3 ~beta ~gamma ~delta ~x ~z ~zg ~wires =
    let open L in
    let a_list, b_list, c_list =
      let rec aux (acc_a, acc_b, acc_c) = function
        | [] -> List.(rev acc_a, rev acc_b, rev acc_c)
        | [ a; b; c ] :: tl -> aux (a :: acc_a, b :: acc_b, c :: acc_c) tl
        | _ -> failwith "Unexpected wires format"
      in
      aux ([], [], []) wires
    in
    let* cs_perm_a = Num.custom ~qr:Scalar.(negate one) ~qm:Scalar.(one) z l1 in

    let* a = sum_alpha_i a_list delta in
    let* b = sum_alpha_i b_list delta in
    let* c = sum_alpha_i c_list delta in
    let* betaid1 = Num.mul ~qm:(get_k 0) beta x in
    let* betaid2 = Num.mul ~qm:(get_k 1) beta x in
    let* betaid3 = Num.mul ~qm:(get_k 2) beta x in
    let* betasigma1 = Num.mul beta ss1 in
    let* betasigma2 = Num.mul beta ss2 in
    let* betasigma3 = Num.mul beta ss3 in

    let* aid = Num.add_list (to_list [ a; betaid1; gamma ]) in
    let* bid = Num.add_list (to_list [ b; betaid2; gamma ]) in
    let* cid = Num.add_list (to_list [ c; betaid3; gamma ]) in
    let* asigma = Num.add_list (to_list [ a; betasigma1; gamma ]) in
    let* bsigma = Num.add_list (to_list [ b; betasigma2; gamma ]) in
    let* csigma = Num.add_list (to_list [ c; betasigma3; gamma ]) in

    let* left_term = Num.mul_list (to_list [ aid; bid; cid; z ]) in

    let* right_term = Num.mul_list (to_list [ asigma; bsigma; csigma; zg ]) in
    let* cs_perm_b = Num.add ~qr:Scalar.(negate one) left_term right_term in
    ret (cs_perm_a, cs_perm_b)
end

module type S = sig
  module PP : Polynomial_protocol.S

  val z_name : string
  val zg_name : string
  val srs_size : zero_knowledge:bool -> n:int -> int
  val polynomials_degree : nb_wires:int -> int

  val preprocessing :
    ?circuit_name:string ->
    domain:PP.PC.Polynomial.Domain.t ->
    permutation:int array ->
    nb_wires:int ->
    unit ->
    PP.PC.Polynomial.Polynomial.t SMap.t

  val common_preprocessing :
    compute_l1:bool ->
    domain:PP.PC.Polynomial.Domain.t ->
    nb_wires:int ->
    domain_evals:PP.Evaluations.domain ->
    PP.PC.Polynomial.Polynomial.t SMap.t * PP.Evaluations.t SMap.t

  val prover_identities :
    ?circuit_name:string ->
    wire_names:string list ->
    generator:PP.PC.Scalar.t ->
    beta:PP.PC.Scalar.t ->
    gamma:PP.PC.Scalar.t ->
    n:int ->
    unit ->
    PP.prover_identities

  val verifier_identities :
    ?circuit_name:string ->
    nb_proofs:int ->
    generator:PP.PC.Scalar.t ->
    n:int ->
    wire_names:string list ->
    beta:PP.PC.Scalar.t ->
    gamma:PP.PC.Scalar.t ->
    delta:PP.PC.Scalar.t ->
    unit ->
    PP.verifier_identities

  val f_map_contribution :
    permutation:int array ->
    values:PP.Evaluations.t ->
    indices:int array SMap.t ->
    beta:PP.PC.Scalar.t ->
    gamma:PP.PC.Scalar.t ->
    domain:PP.PC.Polynomial.Domain.t ->
    PP.PC.Polynomial.Polynomial.t SMap.t

  val cs :
    sum_alpha_i:(L.scalar L.repr list -> L.scalar L.repr -> L.scalar L.repr L.t) ->
    l1:L.scalar L.repr ->
    ss1:L.scalar L.repr ->
    ss2:L.scalar L.repr ->
    ss3:L.scalar L.repr ->
    beta:L.scalar L.repr ->
    gamma:L.scalar L.repr ->
    delta:L.scalar L.repr ->
    x:L.scalar L.repr ->
    z:L.scalar L.repr ->
    zg:L.scalar L.repr ->
    wires:L.scalar L.repr list list ->
    (L.scalar L.repr * L.scalar L.repr) L.t
end

module Permutation_gate
    (PP : Polynomial_protocol.S with type PC.Scalar.t = Plompiler.S.t) :
  S with module PP = PP =
  Permutation_gate_impl (PP)