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

(* Implements a batched version of the KZG10 scheme, described in Section 3 of
   the PlonK paper: https://eprint.iacr.org/2019/953.pdf *)
module Kzg_impl = struct
  module Scalar = Bls12_381.Fr
  module G1 = Bls12_381.G1
  module G2 = Bls12_381.G2
  module GT = Bls12_381.GT
  module Pairing = Bls12_381.Pairing
  module Fr_generation = Fr_generation.Make (Scalar)
  module Polynomial = Polynomial
  module Poly = Polynomial.Polynomial
  module Scalar_map = Map.Make (Scalar)

  (* polynomials to be committed *)
  type secret = Poly.t SMap.t

  (* maps evaluation point names to evaluation point values *)
  type query = Scalar.t SMap.t

  (* maps evaluation point names to (map from polynomial names to evaluations) *)
  type answer = Scalar.t SMap.t SMap.t

  let answer_encoding : answer Data_encoding.t =
    SMap.encoding @@ SMap.encoding Encodings.fr_encoding

  let bytes_of_query =
    Data_encoding.Binary.to_bytes_exn @@ SMap.encoding Encodings.fr_encoding

  let bytes_of_answer = Data_encoding.Binary.to_bytes_exn answer_encoding

  type transcript = Bytes.t

  let pippenger ?(start = 0) ?len ps ss =
    try G1.pippenger ~start ?len ps ss
    with Invalid_argument s ->
      raise (Invalid_argument (Printf.sprintf "KZG.pippenger : %s" s))

  module Public_parameters = struct
    (* Structured Reference String
       - srs1 : [[1]₁, [x¹]₁, …, [x^(d-1)]₁] ;
       - encoding_1 : [1]₂;
       - encoding_x : [x]₂;
       - d : srs length (integer) *)
    type prover = {
      srs1 : Polynomial.Srs.t;
      encoding_1 : G2.t;
      encoding_x : G2.t;
      d : int;
    }

    let to_bytes srs =
      let open Utils.Hash in
      let st = init () in
      update st (G2.to_bytes srs.encoding_1);
      update st (G2.to_bytes srs.encoding_x);
      let srs1 = Polynomial.Srs.to_array srs.srs1 in
      Array.iter (fun key -> update st (G1.to_bytes key)) srs1;
      finish st

    type verifier = { encoding_1 : G2.t; encoding_x : G2.t }

    let verifier_encoding : verifier Data_encoding.t =
      let open Encodings in
      let open Data_encoding in
      conv
        (fun { encoding_1; encoding_x } -> (encoding_1, encoding_x))
        (fun (encoding_1, encoding_x) -> { encoding_1; encoding_x })
        (obj2 (req "encoding_1" g2_encoding) (req "encoding_x" g2_encoding))

    type setup_params = int * int

    let get_d srs = srs.d

    let create_srs1 d x =
      let xi = ref G1.one in
      Array.init d (fun _ ->
          let res = !xi in
          xi := G1.mul !xi x;
          res)

    let encoding x = G1.mul G1.one x
    let encoding_2 x = G2.mul G2.one x

    let setup ?state (d, _) =
      let x = Scalar.random ?state () in
      let srs1 = create_srs1 d x in
      let pippenger_ctxt = Polynomial.Srs.of_array srs1 in
      let encoding_1 = G2.one in
      let encoding_x = G2.(mul one x) in
      ( { srs1 = pippenger_ctxt; encoding_1; encoding_x; d },
        { encoding_1; encoding_x } )

    let prv_import d (srsfile_g1, srsfile_g2) =
      let srs1 = Polynomial.Srs.load_from_file srsfile_g1 ~offset:0 d in
      let srs2 = Utils.import_srs_g2 2 srsfile_g2 in
      let encoding_1, encoding_x = (srs2.(0), srs2.(1)) in
      { srs1; encoding_1; encoding_x; d }

    let vrf_import =
      let encoding_1 =
        Bytes.of_string
          "\019\224+`Rq\159`}\172\211\160\136\'OeYk\208\208\153 \
           \182\026\181\218a\187\220\127PI3L\241\018\019\148]W\229\172}\005]\004+~\002J\162\178\240\143\n\
           \145&\b\005\'-\197\016Q\198\228z\212\250@;\
           \002\180Q\011dz\227\209w\011\172\003&\168\005\187\239\212\128V\200\193!\189\184\006\006\196\160.\1674\2042\172\210\176+\194\139\153\203>(~\133\167c\175&t\146\171W.\153\171?7\r\'\\\236\029\161\170\169\007_\240_y\190\012\229\213\'r}n\017\140\201\205\198\218.5\026\173\253\155\170\140\189\211\167mB\154iQ`\209,\146:\201\204;\172\162\137\225\147T\134\b\184(\001"
        |> G2.of_bytes_exn
      in
      let encoding_x =
        Bytes.of_string
          "\017k\"P\252K0\152\204\173\249\236XRj\236\243\007\210\030#?e\017\005T\238\004F\015r\166\022\244\208p\148\027\007\150\222]\239\218)\025n@\0016E\184Q\142Uh\231]2\r\168\184:\135\225$\006\149k\156tI\030\229\236\240m\194\151\186\017\134\001b\t\215\143\2162\193n\246/\196\158\255\011\2002\251\135\140/f\173\221\132j\168YT\184V\195\188p\216sD\249\000h\137\229\142\176\236W4\021\031\153\185\165\182S\233\145$z\213\006\208\251\018\219>.\131\024\182GS\127\000\230\183\002\252\133&\232\201\016\026{|\188A\015\191\143`!\242\228S\147\200id1\193\166\252\199\245\167\217\n\
           \185\153"
        |> G2.of_bytes_exn
      in
      { encoding_1; encoding_x }

    let import (d, _) (srsfiles, _) =
      let prv = prv_import d srsfiles in
      let vrf = vrf_import in
      (prv, vrf)
  end

  module Commitment = struct
    type t = G1.t SMap.t

    let encoding : t Data_encoding.t = SMap.encoding Encodings.g1_encoding

    type prover_aux = unit

    let expand_transcript = SMap.expand_transcript ~to_bytes:G1.to_bytes

    let commit_kate_amortized pippenger_ctxt p =
      let poly_size = Poly.degree p + 1 in
      let srs_size = Polynomial.Srs.size pippenger_ctxt in
      if poly_size = 0 then G1.zero
      else if poly_size > srs_size then
        raise
          (Failure
             (Printf.sprintf
                "Kzg.compute_encoded_polynomial : Polynomial degree, %i, \
                 exceeds srs’ length, %i."
                poly_size srs_size))
      else Polynomial.Srs.pippenger pippenger_ctxt p

    let commit_single srs p =
      commit_kate_amortized Public_parameters.(srs.srs1) p

    let commit srs f_map =
      let cmt = SMap.map (commit_single srs) f_map in
      let prover_aux = () in
      (cmt, prover_aux)

    let cardinal cmt = SMap.cardinal cmt
  end

  type proof = G1.t SMap.t

  let proof_encoding : proof Data_encoding.t =
    SMap.encoding Encodings.g1_encoding

  let expand_with_proof = SMap.expand_transcript ~to_bytes:G1.to_bytes
  let expand_with_query = Utils.list_expand_transcript ~to_bytes:bytes_of_query

  let expand_with_answer =
    Utils.list_expand_transcript ~to_bytes:bytes_of_answer

  (* compute W := (f(x) - s) / (x - z), where x is the srs secret exponent,
     for every evaluation point [zname], key of the [query] map, where
       z := SMap.find zname query
       s := SMap.find zname batched_answer
       f := SMap.find zname batched_polys
     the computation is performed by first calculating polynomial
     (f(X) - s) / (X - z) and then committing to it using the srs.
     Here, f (respecitvely s) is a batched polynomial (respecively batched
     evaluation) of all polynomials (and their respective evaluations) that
     are evaluated at a common point z. They have been batched with the
     uniformly sampled randomness from [y_map], see {!sample_ymap} *)
  let compute_Ws srs batched_polys batched_answer query =
    SMap.mapi
      (fun x z ->
        let f = SMap.find x batched_polys in
        let s = SMap.find x batched_answer in
        (* WARNING: This modifies [batched_polys], but we won't use it again: *)
        Poly.sub_inplace f f @@ Poly.constant s;
        let h = Poly.division_x_z f z in
        Commitment.commit_single srs h)
      query

  (* verify the KZG equation: e(F - [s]₁ + z W, [1]₂) = e(W, [x]₂)
     for every evaluation point [zname], key of the [query] map, where
       z := SMap.find zname query
       s := SMap.find zname s_map
       W := SMap.find zname w_map
     and F is computed as a linear combination (determined by [coeffs])
     of the commitments in [SMap.find zname cmt_map].
     All verification equations are checked at once by batching them
     with fresh randomness sampled in [r_map].
     The combination of [cmt_map] and other G1.mul is delayed as much
     as possible, in order to combine all of them with a single pippenger *)
  let verifier_check srs cmt_map coeffs query s_map w_map =
    let r_map = SMap.map (fun _ -> Scalar.random ()) w_map in
    let cmts = SMap.bindings cmt_map |> List.map snd in
    let exponents =
      SMap.fold
        (fun x r exponents ->
          let x_coeffs = SMap.find x coeffs in
          SMap.mapi
            (fun name exp ->
              match SMap.find_opt name x_coeffs with
              | None -> exp
              | Some c -> Scalar.(exp + (r * c)))
            exponents)
        r_map
        (SMap.map (fun _ -> Scalar.zero) cmt_map)
      |> SMap.bindings |> List.map snd
    in
    let s =
      SMap.fold
        (fun x r s -> Scalar.(sub s (r * SMap.find x s_map)))
        r_map Scalar.zero
    in
    let w_left_exps =
      List.map (fun (x, r) -> Scalar.mul r @@ SMap.find x query)
      @@ SMap.bindings r_map
    in
    let w_right_exps =
      (* We negate them before the pairing_check, which is done on the lhs *)
      SMap.bindings r_map |> List.map snd |> List.map Scalar.negate
    in

    let ws = SMap.bindings w_map |> List.map snd in
    let left =
      pippenger
        (Array.of_list @@ (G1.one :: ws) @ cmts)
        (Array.of_list @@ (s :: w_left_exps) @ exponents)
    in
    let right = pippenger (Array.of_list ws) (Array.of_list w_right_exps) in
    Public_parameters.[ (left, srs.encoding_1); (right, srs.encoding_x) ]
    |> Pairing.pairing_check

  (* return a map between evaluation point names (from [query]) and uniformly
     sampled scalars, used for batching; also return an updated transcript *)
  let sample_ys transcript query =
    let n = SMap.cardinal query in
    let ys, transcript = Fr_generation.random_fr_list transcript n in
    let y_map =
      SMap.of_list
        (List.map2 (fun y (name, _) -> (name, y)) ys @@ SMap.bindings query)
    in
    (y_map, transcript)

  (* On input a scalar map [y_map] and [answer], e.g.,
      y_map := { 'x0' -> y₀; 'x1' -> y₁ }
     answer := { 'x0' -> { 'a' -> a(x₀); 'b' -> b(x₀); 'c' -> c(x₀); ... };
                 'x1' -> { 'a' -> a(x₁); 'c' -> c(x₁); 'd' -> d(x₁); ... }; }
     outputs a map of batched evaluations:
       { 'x0' -> a(x₀) + y₀b(x0) + y₀²c(x₀) + ...);
         'x1' -> a(x₁) + y₁c(x1) + y₁²d(x₁) + ...); }
     and a map of batching coefficients:
       { 'x0' -> { 'a' -> 1; 'b' -> y₀; 'c' -> y₀²; ... };
         'x1' -> { 'a' -> 1; 'c' -> y₁; 'd' -> y₁²; ... }; } *)
  let batch_answer y_map answer =
    let couples =
      SMap.mapi
        (fun x s_map ->
          let y = SMap.find x y_map in
          let s, coeffs, _yk =
            SMap.fold
              (fun name s (acc_s, coeffs, yk) ->
                let acc_s = Scalar.(add acc_s @@ mul yk s) in
                let coeffs = SMap.add name yk coeffs in
                let yk = Scalar.mul yk y in
                (acc_s, coeffs, yk))
              s_map
              (Scalar.zero, SMap.empty, Scalar.one)
          in
          (s, coeffs))
        answer
    in
    (SMap.map fst couples, SMap.map snd couples)

  (* On input batching coefficients [coeffs] and a map of polys [f_map], e.g.,
      coeffs := { 'x0' -> { 'a' -> 1; 'b' -> y₀; 'c' -> y₀²; ... };
                  'x1' -> { 'a' -> 1; 'c' -> y₁; 'd' -> y₁²; ... }; }
       f_map := { 'a' -> a(X); 'b' -> b(X); 'c' -> c(X); ... },
     outputs a map of batched polynomials:
       { 'x0' -> a(X) + y₀b(X) + y₀²c(X) + ...);
         'x1' -> a(X) + y₁c(X) + y₁²d(X) + ...); } *)
  let batch_polys coeffs f_map =
    SMap.map
      (fun f_coeffs ->
        let f_coeffs = SMap.bindings f_coeffs in
        let polys =
          List.map (fun name -> SMap.find name f_map) @@ List.map fst f_coeffs
        in
        Poly.linear polys @@ List.map snd f_coeffs)
      coeffs

  let prove_single srs transcript f_map query answer =
    let y_map, transcript = sample_ys transcript query in
    let batched_answer, coeffs = batch_answer y_map answer in
    let batched_polys = batch_polys coeffs f_map in
    let proof = compute_Ws srs batched_polys batched_answer query in
    (proof, expand_with_proof transcript proof)

  let verify_single srs transcript cmt_map query answer proof =
    let y_map, transcript = sample_ys transcript query in
    let batched_answer, coeffs = batch_answer y_map answer in
    let b = verifier_check srs cmt_map coeffs query batched_answer proof in
    (b, expand_with_proof transcript proof)

  (* group functions allow [prove] and [verify] rely on [prove_single] and
     [verify_single] respectively *)

  let group_secrets : secret list -> secret = SMap.union_disjoint_list
  let group_cmts : Commitment.t list -> Commitment.t = SMap.union_disjoint_list

  let group_queries : query list -> query =
   fun query_list ->
    let union =
      SMap.union (fun _ z z' ->
          if Scalar.eq z z' then Some z
          else
            failwith "group_query: equal query names must map to equal values")
    in
    List.fold_left union (List.hd query_list) (List.tl query_list)

  let group_answers : answer list -> answer =
   fun answer_list ->
    List.fold_left
      (SMap.union (fun _ m1 m2 -> Some (SMap.union_disjoint m1 m2)))
      (List.hd answer_list) (List.tl answer_list)

  (* evaluate every polynomial in [f_map] at all evaluation points in [query] *)
  let evaluate : Poly.t SMap.t -> query -> answer =
   fun f_map query ->
    SMap.map (fun z -> SMap.map (fun f -> Poly.evaluate f z) f_map) query

  let prove srs transcript f_map_list _prover_aux_list query_list answer_list =
    let transcript = expand_with_query transcript query_list in
    let transcript = expand_with_answer transcript answer_list in
    let f_map = group_secrets f_map_list in
    let query = group_queries query_list in
    let answer = group_answers answer_list in
    prove_single srs transcript f_map query answer

  let verify srs transcript cmt_map_list query_list answer_list proof =
    let transcript = expand_with_query transcript query_list in
    let transcript = expand_with_answer transcript answer_list in
    let cmt_map = group_cmts cmt_map_list in
    let query = group_queries query_list in
    let answer = group_answers answer_list in
    verify_single srs transcript cmt_map query answer proof
