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 Permutation_gate_impl (PP : Polynomial_protocol.Polynomial_protocol_sig) =
struct
  module PP = PP
  module MP = PP.MP
  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 monomial_of_list = SMap.monomial_of_list
  let z = "Z"

  (* 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.")

  (* prime n-th roots of unity (g_2_i^(2^i)=1)
     n is of form 2^i
      generator is an nth root of unity
      permutations = array ; permutations.(i) = σ(i)
  *)

  (* Returns the list of monomial corresponding to the permutation identity*)
  let gate_identity ~prefix wires_name beta gamma =
    let wires_name = Array.map (fun x -> prefix ^ x) wires_name in
    (* Returns for a given tbit and position the labels for the permutation identity, as well as the gamma and beta coefficient *)
    let choose_list wires_name nb n =
      let index = string_of_int (n + 1) in
      match nb with
      | 0 ->
          let f = wires_name.(n) in
          (true, f, f, 0, 0)
      | 1 -> (true, "Si" ^ index, prefix ^ "Ss" ^ index, 1, 0)
      | 2 -> (false, "", "", 0, 1)
      | _ -> raise (Failure "choose_list: Error in ternary decomposition")
    in
    (* Returns the ternary decomposition of a number $nb$ as a $length_array$-element long int list*)
    let ternary nb length_array =
      let rec cpt_three acc length n =
        match n with
        | 0 -> (acc, length)
        | _ -> cpt_three ((n mod 3) :: acc) (length + 1) (n / 3)
      in
      let suffix, suffix_length = cpt_three [] 0 nb in
      List.rev_append
        (List.init (length_array - suffix_length) (fun _ -> 0))
        suffix
    in
    let nb_wires = Array.length wires_name in
    let array_betas = Fr_generation.powers (nb_wires + 1) beta in
    let array_gammas = Fr_generation.powers (nb_wires + 1) gamma in
    let size_identity = Z.(to_int (pow (of_int 3) nb_wires)) in
    let rec compute_identity_zfg acc size =
      let rec cpt_mnmls (acc_z, acc_zg, acc_b, acc_g) base3 ctr =
        match base3 with
        | t :: sub3 ->
            let add, item_z, item_zg, b, g = choose_list wires_name t ctr in
            let acc_b = acc_b + b in
            let acc_g = acc_g + g in
            let acc_z = if add then item_z :: acc_z else acc_z in
            let acc_zg = if add then item_zg :: acc_zg else acc_zg in
            cpt_mnmls (acc_z, acc_zg, acc_b, acc_g) sub3 (ctr + 1)
        | [] ->
            let coeff = Scalar.mul array_betas.(acc_b) array_gammas.(acc_g) in
            ( (SMap.monomial_of_list acc_z, coeff),
              (SMap.monomial_of_list acc_zg, Scalar.negate coeff) )
      in
      if size = size_identity then acc
      else
        let i_base3 = ternary size nb_wires in
        let mnml_z, mnml_zg =
          cpt_mnmls ([ prefix ^ "Z" ], [ prefix ^ "Zg" ], 0, 0) i_base3 0
        in
        compute_identity_zfg (mnml_z :: mnml_zg :: acc) (size + 1)
    in
    let identity_zfg = MP.Polynomial.of_list (compute_identity_zfg [] 0) in
    let identity_l1_z =
      MP.Polynomial.of_list
        [
          (SMap.monomial_of_list [ "L1"; prefix ^ "Z" ], one);
          (SMap.monomial_of_list [ "L1" ], mone);
        ]
    in
    SMap.of_list
      [ (prefix ^ "Perm.a", identity_l1_z); (prefix ^ "Perm.b", identity_zfg) ]

  (* v_map must contain all polynomials that are not involved in permutation checking ;
     those involved in permutation checking are fl, fr, fo *)
  let v_map ~prefix generator =
    let g_poly = Poly.of_coefficients [ (generator, 1) ] in
    (* (name of h∘v, (name of h, value of v))
       h∘v name will be used in identities
       h name is the name of the polynomial in g_map or f_map to compose with v
    *)
    SMap.singleton (prefix ^ "Zg") (prefix ^ "Z", g_poly)

  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
          Fr_generation.[| fr_of_int_safe 0; fr_of_int_safe 1 |]
          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) ])))

    (* Add a [not_committed] variant to represent Scalar.mul k x.(0) *)
    type PP.not_committed += MulK of Evaluations.scalar

    let () =
      (* NB: do not change tag, it will break the encoding *)
      let open Encodings in
      PP.register_nc_eval_and_encoding
        (function
          | MulK k -> Some (fun x_array -> Scalar.mul k x_array.(0)) | _ -> None)
        ~title:"permutation" ~tag:1
        Data_encoding.(obj1 (req "permutation" fr_encoding))
        (function MulK k -> Some k | _ -> None)
        (fun k -> MulK k)

    let sid_map_non_quadratic_residues_verifier 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), MulK k)))

    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 res_evaluation tmp_evaluation beta gamma evaluations
        wires_names s_names z_name n =
      let wires_names = Array.to_list wires_names in
      let z_evaluation = Evaluations.find_evaluation evaluations "Z" in

      let _i, res_evaluation =
        let f_fold (i, acc_evaluation) wire_name s_name =
          let comp = if i = 0 && z_name = "Zg" then 1 else 0 in
          let res_evaluation =
            let evaluation_linear_i =
              Evaluations.linear ~res:tmp_evaluation ~evaluations
                ~poly_names:[ wire_name; s_name ] ~linear_coeffs:[ one; beta ]
                ~add_constant:gamma ()
            in
            Evaluations.mul_c ~res:res_evaluation
              ~evaluations:[ evaluation_linear_i; acc_evaluation ]
              ~composition_gx:([ 0; comp ], n)
              ()
          in
          (i + 1, res_evaluation)
        in
        List.fold_left2 f_fold (0, z_evaluation) wires_names s_names
      in
      res_evaluation

    (* evaluations must contain z’s evaluation *)
    let precompute_perm_identity_poly wires_name beta gamma evaluations n =
      let z_evaluation = Evaluations.find_evaluation evaluations "Z" in
      let z_evaluation_len = Evaluations.length z_evaluation in
      let tmp_evaluation = Evaluations.find_evaluation evaluations "tmp_eval" in
      let id1_evaluation = Evaluations.create z_evaluation_len in
      let id2_evaluation = Evaluations.create z_evaluation_len in

      let identity_zfg =
        let nb_wires = Array.length wires_name 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 tmp_evaluation id2_evaluation beta gamma evaluations
            wires_name sid_names "Z" n
        in
        let g_evaluation =
          let ss_names =
            List.init nb_wires (fun i -> "Ss" ^ string_of_int (i + 1))
          in
          compute_prime id2_evaluation id1_evaluation beta gamma evaluations
            wires_name ss_names "Zg" 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 [ ("Perm.a", identity_l1_z); ("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 blinds =
      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 = Bls12_381.Fr.(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.(values.(indices_list_j_i) + gamma) in
            let f_coeff =
              let gi = Domain.get domain i in
              Bls12_381.Fr.(
                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
              Bls12_381.Fr.(
                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
              Bls12_381.Fr.(
                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
          Bls12_381.Fr.(
            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

      let z_poly = Evaluations.interpolation_fft2 domain scalar_array_Z in
      match blinds with
      | None -> z_poly
      | Some b ->
          let p_poly =
            Poly.of_coefficients [ (b.(0), 2); (b.(1), 1); (b.(2), 0) ]
          in
          let zs_poly =
            Poly.of_coefficients [ (one, size_domain); (mone, 0) ]
          in
          Poly.(add (mul zs_poly p_poly) z_poly)
  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 ?(prefix = "") ~domain ~permutation ~nb_wires () =
    Preprocessing.ssigma_map_non_quadratic_residues ~prefix permutation domain
      nb_wires

  let sid_not_committed = Preprocessing.sid_map_non_quadratic_residues_verifier

  let common_preprocessing ~compute_l1 ~domain ~nb_wires ~domain_evals =
    let sid_evals = Preprocessing.evaluations_sid nb_wires domain_evals in
    let sid_query =
      let sid_func =
        Preprocessing.sid_map_non_quadratic_residues_verifier nb_wires
      in
      PP.{ empty_verifier_query with not_committed = sid_func }
    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, sid_query)

  let prover_query ?(prefix = "") ~wires_name ~generator ~beta ~gamma
      ~evaluations ~n () =
    PP.
      {
        v_map = v_map ~prefix generator;
        precomputed_polys =
          Permutation_poly.precompute_perm_identity_poly wires_name beta gamma
            evaluations n;
      }

  let verifier_query ?(compute_sid = true) ?(prefix = "") ~wires_name ~generator
      ~beta ~gamma ~nb_wires () =
    PP.
      {
        v_map = v_map ~prefix generator;
        identities = gate_identity ~prefix wires_name beta gamma;
        not_committed =
          (if compute_sid then
           Preprocessing.sid_map_non_quadratic_residues_verifier nb_wires
          else SMap.empty);
      }

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

module type Permutation_gate_sig = sig
  module PP : Polynomial_protocol.Polynomial_protocol_sig

  val srs_size : zero_knowledge:bool -> n:int -> int
  val polynomials_degree : nb_wires:int -> int

  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
    * PP.verifier_query

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

  val prover_query :
    ?prefix:string ->
    wires_name:string array ->
    generator:PP.PC.Scalar.t ->
    beta:PP.PC.Scalar.t ->
    gamma:PP.PC.Scalar.t ->
    evaluations:PP.Evaluations.t SMap.t ->
    n:int ->
    unit ->
    PP.prover_query

  val verifier_query :
    ?compute_sid:bool ->
    ?prefix:string ->
    wires_name:string array ->
    generator:PP.PC.Scalar.t ->
    beta:PP.PC.Scalar.t ->
    gamma:PP.PC.Scalar.t ->
    nb_wires:int ->
    unit ->
    PP.verifier_query

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

module Permutation_gate (PP : Polynomial_protocol.Polynomial_protocol_sig) :
  Permutation_gate_sig with module PP = PP =
  Permutation_gate_impl (PP)