Merge pull request #14284 from MinaProtocol/feature/more-chunks

Expose chunks in `Plonk_verification_evals.t`

src/lib/crypto/kimchi_backend/common/dlog_plonk_based_keypair.ml

@@ -225,13 +225,13 @@ module Make (Inputs : Inputs_intf) = struct
     }
 
   let full_vk_commitments (t : Inputs.Verifier_index.t) :
-      ( Inputs.Curve.Affine.t
-      , Inputs.Curve.Affine.t option )
+      ( Inputs.Curve.Affine.t array
+      , Inputs.Curve.Affine.t array option )
       Pickles_types.Plonk_verification_key_evals.Step.t =
-    let g c : Inputs.Curve.Affine.t =
+    let g c : Inputs.Curve.Affine.t array =
       match Inputs.Poly_comm.of_backend_without_degree_bound c with
       | `Without_degree_bound x ->
-          x.(0)
+          x
       | `With_degree_bound _ ->
           assert false
     in

src/lib/pickles/common.ml

@@ -49,7 +49,8 @@ let hash_messages_for_next_step_proof ~app_state
   let g (x, y) = [ x; y ] in
   Tick_field_sponge.digest Tick_field_sponge.params
     (Types.Step.Proof_state.Messages_for_next_step_proof.to_field_elements t ~g
-       ~comm:(fun (x : Tock.Curve.Affine.t) -> Array.of_list (g x))
+       ~comm:(fun (x : Tock.Curve.Affine.t array) ->
+         Array.concat_map x ~f:(fun x -> Array.of_list (g x)) )
        ~app_state )
 
 let dlog_pcs_batch (type nat proofs_verified total)
@@ -230,21 +231,23 @@ let tick_public_input_of_statement ~max_proofs_verified
     ~f:(Backend.Tick.Field.Vector.get input)
 
 let ft_comm ~add:( + ) ~scale ~endoscale:_ ~negate
-    ~verification_key:(m : _ Plonk_verification_key_evals.t) ~alpha:_
+    ~verification_key:(m : _ array Plonk_verification_key_evals.t) ~alpha:_
     ~(plonk : _ Types.Wrap.Proof_state.Deferred_values.Plonk.In_circuit.t)
     ~t_comm =
-  let _, [ sigma_comm_last ] =
-    Vector.split m.sigma_comm (snd (Plonk_types.Permuts_minus_1.add Nat.N1.n))
-  in
-  let f_comm = List.reduce_exn ~f:( + ) [ scale sigma_comm_last plonk.perm ] in
-  let chunked_t_comm =
-    let n = Array.length t_comm in
-    let res = ref t_comm.(n - 1) in
+  let reduce_chunks comm =
+    let n = Array.length comm in
+    let res = ref comm.(n - 1) in
     for i = n - 2 downto 0 do
-      res := t_comm.(i) + scale !res plonk.zeta_to_srs_length
+      res := comm.(i) + scale !res plonk.zeta_to_srs_length
     done ;
     !res
   in
+  let _, [ sigma_comm_last ] =
+    Vector.split m.sigma_comm (snd (Plonk_types.Permuts_minus_1.add Nat.N1.n))
+  in
+  let sigma_comm_last = reduce_chunks sigma_comm_last in
+  let f_comm = List.reduce_exn ~f:( + ) [ scale sigma_comm_last plonk.perm ] in
+  let chunked_t_comm = reduce_chunks t_comm in
   f_comm + chunked_t_comm
   + negate (scale chunked_t_comm plonk.zeta_to_domain_size)
 

src/lib/pickles/common.mli

