Source file main_protocol.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

module type S = sig
  module PP : Polynomial_protocol.S

  type prover_common_pp [@@deriving repr]
  type prover_circuit_pp [@@deriving repr]

  type prover_public_parameters = {
    common_pp : prover_common_pp;
    circuits_map : prover_circuit_pp SMap.t;
    transcript : PP.transcript;
  }
  [@@deriving repr]

  include
    Plonk.Main_protocol.S
      with type prover_public_parameters := prover_public_parameters

  type worker_inputs [@@deriving repr]

  val split_inputs_map :
    nb_workers:int ->
    circuit_prover_input list SMap.t ->
    worker_inputs SMap.t list

  type commit_to_wires_reply = PP.PC.Commitment.t [@@deriving repr]
  type wires_info

  type commit_to_wires_remember = {
    f_wires_map : PP.Evaluations.polynomial SMap.t;
    wires_map : wires_info list SMap.t;
    inputs_map : circuit_prover_input list SMap.t;
    shifts_map : (int * int) SMap.t;
    srs : PP.prover_public_parameters;
    f_wires_map_list_map : PP.Evaluations.polynomial SMap.t list SMap.t;
    wires_prover_aux : PP.PC.Commitment.prover_aux;
  }

  val worker_commit_to_wires :
    ?zero_knowledge:bool ->
    prover_common_pp * prover_circuit_pp SMap.t ->
    worker_inputs SMap.t ->
    commit_to_wires_reply * commit_to_wires_remember

  type commit_to_plook_reply = {
    batched_witness_map : PP.Evaluations.t SMap.t;
    cmt_plookup : PP.PC.Commitment.t;
    f_plook_map : PP.PC.Polynomial.Polynomial.t SMap.t;
    plook_prover_aux : PP.PC.Commitment.prover_aux;
  }
  [@@deriving repr]

  type commit_to_plook_remember = {
    f_wires_map_list_map : PP.Evaluations.polynomial SMap.t list SMap.t;
    f_plook_map_list_map : PP.Evaluations.polynomial SMap.t list SMap.t;
    beta_plookup : scalar;
    gamma_plookup : scalar;
  }

  val commit_to_plook :
    ?zero_knowledge:bool ->
    prover_common_pp * prover_circuit_pp SMap.t ->
    wires_info list SMap.t ->
    (int * int) SMap.t ->
    PP.prover_public_parameters ->
    PP.transcript ->
    PP.Evaluations.polynomial SMap.t list SMap.t ->
    commit_to_plook_reply * commit_to_plook_remember

  val batch_evaluated_ids :
    alpha:scalar -> PP.Evaluations.t SMap.t -> string list -> PP.Evaluations.t

  val worker_build_evaluations :
    prover_common_pp * prover_circuit_pp SMap.t ->
    f_wires_map:PP.Evaluations.polynomial SMap.t ->
    f_plook_map:PP.Evaluations.polynomial SMap.t ->
    PP.Evaluations.t SMap.t

  val worker_build_identities :
    shifts_map:(int * int) SMap.t ->
    prover_common_pp * prover_circuit_pp SMap.t ->
    f_wires_map_list_map:PP.Evaluations.polynomial SMap.t list SMap.t ->
    f_plook_map_list_map:PP.Evaluations.polynomial SMap.t list SMap.t ->
    inputs_map:circuit_prover_input list SMap.t ->
    beta_plookup:scalar ->
    gamma_plookup:scalar ->
    PP.prover_identities

  val kzg_eval_at_x :
    PP.transcript ->
    (PP.PC.secret * PP.PC.Commitment.prover_aux) list ->
    scalar ->
    bool ->
    PP.PC.answer list
end

module Make
    (PP : Polynomial_protocol.S with type PC.BasePC.Scalar.t = Plompiler.S.t) =
