Source file polynomial_commitment.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 Plonk

module Make_impl
    (Pack : Pack.Aggregator)
    (PC : Polynomial_commitment.S
            with type Scalar.t = Pack.scalar
             and type Commitment.t = Pack.g1 SMap.t) =
struct
  module Scalar = PC.Scalar
  module Fr_generation = PC.Fr_generation
  module Polynomial = PC.Polynomial
  module Domain = PC.Polynomial.Domain
  module Poly = PC.Polynomial.Polynomial
  module Scalar_map = PC.Scalar_map

  type secret = PC.secret
  type query = PC.query [@@deriving repr]
  type answer = PC.answer [@@deriving repr]
  type transcript = PC.transcript

  module Public_parameters = struct
    type prover = {
      pp_pc_prover : PC.Public_parameters.prover;
      pp_pack_prover : Pack.prover_public_parameters;
    }
    [@@deriving repr]

    type verifier = {
      pp_pc_verifier : PC.Public_parameters.verifier;
      pp_pack_verifier : Pack.verifier_public_parameters;
    }
    [@@deriving repr]

    type setup_params = int

    let setup setup_params srs =
      let pp_pc_prover, pp_pc_verifier =
        PC.Public_parameters.setup setup_params srs
      in
      let pp_pack_prover, pp_pack_verifier =
        Pack.setup setup_params (snd srs)
      in
      let pp_prover = { pp_pc_prover; pp_pack_prover } in
      let pp_verifier = { pp_pc_verifier; pp_pack_verifier } in
      (pp_prover, pp_verifier)

    let to_bytes d ({ pp_pc_prover; pp_pack_prover } : prover) =
      Utils.Hash.hash_bytes
        [
          PC.Public_parameters.to_bytes d pp_pc_prover;
          Pack.public_parameters_to_bytes pp_pack_prover;
        ]
  end

  module Commitment = struct
    type t = Pack.commitment [@@deriving repr]

    (* [PC.Commitment.t] is required to be [Bls12_381.G1.t SMap.t],
       containing all the commitments that were packed *)
    type prover_aux = PC.Commitment.t * PC.Commitment.prover_aux
    [@@deriving repr]

    let commit ?all_keys (pp : Public_parameters.prover) f_map =
      let relevant_positions =
        match all_keys with
        | None -> List.init (SMap.cardinal f_map) Fun.id
        | Some ks ->
            List.mapi (fun i x -> (i, x)) ks
            |> List.filter_map (fun (i, x) ->
                   Option.map (Fun.const i) (SMap.find_opt x f_map))
      in
      let prover_aux = PC.Commitment.commit pp.pp_pc_prover f_map in
      let cm_list = List.map snd @@ SMap.bindings (fst prover_aux) in
      let pack_cmt =
        Pack.partial_commit ~relevant_positions pp.pp_pack_prover
          (Array.of_list cm_list)
      in
      (pack_cmt, prover_aux)

    let cardinal = Pack.commitment_cardinal
  end

  type proof = {
    pc_proof : PC.proof;
    packed_values : Pack.packed list;
    pack_proof : Pack.proof;
  }
  [@@deriving repr]

  type prover_aux = {
    r : Scalar.t;
    s_list : Scalar.t SMap.t list;
    cm_answers : Scalar.t;
  }

  let expand_with_proof = Utils.expand_transcript proof_t
  let expand_with_query = Utils.list_expand_transcript query_t
  let expand_with_answer = Utils.list_expand_transcript answer_t

  let batch_p ~zero ~add ~mul r map =
    SMap.fold
      (fun _ x (acc, rk) -> (add acc (mul rk x), Scalar.mul r rk))
      map (zero, Scalar.one)
    |> fst

  let batch ~zero ~add ~mul r map =
    SMap.fold
      (fun _ x (acc, rk) -> (add acc (mul rk x), Scalar.mul r rk))
      map (zero, Scalar.one)
    |> fst

  let batch_polys r map =
    let polys = SMap.bindings map |> List.map snd in
    Poly.linear_with_powers polys r

  let batch_answers r = SMap.map Scalar.(batch ~zero ~add ~mul r)
  let evaluate = PC.evaluate

  let poseidon array =
    let module Poseidon = Mec.Hash.Poseidon128.Make (Scalar) in
    Poseidon.Hash.(get (digest (init ()) array))

  let poseidon_answers answers =
    answers
    |> List.map (fun a ->
           SMap.bindings a
           |> List.map (fun (_, m) -> SMap.bindings m |> List.map snd)
           |> List.flatten)
    |> List.flatten |> Array.of_list |> poseidon

  (* compute P := cmt₀ + r cmt₁ + r² cmt₂ + ... for every group of commitments
     in the list [prover_aux_list], and common randomness r (freshly sampled);
     such P values are returned as [packed_values], together with a proof
     [packed_proof] of their correctness;
     also, on input a list of evaluations [answer_list], at the requested points
     in [query_list], produce a proof of their validity: such proof is a PC
     proof (for every group) on the aggregatted commitment P with respect to the
     corresponding aggregated evaluations (we thus batch [answer_list] with [r]
     similarly) *)
  let prove_pack (pp : Public_parameters.prover) transcript f_map_list
      (prover_aux_list : Commitment.prover_aux list) query_list answer_list =
    let r, transcript = Fr_generation.random_fr transcript in
    let f_list = List.map (batch_polys r) f_map_list in
    let s_list = List.map (batch_answers r) answer_list in
    (* [cmts_list] is a list of G1.t SMap.t, containing the PC commitments to
       every polynomial (note that PC.Commitment.t = Bls12_381.G1.t SMap.t) *)
    let cmts_list =
      List.map
        (fun (cmts, _prover_aux) ->
          List.map snd @@ SMap.bindings cmts |> Array.of_list)
        prover_aux_list
    in
    (* [packed_values] has type [G1.t list] and it is the result of batching
       each map in [cmt_list] with powers of [r].
       [pack_proof] asserts that [packed_values] was correctly computed. *)
    let (packed_values, pack_proof), transcript =
      Pack.prove pp.pp_pack_prover transcript r cmts_list
    in
    (* prepare [f_list] and [s_list], the batched version of [f_map_list]
       polys and [answer_list] (using randomness [r]) by selecting a dummy
       name for them [string_of_int i] in order to call the underlying PC *)
    let f_map_list =
      List.mapi (fun i l -> SMap.singleton (string_of_int i) l) f_list
    in
    let s_map_list =
      List.mapi
        (fun i m -> SMap.map (fun s -> SMap.singleton (string_of_int i) s) m)
        s_list
    in
    let prover_aux_list = List.map snd prover_aux_list in
    (* call the underlying PC prover on the batched polynomials/evaluations
       the verifier will verify such proof using [packed_values] as the
       commitments *)
    let pc_proof, transcript =
      PC.prove pp.pp_pc_prover transcript f_map_list prover_aux_list query_list
        s_map_list
    in
    let proof = { pc_proof; packed_values; pack_proof } in
    let transcript = expand_with_proof proof transcript in
    (proof, transcript, r, s_list)

  let prove (pp : Public_parameters.prover) transcript f_map_list
      (prover_aux_list : Commitment.prover_aux list) query_list answer_list =
    let transcript = expand_with_query query_list transcript in
    let transcript = expand_with_answer answer_list transcript in
    let proof, transcript, _, _ =
      prove_pack pp transcript f_map_list prover_aux_list query_list answer_list
    in
    (proof, transcript)

  let prove_super_aggregation (pp : Public_parameters.prover) transcript
      f_map_list (prover_aux_list : Commitment.prover_aux list) query_list
      answer_list =
    let transcript = expand_with_query query_list transcript in
    let cm_answers = poseidon_answers answer_list in
    let transcript = Utils.expand_transcript Scalar.t cm_answers transcript in
    let proof, transcript, r, s_list =
      prove_pack pp transcript f_map_list prover_aux_list query_list answer_list
    in
    ((proof, { r; s_list; cm_answers }), transcript)

  let verify_pack (pp : Public_parameters.verifier) r transcript cmt_list
      query_list s_list proof =
    (* verify that the [packed_values] are correct, they will be used as
       the commitments for the PC proof of (batched) evaluations *)
    let pack_ok, transcript =
      Pack.verify pp.pp_pack_verifier transcript cmt_list r
        (proof.packed_values, proof.pack_proof)
    in
    (* batch the evaluations using [r] and prepare the query to the PC verifier
       by selecting the default dummy names [string_of_int i] names *)
    let s_map_list =
      List.mapi
        (fun i m -> SMap.map (fun s -> SMap.singleton (string_of_int i) s) m)
        s_list
    in
    let cmt_map_list =
      List.mapi
        (fun i l -> SMap.singleton (string_of_int i) l)
        proof.packed_values
    in
    (* verify that the batched evaluations are correct *)
    let pc_ok, transcript =
      PC.verify pp.pp_pc_verifier transcript cmt_map_list query_list s_map_list
        proof.pc_proof
    in
    (pack_ok && pc_ok, expand_with_proof proof transcript)

  let verify (pp : Public_parameters.verifier) transcript cmt_list query_list
      s_map_list proof =
    let transcript = expand_with_query query_list transcript in
    let transcript = expand_with_answer s_map_list transcript in
    let r, transcript = Fr_generation.random_fr transcript in
    let s_list = List.map (batch_answers r) s_map_list in
    verify_pack pp r transcript cmt_list query_list s_list proof

  let verify_super_aggregation (pp : Public_parameters.verifier) transcript
      cmt_list query_list s_list cm_answers proof =
    let transcript = expand_with_query query_list transcript in
    let transcript = Utils.expand_transcript Scalar.t cm_answers transcript in
    let r, transcript = Fr_generation.random_fr transcript in
    let ok, transcript =
      verify_pack pp r transcript cmt_list query_list s_list proof
    in
    (ok, r, transcript)
end

module type S = sig
  include Polynomial_commitment.S

  type prover_aux = {
    r : Scalar.t;
    s_list : Scalar.t SMap.t list;
    cm_answers : Scalar.t;
  }

  val poseidon : Scalar.t array -> Scalar.t

  val prove_super_aggregation :
    Public_parameters.prover ->
    transcript ->
    Polynomial.Polynomial.t SMap.t list ->
    Commitment.prover_aux list ->
    query list ->
    Scalar.t SMap.t SMap.t list ->
    (proof * prover_aux) * transcript

  val verify_super_aggregation :
    Public_parameters.verifier ->
    transcript ->
    Commitment.t list ->
    query list ->
    Scalar.t SMap.t list ->
    Scalar.t ->
    proof ->
    bool * Scalar.t * transcript
end

module Make : functor
  (Pack : Pack.Aggregator)
  (PC : Polynomial_commitment.S
          with type Scalar.t = Pack.scalar
           and type Commitment.t = Pack.g1 SMap.t)
  -> S with module Scalar = PC.Scalar =
  Make_impl

include Make (Pack) (Polynomial_commitment.Kzg_impl)