end

module type Polynomial_commitment_sig = sig
  module Scalar : Ff_sig.PRIME with type t = Bls12_381.Fr.t
  module Polynomial : Polynomial.S with type scalar = Scalar.t
  module Scalar_map : Map.S with type key = Scalar.t

  module Fr_generation :
    Fr_generation.Fr_generation_sig with type scalar = Scalar.t

  (* polynomials to be committed *)
  type secret = Polynomial.Polynomial.t SMap.t

  (* maps evaluation point names to evaluation point values *)
  type query = Scalar.t SMap.t

  (* maps evaluation point names to (map from polynomial names to evaluations) *)
  type answer = Scalar.t SMap.t SMap.t

  val answer_encoding : answer Data_encoding.t

  type proof

  val proof_encoding : proof Data_encoding.t

  type transcript = Bytes.t

  module Public_parameters : sig
    type prover
    type verifier

    val verifier_encoding : verifier Data_encoding.t

    type setup_params = int * int

    (* [setup ~state d] returns an SRS of size d generated randomly *)
    val setup : ?state:Random.State.t -> setup_params -> prover * verifier

    (* [get_d srs] returns the size of srs *)
    val get_d : prover -> int

    (* [import d srsfile] returns the SRS of size d imported from srsfile *)
    val import :
      setup_params -> (string * string) * (string * string) -> prover * verifier

    val to_bytes : prover -> Bytes.t
  end

  module Commitment : sig
    type t
    type prover_aux

    val encoding : t Data_encoding.t
    val expand_transcript : transcript -> t -> transcript
    val commit : Public_parameters.prover -> secret -> t * prover_aux
    val cardinal : t -> int
  end

  val evaluate : secret -> query -> answer

  val prove :
    Public_parameters.prover ->
    transcript ->
    secret list ->
    Commitment.prover_aux list ->
    query list ->
    answer list ->
    proof * transcript

  val verify :
    Public_parameters.verifier ->
    transcript ->
    Commitment.t list ->
    query list ->
    answer list ->
    proof ->
    bool * transcript
end

include (
  Kzg_impl :
    Polynomial_commitment_sig with type Commitment.t = Bls12_381.G1.t SMap.t)