struct
  module MP = Plonk.Main_protocol.Make_impl (PP)
  include (MP : module type of MP with module PP := PP)
  module PP = PP
  module Commitment = Commitment

  type commit_to_wires_reply = Commitment.t [@@deriving repr]

  type worker_inputs = { inputs : circuit_prover_input list; shift : int * int }
  [@@deriving repr]

  include Prover

  let split_inputs_map ~nb_workers inputs_map =
    let list_range i1 i2 = List.filteri (fun i _ -> i1 <= i && i < i2) in
    List.map
      (fun i ->
        SMap.map
          (fun l ->
            let n = List.length l in
            let chunk_size =
              Z.(cdiv (of_int n) (of_int nb_workers) |> to_int)
            in
            let inputs = list_range (chunk_size * i) (chunk_size * (i + 1)) l in
            let shift = (chunk_size * i, n) in
            { inputs; shift })
          inputs_map)
      (List.init nb_workers Fun.id)

  type commit_to_plook_reply = {
    batched_witness_map : Evaluations.t SMap.t;
    cmt_plookup : Commitment.t;
    f_plook_map : PP.PC.Polynomial.Polynomial.t SMap.t;
    plook_prover_aux : Commitment.prover_aux;
  }
  [@@deriving repr]

  type commit_to_plook_remember = {
    f_wires_map_list_map : PP.Evaluations.polynomial SMap.t list SMap.t;
    f_plook_map_list_map : PP.Evaluations.polynomial SMap.t list SMap.t;
    beta_plookup : scalar;
    gamma_plookup : scalar;
  }

  type commit_to_wires_remember = {
    f_wires_map : PP.Evaluations.polynomial SMap.t;
    wires_map : wires_info list SMap.t;
    inputs_map : circuit_prover_input list SMap.t;
    shifts_map : (int * int) SMap.t;
    srs : PP.prover_public_parameters;
    f_wires_map_list_map : PP.Evaluations.polynomial SMap.t list SMap.t;
    wires_prover_aux : PP.PC.Commitment.prover_aux;
  }

  let worker_commit_to_wires ?(zero_knowledge = false) (common_pp, circuits_map)
      worker_inputs_map =
    let inputs_map = SMap.map (fun wi -> wi.inputs) worker_inputs_map in
    let shifts_map = SMap.map (fun wi -> wi.shift) worker_inputs_map in
    let srs = common_pp.pp_public_parameters in
    let ( wires_map,
          commitment_wires,
          wires_prover_aux,
          f_wires_map_list_map,
          f_wires_map ) =
      commit_to_wires ~shifts_map ~zero_knowledge ~common_pp ~circuits_map
        ~inputs_map ()
    in
    ( commitment_wires,
      {
        f_wires_map;
        wires_map;
        inputs_map;
        shifts_map;
        srs;
        f_wires_map_list_map;
        wires_prover_aux;
      } )

  let worker_build_all_batched_witness delta wires_map =
    SMap.map (build_batched_witness delta) wires_map

  let commit_to_plook ?(zero_knowledge = false) (common_pp, circuits_map)
      wires_map shifts_map srs transcript f_wires_map_list_map =
    let betas_gammas, _transcript = Fr_generation.random_fr_list transcript 5 in
    let beta_plookup = List.nth betas_gammas 2 in
    let gamma_plookup = List.nth betas_gammas 3 in
    let delta = List.nth betas_gammas 4 in
    let batched_witness_map =
      worker_build_all_batched_witness delta wires_map
    in
    (* ******************************************* *)
    let f_plook_map_list_map =
      build_all_f_map_evaluation_plook ~zero_knowledge common_pp beta_plookup
        gamma_plookup circuits_map wires_map
    in
    let f_plook_map = gather_maps ~shifts_map f_plook_map_list_map in
    (* commit to the plookup polynomials *)
    let cmt_plookup, plook_prover_aux =
      (* FIXME: implement Plookup *)
      let all_keys = build_all_keys [ "plook" ] shifts_map in
      PP.PC.Commitment.commit ~all_keys srs f_plook_map
    in
    ( { batched_witness_map; cmt_plookup; f_plook_map; plook_prover_aux },
      {
        f_wires_map_list_map;
        f_plook_map_list_map;
        beta_plookup;
        gamma_plookup;
      } )

  let batch_evaluated_ids ~alpha evaluated_ids all_ids_keys =
    let powers_map =
      SMap.of_list @@ List.mapi (fun i s -> (s, i)) all_ids_keys
    in
    let ids_keys, evaluations = List.split @@ SMap.bindings evaluated_ids in
    let powers =
      List.map (fun s -> SMap.find s powers_map) ids_keys
      |> List.map (fun i -> Scalar.pow alpha @@ Z.of_int i)
    in
    Evaluations.linear_c ~evaluations ~linear_coeffs:powers ()

  let main_build_aggregated_wires_map ~zero_knowledge common_pp circuits_map
      circuit_name batched_witness =
    if zero_knowledge then failwith "Distribution with ZK is not supported";
    build_aggregated_wires_map ~zero_knowledge Scalar.zero common_pp
      circuits_map SMap.empty circuit_name batched_witness

  let worker_build_evaluations ((common_pp : prover_common_pp), circuits_map)
      ~f_wires_map ~f_plook_map =
    build_evaluations (common_pp, circuits_map)
      ~f_map:(SMap.union_disjoint f_wires_map f_plook_map)

  let worker_build_identities ~shifts_map = build_identities ~shifts_map

  let kzg_eval_at_x transcript secrets_worker generator gx_gates =
    let eval_points_worker =
      [ List.hd @@ List.rev @@ eval_points false gx_gates ]
    in
    let x, _transcript = Fr_generation.random_fr transcript in
    let polys_list_worker = List.map fst secrets_worker in
    let query_list_worker =
      List.map (PP.convert_eval_points ~generator ~x) eval_points_worker
    in
    List.map2 PP.PC.evaluate polys_list_worker query_list_worker
end