Source file kzg_pack.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 SMap = Plonk.SMap

(** Extension of the KZG_pack implementation with additional
    types and functions used in by Distributed_prover  *)
module Make
    (Pack : Aggregation.Pack.Aggregator)
    (PC : Kzg.PC_for_distribution_sig
            with type BasePC.Scalar.t = Pack.scalar
             and type Commitment.t = Pack.g1 SMap.t) =
struct
  module BasePC = Aggregation.Polynomial_commitment.Make_impl (Pack) (PC)

  module Commitment = struct
    include BasePC.Commitment

    let recombine l = List.fold_left Pack.combine (List.hd l) (List.tl l)

    let recombine_prover_aux l =
      let cm = PC.Commitment.recombine (List.map fst l) in
      let p_a = PC.Commitment.recombine_prover_aux (List.map snd l) in
      (cm, p_a)

    let empty = Pack.empty_commitment
    let empty_prover_aux = (PC.Commitment.empty, PC.Commitment.empty_prover_aux)
  end

  include (BasePC : module type of BasePC with module Commitment := Commitment)

  type worker_msg = Scalar.t * string list list [@@deriving repr]

  type main_prover_msg = Poly.t list * Commitment.prover_aux list
  [@@deriving repr]

  type main_prover_state =
    Public_parameters.prover
    * transcript
    * Scalar.t
    * query list
    * Scalar.t SMap.t list
    * main_prover_msg

  let merge_answers : answer list -> answer =
    let open SMap in
    List.fold_left (union (fun _k m1 m2 -> Some (union_disjoint m1 m2))) empty

  let distributed_prove_worker f_map_list prover_aux_list (r, poly_keys_list) =
    let gen_powers r l =
      List.mapi (fun i x -> (x, Scalar.pow r @@ Z.of_int i)) l |> SMap.of_list
    in
    let r_powers_list = List.map (gen_powers r) poly_keys_list in
    let f_list =
      List.map2
        (fun f_map r_map ->
          let polys = SMap.bindings f_map in
          let coeffs = List.map (fun (name, _) -> SMap.find name r_map) polys in
          Poly.linear (List.map snd polys) coeffs)
        f_map_list r_powers_list
    in
    (f_list, prover_aux_list)

  let distributed_prove_main1 (pp : Public_parameters.prover) transcript
      query_list answer_list secret_list prover_aux_list =
    let transcript = expand_with_query query_list transcript in
    let transcript = expand_with_answer answer_list transcript in
    let r, transcript = Fr_generation.random_fr transcript in
    let s_list = List.map (batch_answers r) answer_list in

    let get_keys map_map =
      SMap.fold
        (fun _ m acc -> SMap.union (fun _ _ x -> Some x) m acc)
        map_map SMap.empty
      |> SMap.bindings |> List.map fst
    in
    let poly_keys_list = List.map get_keys answer_list in
    let worker_message = (r, poly_keys_list) in
    (* The main thread simulates a worker, since it is the only one who knows
         the information about t_map, g_map, plook_map. We need to pad a dummy
         secret and a dummy prover_aux at the end, corresponding to f_map, which
         the main thread does not have information about. *)
    let main_msg =
      distributed_prove_worker secret_list prover_aux_list worker_message
    in
    let state = (pp, transcript, r, query_list, s_list, main_msg) in
    (worker_message, state)

  let distributed_prove_main2
      ( (pp : Public_parameters.prover),
        transcript,
        r,
        query_list,
        s_list,
        main_msg ) worker_msg_list =
    let worker_msg_list = main_msg :: worker_msg_list in
    let f_list_list = List.map fst worker_msg_list in
    let f_list =
      Plonk.List.mapn (List.fold_left Poly.add Poly.zero) f_list_list
    in
    let prover_aux_list_list = List.map snd worker_msg_list in
    let prover_aux_list =
      Plonk.List.mapn Commitment.recombine_prover_aux prover_aux_list_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)
end

module Kzg_pack_impl = Make (Aggregation.Pack) (Kzg.Kzg_impl)