@@ -40,7 +40,7 @@ val ft_comm :
   -> scale:('comm -> 'scalar -> 'comm)
   -> endoscale:('comm -> 'c -> 'comm)
   -> negate:('comm -> 'comm)
-  -> verification_key:'comm Pickles_types.Plonk_verification_key_evals.t
+  -> verification_key:'comm array Pickles_types.Plonk_verification_key_evals.t
   -> alpha:'c
   -> plonk:
        ( 'd
@@ -145,7 +145,7 @@ end
 
 val hash_messages_for_next_step_proof :
      app_state:('a -> Kimchi_pasta.Basic.Fp.Stable.Latest.t Core_kernel.Array.t)
-  -> ( Backend.Tock.Curve.Affine.t
+  -> ( Backend.Tock.Curve.Affine.t array
      , 'a
      , ( Kimchi_pasta.Basic.Fp.Stable.Latest.t
          * Kimchi_pasta.Basic.Fp.Stable.Latest.t

src/lib/pickles/compile.ml

@@ -117,8 +117,8 @@ type ('max_proofs_verified, 'branches, 'prev_varss) wrap_main_generic =
          , 'max_local_max_proofs_verifieds )
          Full_signature.t
       -> ('prev_varss, 'branches) Hlist.Length.t
-      -> ( ( Wrap_main_inputs.Inner_curve.Constant.t
-           , Wrap_main_inputs.Inner_curve.Constant.t option )
+      -> ( ( Wrap_main_inputs.Inner_curve.Constant.t array
+           , Wrap_main_inputs.Inner_curve.Constant.t array option )
            Wrap_verifier.index'
          , 'branches )
          Vector.t
@@ -727,7 +727,11 @@ struct
         let step handler next_state =
           let wrap_vk = Lazy.force wrap_vk in
           S.f ?handler branch_data next_state ~prevs_length:prev_vars_length
-            ~self ~step_domains ~self_dlog_plonk_index:wrap_vk.commitments
+            ~self ~step_domains
+            ~self_dlog_plonk_index:
+              ((* TODO *) Plonk_verification_key_evals.map
+                 ~f:(fun x -> [| x |])
+                 wrap_vk.commitments )
             ~public_input ~auxiliary_typ ~feature_flags
             (Impls.Step.Keypair.pk (fst (Lazy.force step_pk)))
             wrap_vk.index
@@ -771,8 +775,12 @@ struct
             Wrap.wrap ~proof_cache ~max_proofs_verified:Max_proofs_verified.n
               ~feature_flags ~actual_feature_flags:b.feature_flags
               full_signature.maxes wrap_requests ?tweak_statement
-              ~dlog_plonk_index:wrap_vk.commitments wrap_main ~typ ~step_vk
-              ~step_plonk_indices:(Lazy.force step_vks) ~actual_wrap_domains
+              ~dlog_plonk_index:
+                ((* TODO *) Plonk_verification_key_evals.map
+                   ~f:(fun x -> [| x |])
+                   wrap_vk.commitments )
+              wrap_main ~typ ~step_vk ~step_plonk_indices:(Lazy.force step_vks)
+              ~actual_wrap_domains
               (Impls.Wrap.Keypair.pk (fst (Lazy.force wrap_pk)))
               proof
           in
@@ -819,7 +827,10 @@ struct
       ; proofs_verifieds
       ; max_proofs_verified
       ; public_input = typ
-      ; wrap_key = Lazy.map wrap_vk ~f:Verification_key.commitments
+      ; wrap_key =
+          Lazy.map wrap_vk ~f:(fun x ->
+              Plonk_verification_key_evals.map (Verification_key.commitments x)
+                ~f:(fun x -> [| x |]) )
       ; wrap_vk = Lazy.map wrap_vk ~f:Verification_key.index
       ; wrap_domains
       ; step_domains
@@ -847,7 +858,9 @@ module Side_loaded = struct
           ~log_2_domain_size:(Lazy.force d.wrap_vk).domain.log_size_of_group
       in
       { wrap_vk = Some (Lazy.force d.wrap_vk)
-      ; wrap_index = Lazy.force d.wrap_key
+      ; wrap_index =
+          Plonk_verification_key_evals.map (Lazy.force d.wrap_key) ~f:(fun x ->
+              x.(0) )
       ; max_proofs_verified =
           Pickles_base.Proofs_verified.of_nat (Nat.Add.n d.max_proofs_verified)
       ; actual_wrap_domain_size

src/lib/pickles/compile.mli

@@ -154,8 +154,8 @@ type ('max_proofs_verified, 'branches, 'prev_varss) wrap_main_generic =
          , 'max_local_max_proofs_verifieds )
          Full_signature.t
       -> ('prev_varss, 'branches) Hlist.Length.t
-      -> ( ( Wrap_main_inputs.Inner_curve.Constant.t
-           , Wrap_main_inputs.Inner_curve.Constant.t option )
+      -> ( ( Wrap_main_inputs.Inner_curve.Constant.t array
+           , Wrap_main_inputs.Inner_curve.Constant.t array option )
            Wrap_verifier.index'
          , 'branches )
          Vector.t

src/lib/pickles/pickles.ml

@@ -224,7 +224,9 @@ module Make_str (_ : Wire_types.Concrete) = struct
             ~log_2_domain_size:(Lazy.force d.wrap_vk).domain.log_size_of_group
         in
         { wrap_vk = Some (Lazy.force d.wrap_vk)
-        ; wrap_index = Lazy.force d.wrap_key
+        ; wrap_index =
+            Plonk_verification_key_evals.map (Lazy.force d.wrap_key)
+              ~f:(fun x -> x.(0))
         ; max_proofs_verified =
             Pickles_base.Proofs_verified.of_nat
               (Nat.Add.n d.max_proofs_verified)
@@ -1291,7 +1293,10 @@ module Make_str (_ : Wire_types.Concrete) = struct
                 S.f branch_data () ~feature_flags ~prevs_length:prev_vars_length
                   ~self ~public_input:(Input typ)
                   ~auxiliary_typ:Impls.Step.Typ.unit ~step_domains
-                  ~self_dlog_plonk_index:wrap_vk.commitments
+                  ~self_dlog_plonk_index:
+                    ((* TODO *) Plonk_verification_key_evals.map
+                       ~f:(fun x -> [| x |])
+                       wrap_vk.commitments )
                   (Impls.Step.Keypair.pk (fst (Lazy.force step_pk)))
                   wrap_vk.index
               in
@@ -1772,7 +1777,11 @@ module Make_str (_ : Wire_types.Concrete) = struct
                       : _ P.Base.Wrap.t )
                   in
                   wrap ~max_proofs_verified:Max_proofs_verified.n
-                    full_signature.maxes ~dlog_plonk_index:wrap_vk.commitments
+                    full_signature.maxes
+                    ~dlog_plonk_index:
+                      ((* TODO *) Plonk_verification_key_evals.map
+                         ~f:(fun x -> [| x |])
+                         wrap_vk.commitments )
                     wrap_main A_value.to_field_elements ~pairing_vk
                     ~step_domains:b.domains
                     ~pairing_plonk_indices:(Lazy.force step_vks) ~wrap_domains
@@ -1800,7 +1809,12 @@ module Make_str (_ : Wire_types.Concrete) = struct
             ; proofs_verifieds
             ; max_proofs_verified = (module Max_proofs_verified)
             ; public_input = typ
-            ; wrap_key = Lazy.map wrap_vk ~f:Verification_key.commitments
+            ; wrap_key =
+                Lazy.map wrap_vk ~f:(fun x ->
+                    (* TODO *)
+                    Plonk_verification_key_evals.map
+                      ~f:(fun x -> [| x |])
+                      (Verification_key.commitments x) )
             ; wrap_vk = Lazy.map wrap_vk ~f:Verification_key.index
             ; wrap_domains
             ; step_domains

src/lib/pickles/requests.ml

@@ -129,7 +129,7 @@ module Step = struct
           , local_branches )
           H3.T(Per_proof_witness.Constant.No_app_state).t
           t
-      | Wrap_index : Tock.Curve.Affine.t Plonk_verification_key_evals.t t
+      | Wrap_index : Tock.Curve.Affine.t array Plonk_verification_key_evals.t t
       | App_state : statement t
       | Return_value : return_value -> unit t
       | Auxiliary_value : auxiliary_value -> unit t
@@ -183,7 +183,8 @@ module Step = struct
             , local_branches )
             H3.T(Per_proof_witness.Constant.No_app_state).t
             t
-        | Wrap_index : Tock.Curve.Affine.t Plonk_verification_key_evals.t t
+        | Wrap_index :
+            Tock.Curve.Affine.t array Plonk_verification_key_evals.t t
         | App_state : statement t
         | Return_value : return_value -> unit t
         | Auxiliary_value : auxiliary_value -> unit t

src/lib/pickles/requests.mli

@@ -35,7 +35,7 @@ module Step : sig
           Hlist.H3.T(Per_proof_witness.Constant.No_app_state).t
           Snarky_backendless.Request.t
       | Wrap_index :
-          Backend.Tock.Curve.Affine.t Plonk_verification_key_evals.t
+          Backend.Tock.Curve.Affine.t array Plonk_verification_key_evals.t
           Snarky_backendless.Request.t
       | App_state : statement Snarky_backendless.Request.t
       | Return_value : return_value -> unit Snarky_backendless.Request.t

src/lib/pickles/step.ml

@@ -121,7 +121,7 @@ struct
     let expand_proof :
         type var value local_max_proofs_verified m.
            Impls.Wrap.Verification_key.t
-        -> 'a
+        -> _ array Plonk_verification_key_evals.t
         -> value
         -> (local_max_proofs_verified, local_max_proofs_verified) Proof.t
         -> (var, value, local_max_proofs_verified, m) Tag.t

src/lib/pickles/step.mli

@@ -33,7 +33,7 @@ module Make
     -> step_domains:(Import.Domains.t, 'self_branches) Pickles_types.Vector.t
     -> feature_flags:Opt.Flag.t Plonk_types.Features.Full.t
     -> self_dlog_plonk_index:
-         Backend.Tick.Inner_curve.Affine.t
+         Backend.Tick.Inner_curve.Affine.t array
          Pickles_types.Plonk_verification_key_evals.t
     -> public_input:
          ( 'var

src/lib/pickles/step_main.ml

@@ -339,9 +339,11 @@ let step_main :
             in
             Req.Compute_prev_proof_parts previous_proof_statements ) ;
         let dlog_plonk_index =
+          let num_chunks = (* TODO *) 1 in
           exists
             ~request:(fun () -> Req.Wrap_index)
-            (Plonk_verification_key_evals.typ Inner_curve.typ)
+            (Plonk_verification_key_evals.typ
+               (Typ.array ~length:num_chunks Inner_curve.typ) )
         and prevs =
           exists (Prev_typ.f prev_proof_typs) ~request:(fun () ->
               Req.Proof_with_datas )

src/lib/pickles/step_verifier.ml

@@ -241,6 +241,13 @@ struct
         let combined_polynomial (* Corresponds to xi in figure 7 of WTS *) =
           with_label "combined_polynomial" (fun () ->
               Pcs_batch.combine_split_commitments pcs_batch
+                ~reduce_without_degree_bound:Array.to_list
+                ~reduce_with_degree_bound:(fun { Plonk_types.Poly_comm
+                                                 .With_degree_bound
+                                                 .unshifted
+                                               ; shifted
+                                               } ->
+                  Array.to_list unshifted @ [ shifted ] )
                 ~scale_and_add:(fun ~(acc :
                                        [ `Maybe_finite of
                                          Boolean.var * Inner_curve.t
@@ -572,14 +579,11 @@ struct
           let without_degree_bound =
             Vector.append
               (Vector.map sg_old ~f:(fun g -> [| g |]))
-              ( [| x_hat |] :: [| ft_comm |] :: z_comm :: [| m.generic_comm |]
-              :: [| m.psm_comm |] :: [| m.complete_add_comm |]
-              :: [| m.mul_comm |] :: [| m.emul_comm |]
-              :: [| m.endomul_scalar_comm |]
+              ( [| x_hat |] :: [| ft_comm |] :: z_comm :: m.generic_comm
+              :: m.psm_comm :: m.complete_add_comm :: m.mul_comm :: m.emul_comm
+              :: m.endomul_scalar_comm
               :: Vector.append w_comm
-                   (Vector.append
-                      (Vector.map m.coefficients_comm ~f:(fun g -> [| g |]))
-                      (Vector.map sigma_comm_init ~f:(fun g -> [| g |]))
+                   (Vector.append m.coefficients_comm sigma_comm_init
                       (snd Plonk_types.(Columns.add Permuts_minus_1.n)) )
                    (snd
                       Plonk_types.(
@@ -1107,8 +1111,9 @@ struct
     let sponge = Sponge.create sponge_params in
     Array.iter
       (Types.index_to_field_elements
-         ~g:(fun (z : Inputs.Inner_curve.t) ->
-           List.to_array (Inner_curve.to_field_elements z) )
+         ~g:
+           (Array.concat_map ~f:(fun (z : Inputs.Inner_curve.t) ->
+                List.to_array (Inner_curve.to_field_elements z) ) )
          index )
       ~f:(fun x -> Sponge.absorb sponge (`Field x)) ;
     sponge

src/lib/pickles/step_verifier.mli

@@ -76,7 +76,7 @@ val finalize_other_proof :
 
 val hash_messages_for_next_step_proof :
      index:
-       Step_main_inputs.Inner_curve.t
+       Step_main_inputs.Inner_curve.t array
        Pickles_types.Plonk_verification_key_evals.t
   -> ('s -> Step_main_inputs.Impl.Field.t array)
   -> (   ( 'a
@@ -91,7 +91,7 @@ val hash_messages_for_next_step_proof :
 
 val hash_messages_for_next_step_proof_opt :
      index:
-       Step_main_inputs.Inner_curve.t
+       Step_main_inputs.Inner_curve.t array
        Pickles_types.Plonk_verification_key_evals.t
   -> ('s -> Step_main_inputs.Impl.Field.t array)
   -> Step_main_inputs.Sponge.t
@@ -135,7 +135,7 @@ val verify :
          Step_main_inputs.Impl.field
          Composition_types.Branch_data.Proofs_verified.One_hot.Checked.t ]
   -> wrap_verification_key:
-       Step_main_inputs.Inner_curve.t
+       Step_main_inputs.Inner_curve.t array
        Pickles_types.Plonk_verification_key_evals.t
   -> ( Step_main_inputs.Impl.field Limb_vector.Challenge.t
      , Step_main_inputs.Impl.field Limb_vector.Challenge.t

src/lib/pickles/types_map.ml

@@ -16,7 +16,7 @@ module Basic = struct
     ; public_input : ('var, 'value) Impls.Step.Typ.t
     ; branches : 'n2 Nat.t
     ; wrap_domains : Domains.t
-    ; wrap_key : Tick.Inner_curve.Affine.t Plonk_verification_key_evals.t
+    ; wrap_key : Tick.Inner_curve.Affine.t array Plonk_verification_key_evals.t
     ; wrap_vk : Impls.Wrap.Verification_key.t
     ; feature_flags : Opt.Flag.t Plonk_types.Features.Full.t
     }
@@ -59,7 +59,10 @@ module Side_loaded = struct
     let wrap_key, wrap_vk =
       match ephemeral with
       | Some { index = `In_prover i | `In_both (i, _) } ->
-          (i.wrap_index, i.wrap_vk)
+          let wrap_index =
+            Plonk_verification_key_evals.map i.wrap_index ~f:(fun x -> [| x |])
+          in
+          (wrap_index, i.wrap_vk)
       | _ ->
           failwithf "Side_loaded.to_basic: Expected `In_prover (%s)" __LOC__ ()
     in
@@ -95,7 +98,8 @@ module Compiled = struct
     ; proofs_verifieds : (int, 'branches) Vector.t
           (* For each branch in this rule, how many predecessor proofs does it have? *)
     ; public_input : ('a_var, 'a_value) Impls.Step.Typ.t
-    ; wrap_key : Tick.Inner_curve.Affine.t Plonk_verification_key_evals.t Lazy.t
+    ; wrap_key :
+        Tick.Inner_curve.Affine.t array Plonk_verification_key_evals.t Lazy.t
     ; wrap_vk : Impls.Wrap.Verification_key.t Lazy.t
     ; wrap_domains : Domains.t
     ; step_domains : (Domains.t, 'branches) Vector.t
@@ -134,7 +138,7 @@ module For_step = struct
     ; proofs_verifieds :
         [ `Known of (Impls.Step.Field.t, 'branches) Vector.t | `Side_loaded ]
     ; public_input : ('a_var, 'a_value) Impls.Step.Typ.t
-    ; wrap_key : inner_curve_var Plonk_verification_key_evals.t
+    ; wrap_key : inner_curve_var array Plonk_verification_key_evals.t
     ; wrap_domain :
         [ `Known of Domain.t
         | `Side_loaded of
@@ -157,11 +161,14 @@ module For_step = struct
           failwithf "For_step.side_loaded: Expected `In_circuit (%s)" __LOC__ ()
     in
     let T = Nat.eq_exn branches Side_loaded_verification_key.Max_branches.n in
+    let wrap_key =
+      Plonk_verification_key_evals.map index.wrap_index ~f:(fun x -> [| x |])
+    in
     { branches
     ; max_proofs_verified
     ; public_input
     ; proofs_verifieds = `Side_loaded
-    ; wrap_key = index.wrap_index
+    ; wrap_key
     ; wrap_domain = `Side_loaded index.actual_wrap_domain_size
     ; step_domains = `Side_loaded
     ; feature_flags
@@ -186,7 +193,7 @@ module For_step = struct
     ; public_input
     ; wrap_key =
         Plonk_verification_key_evals.map (Lazy.force wrap_key)
-          ~f:Step_main_inputs.Inner_curve.constant
+          ~f:(Array.map ~f:Step_main_inputs.Inner_curve.constant)
     ; wrap_domain = `Known wrap_domains.h
     ; step_domains = `Known step_domains
     ; feature_flags