Merge pull request #14331 from MinaProtocol/feature/surely-not-more-c…

…hunks

Add support to wrap prover for verifying chunked step proofs

src/lib/crypto/kimchi_bindings/stubs/kimchi_types.ml

@@ -254,5 +254,6 @@ module VerifierIndex = struct
     ; evals : 'poly_comm verification_evals
     ; shifts : 'fr array
     ; lookup_index : 'poly_comm Lookup.t option
+    ; zk_rows : int
     }
 end

src/lib/crypto/kimchi_bindings/stubs/src/linearization.rs

@@ -34,7 +34,7 @@ pub fn linearization_strings<F: ark_ff::PrimeField + ark_ff::SquareRootField>(
         })
     };
     let evaluated_cols = linearization_columns::<F>(features.as_ref());
-    let (linearization, _powers_of_alpha) = constraints_expr::<F>(features.as_ref(), true, 3);
+    let (linearization, _powers_of_alpha) = constraints_expr::<F>(features.as_ref(), true);
 
     let Linearization {
         constant_term,

src/lib/crypto/kimchi_bindings/stubs/src/pasta_fp_plonk_index.rs

@@ -164,7 +164,7 @@ pub fn caml_pasta_fp_plonk_index_read(
     )?;
     t.srs = srs.clone();
 
-    let (linearization, powers_of_alpha) = expr_linearization(Some(&t.cs.feature_flags), true, 3);
+    let (linearization, powers_of_alpha) = expr_linearization(Some(&t.cs.feature_flags), true);
     t.linearization = linearization;
     t.powers_of_alpha = powers_of_alpha;
 

src/lib/crypto/kimchi_bindings/stubs/src/pasta_fp_plonk_verifier_index.rs

@@ -59,6 +59,7 @@ impl From<VerifierIndex<Vesta, OpeningProof<Vesta>>> for CamlPastaFpPlonkVerifie
             },
             shifts: vi.shift.to_vec().iter().map(Into::into).collect(),
             lookup_index: vi.lookup_index.map(Into::into),
+            zk_rows: vi.zk_rows as isize,
         }
     }
 }
@@ -110,7 +111,8 @@ impl From<CamlPastaFpPlonkVerifierIndex> for VerifierIndex<Vesta, OpeningProof<V
         };
 
         // TODO dummy_lookup_value ?
-        let (linearization, powers_of_alpha) = expr_linearization(Some(&feature_flags), true, 3);
+        let (linearization, powers_of_alpha) =
+            expr_linearization(Some(&feature_flags), true);
 
         VerifierIndex::<Vesta, OpeningProof<Vesta>> {
             domain,
@@ -120,7 +122,7 @@ impl From<CamlPastaFpPlonkVerifierIndex> for VerifierIndex<Vesta, OpeningProof<V
             powers_of_alpha,
             srs: { Arc::clone(&index.srs.0) },
 
-            zk_rows: 3,
+            zk_rows: index.zk_rows as u64,
 
             sigma_comm,
             coefficients_comm,
@@ -143,12 +145,16 @@ impl From<CamlPastaFpPlonkVerifierIndex> for VerifierIndex<Vesta, OpeningProof<V
             shift,
             permutation_vanishing_polynomial_m: {
                 let res = once_cell::sync::OnceCell::new();
-                res.set(permutation_vanishing_polynomial(domain, 3)).unwrap();
+                res.set(permutation_vanishing_polynomial(
+                    domain,
+                    index.zk_rows as u64,
+                ))
+                .unwrap();
                 res
             },
             w: {
                 let res = once_cell::sync::OnceCell::new();
-                res.set(zk_w(domain, 3)).unwrap();
+                res.set(zk_w(domain, index.zk_rows as u64)).unwrap();
                 res
             },
             endo: endo_q,
@@ -273,6 +279,7 @@ pub fn caml_pasta_fp_plonk_verifier_index_dummy() -> CamlPastaFpPlonkVerifierInd
         },
         shifts: (0..PERMUTS - 1).map(|_| Fp::one().into()).collect(),
         lookup_index: None,
+        zk_rows: 3,
     }
 }
 

src/lib/crypto/kimchi_bindings/stubs/src/pasta_fq_plonk_index.rs

@@ -163,7 +163,7 @@ pub fn caml_pasta_fq_plonk_index_read(
     )?;
     t.srs = srs.clone();
 
-    let (linearization, powers_of_alpha) = expr_linearization(Some(&t.cs.feature_flags), true, 3);
+    let (linearization, powers_of_alpha) = expr_linearization(Some(&t.cs.feature_flags), true);
     t.linearization = linearization;
     t.powers_of_alpha = powers_of_alpha;
 

src/lib/crypto/kimchi_bindings/stubs/src/pasta_fq_plonk_verifier_index.rs

@@ -58,6 +58,7 @@ impl From<VerifierIndex<Pallas, OpeningProof<Pallas>>> for CamlPastaFqPlonkVerif
             },
             shifts: vi.shift.to_vec().iter().map(Into::into).collect(),
             lookup_index: vi.lookup_index.map(Into::into),
+            zk_rows: vi.zk_rows as isize,
         }
     }
 }
@@ -109,7 +110,8 @@ impl From<CamlPastaFqPlonkVerifierIndex> for VerifierIndex<Pallas, OpeningProof<
         };
 
         // TODO dummy_lookup_value ?
-        let (linearization, powers_of_alpha) = expr_linearization(Some(&feature_flags), true, 3);
+        let (linearization, powers_of_alpha) =
+            expr_linearization(Some(&feature_flags), true);
 
         VerifierIndex::<Pallas, OpeningProof<Pallas>> {
             domain,
@@ -119,7 +121,7 @@ impl From<CamlPastaFqPlonkVerifierIndex> for VerifierIndex<Pallas, OpeningProof<
             powers_of_alpha,
             srs: { Arc::clone(&index.srs.0) },
 
-            zk_rows: 3,
+            zk_rows: index.zk_rows as u64,
 
             sigma_comm,
             coefficients_comm,
@@ -142,12 +144,16 @@ impl From<CamlPastaFqPlonkVerifierIndex> for VerifierIndex<Pallas, OpeningProof<
             shift,
             permutation_vanishing_polynomial_m: {
                 let res = once_cell::sync::OnceCell::new();
-                res.set(permutation_vanishing_polynomial(domain, 3)).unwrap();
+                res.set(permutation_vanishing_polynomial(
+                    domain,
+                    index.zk_rows as u64,
+                ))
+                .unwrap();
                 res
             },
             w: {
                 let res = once_cell::sync::OnceCell::new();
-                res.set(zk_w(domain, 3)).unwrap();
+                res.set(zk_w(domain, index.zk_rows as u64)).unwrap();
                 res
             },
             endo: endo_q,
@@ -272,6 +278,7 @@ pub fn caml_pasta_fq_plonk_verifier_index_dummy() -> CamlPastaFqPlonkVerifierInd
         },
         shifts: (0..PERMUTS - 1).map(|_| Fq::one().into()).collect(),
         lookup_index: None,
+        zk_rows: 3,
     }
 }
 

src/lib/crypto/kimchi_bindings/stubs/src/plonk_verifier_index.rs

@@ -194,4 +194,5 @@ pub struct CamlPlonkVerifierIndex<Fr, SRS, PolyComm> {
     pub evals: CamlPlonkVerificationEvals<PolyComm>,
     pub shifts: Vec<Fr>,
     pub lookup_index: Option<CamlLookupVerifierIndex<PolyComm>>,
+    pub zk_rows: ocaml::Int,
 }

src/lib/mina_wire_types/pickles/pickles.ml

@@ -99,6 +99,10 @@ module M = struct
                 ( Snark_params.Tick.Field.t
                 , Snark_params.Tick.Field.t array )
                 Pickles_types.Plonk_types.All_evals.t
+                  (* A job half-done may be worse than not done at all.
+                     TODO: Migrate Plonk_types here, and actually include the
+                     *wire* type, not this in-memory version.
+                  *)
             ; proof : Wrap_wire_proof.V1.t
             }
         end

src/lib/pickles/commitment_lengths.ml

@@ -1,9 +1,8 @@
 open Pickles_types
 
-let default =
-  let one = 1 in
+let default ~num_chunks =
   { Plonk_types.Messages.Poly.w =
-      Vector.init Plonk_types.Columns.n ~f:(fun _ -> one)
-  ; z = one
-  ; t = 7
+      Vector.init Plonk_types.Columns.n ~f:(fun _ -> num_chunks)
+  ; z = num_chunks
+  ; t = 7 * num_chunks
   }

src/lib/pickles/commitment_lengths.mli

@@ -16,7 +16,8 @@
     length of vectors, and avoid runtime checks.
 *)
 val default :
-  ( int Pickles_types.Plonk_types.Columns_vec.t
-  , int
-  , int )
-  Pickles_types.Plonk_types.Messages.Poly.t
+     num_chunks:int
+  -> ( int Pickles_types.Plonk_types.Columns_vec.t
+     , int
+     , int )
+     Pickles_types.Plonk_types.Messages.Poly.t

src/lib/pickles/compile.ml

@@ -348,6 +348,7 @@ struct
       -> ?override_wrap_domain:Pickles_base.Proofs_verified.t
       -> ?override_wrap_main:
            (max_proofs_verified, branches, prev_varss) wrap_main_generic
+      -> ?num_chunks:int
       -> branches:(module Nat.Intf with type n = branches)
       -> max_proofs_verified:
            (module Nat.Add.Intf with type n = max_proofs_verified)
@@ -380,8 +381,9 @@ struct
          * _ =
    fun ~self ~cache ~proof_cache ?disk_keys
        ?(return_early_digest_exception = false) ?override_wrap_domain
-       ?override_wrap_main ~branches:(module Branches) ~max_proofs_verified
-       ~name ~constraint_constants ~public_input ~auxiliary_typ ~choices () ->
+       ?override_wrap_main ?(num_chunks = 1) ~branches:(module Branches)
+       ~max_proofs_verified ~name ~constraint_constants ~public_input
+       ~auxiliary_typ ~choices () ->
     let snark_keys_header kind constraint_system_hash =
       { Snark_keys_header.header_version = Snark_keys_header.header_version
       ; kind
@@ -454,7 +456,7 @@ struct
               (Auxiliary_value)
           in
           M.f full_signature prev_varss_n prev_varss_length ~max_proofs_verified
-            ~feature_flags
+            ~feature_flags ~num_chunks
       | Some override ->
           Common.wrap_domains
             ~proofs_verified:(Pickles_base.Proofs_verified.to_int override)
@@ -618,7 +620,7 @@ struct
       match override_wrap_main with
       | None ->
           let srs = Tick.Keypair.load_urs () in
-          Wrap_main.wrap_main ~feature_flags ~srs full_signature
+          Wrap_main.wrap_main ~num_chunks ~feature_flags ~srs full_signature
             prev_varss_length step_vks proofs_verifieds step_domains
             max_proofs_verified
       | Some { wrap_main; tweak_statement = _ } ->
@@ -743,7 +745,9 @@ struct
                            , return_value
                            , auxiliary_value
                            , actual_wrap_domains ) =
-            step ~proof_cache handler ~maxes:(module Maxes) next_state
+            step ~zk_rows:step_vk.zk_rows ~proof_cache handler
+              ~maxes:(module Maxes)
+              next_state
           in
           let proof =
             { proof with
@@ -946,6 +950,7 @@ let compile_with_wrap_main_override_promise :
     -> ?override_wrap_domain:Pickles_base.Proofs_verified.t
     -> ?override_wrap_main:
          (max_proofs_verified, branches, prev_varss) wrap_main_generic
+    -> ?num_chunks:int
     -> public_input:
          ( var
          , value
@@ -993,7 +998,7 @@ let compile_with_wrap_main_override_promise :
  *)
  fun ?self ?(cache = []) ?proof_cache ?disk_keys
      ?(return_early_digest_exception = false) ?override_wrap_domain
-     ?override_wrap_main ~public_input ~auxiliary_typ ~branches
+     ?override_wrap_main ?num_chunks ~public_input ~auxiliary_typ ~branches
      ~max_proofs_verified ~name ~constraint_constants ~choices () ->
   let self =
     match self with
@@ -1061,7 +1066,7 @@ let compile_with_wrap_main_override_promise :
   in
   let provers, wrap_vk, wrap_disk_key, cache_handle =
     M.compile ~return_early_digest_exception ~self ~proof_cache ~cache
-      ?disk_keys ?override_wrap_domain ?override_wrap_main ~branches
+      ?disk_keys ?override_wrap_domain ?override_wrap_main ?num_chunks ~branches
       ~max_proofs_verified ~name ~public_input ~auxiliary_typ
       ~constraint_constants
       ~choices:(fun ~self -> conv_irs (choices ~self))

src/lib/pickles/compile.mli

@@ -278,6 +278,7 @@ val compile_with_wrap_main_override_promise :
   -> ?override_wrap_domain:Pickles_base.Proofs_verified.t
   -> ?override_wrap_main:
        ('max_proofs_verified, 'branches, 'prev_varss) wrap_main_generic
+  -> ?num_chunks:int
   -> public_input:
        ( 'var
        , 'value

src/lib/pickles/dummy.ml

@@ -16,7 +16,7 @@ let evals =
     in
     let ex =
       { All_evals.With_public_input.evals = e
-      ; public_input = (Ro.tock (), Ro.tock ())
+      ; public_input = ([| Ro.tock () |], [| Ro.tock () |])
       }
     in
     { All_evals.ft_eval1 = Ro.tock (); evals = ex })

src/lib/pickles/pickles.ml

@@ -307,21 +307,21 @@ module Make_str (_ : Wire_types.Concrete) = struct
     Compile.compile_with_wrap_main_override_promise
 
   let compile_promise ?self ?cache ?proof_cache ?disk_keys
-      ?return_early_digest_exception ?override_wrap_domain ~public_input
-      ~auxiliary_typ ~branches ~max_proofs_verified ~name ~constraint_constants
-      ~choices () =
+      ?return_early_digest_exception ?override_wrap_domain ?num_chunks
+      ~public_input ~auxiliary_typ ~branches ~max_proofs_verified ~name
+      ~constraint_constants ~choices () =
     compile_with_wrap_main_override_promise ?self ?cache ?proof_cache ?disk_keys
-      ?return_early_digest_exception ?override_wrap_domain ~public_input
-      ~auxiliary_typ ~branches ~max_proofs_verified ~name ~constraint_constants
-      ~choices ()
+      ?return_early_digest_exception ?override_wrap_domain ?num_chunks
+      ~public_input ~auxiliary_typ ~branches ~max_proofs_verified ~name
+      ~constraint_constants ~choices ()
 
   let compile ?self ?cache ?proof_cache ?disk_keys ?override_wrap_domain
-      ~public_input ~auxiliary_typ ~branches ~max_proofs_verified ~name
-      ~constraint_constants ~choices () =
+      ?num_chunks ~public_input ~auxiliary_typ ~branches ~max_proofs_verified
+      ~name ~constraint_constants ~choices () =
     let self, cache_handle, proof_module, provers =
       compile_promise ?self ?cache ?proof_cache ?disk_keys ?override_wrap_domain
-        ~public_input ~auxiliary_typ ~branches ~max_proofs_verified ~name
-        ~constraint_constants ~choices ()
+        ?num_chunks ~public_input ~auxiliary_typ ~branches ~max_proofs_verified
+        ~name ~constraint_constants ~choices ()
     in
     let rec adjust_provers :
         type a1 a2 a3 s1 s2_inner.
@@ -1125,6 +1125,7 @@ module Make_str (_ : Wire_types.Concrete) = struct
               Wrap_domains.Make (A) (A_value) (A) (A_value) (A) (A_value)
             in
             M.f full_signature prev_varss_n prev_varss_length ~feature_flags
+              ~num_chunks:1
               ~max_proofs_verified:(module Max_proofs_verified)
           in
           let module Branch_data = struct
@@ -1304,7 +1305,8 @@ module Make_str (_ : Wire_types.Concrete) = struct
               let wrap =
                 let wrap_vk = Lazy.force wrap_vk in
                 let%bind.Promise proof, (), (), _ =
-                  step ~proof_cache:None ~maxes:(module Maxes)
+                  step ~zk_rows:pairing_vk.zk_rows ~proof_cache:None
+                    ~maxes:(module Maxes)
                 in
                 let proof =
                   { proof with
@@ -1561,6 +1563,7 @@ module Make_str (_ : Wire_types.Concrete) = struct
                           ~endo:Endo.Step_inner_curve.base
                           ~mds:Tick_field_sponge.params.mds
                           ~srs_length_log2:Common.Max_degree.step_log2
+                          ~zk_rows:3
                           ~field_of_hex:(fun s ->
                             Kimchi_pasta.Pasta.Bigint256.of_hex_string s
                             |> Kimchi_pasta.Pasta.Fp.of_bigint )
@@ -1574,7 +1577,9 @@ module Make_str (_ : Wire_types.Concrete) = struct
                           ~actual_proofs_verified:
                             (Nat.Add.create actual_proofs_verified)
                           { evals = proof.proof.openings.evals
-                          ; public_input = x_hat
+                          ; public_input =
+                              (let x1, x2 = x_hat in
+                               ([| x1 |], [| x2 |]) )
                           }
                           ~r ~xi ~zeta ~zetaw
                           ~old_bulletproof_challenges:prev_challenges
@@ -1770,7 +1775,9 @@ module Make_str (_ : Wire_types.Concrete) = struct
                             ~to_option:Opt.to_option next_statement
                       ; prev_evals =
                           { Plonk_types.All_evals.evals =
-                              { public_input = x_hat
+                              { public_input =
+                                  (let x1, x2 = x_hat in
+                                   ([| x1 |], [| x2 |]) )
                               ; evals = proof.proof.openings.evals
                               }
                           ; ft_eval1 = proof.proof.openings.ft_eval1

src/lib/pickles/pickles_intf.mli

@@ -372,6 +372,7 @@ module type S = sig
          * Cache.Wrap.Key.Verification.t
     -> ?return_early_digest_exception:bool
     -> ?override_wrap_domain:Pickles_base.Proofs_verified.t
+    -> ?num_chunks:int
     -> public_input:
          ( 'var
          , 'value
@@ -426,6 +427,7 @@ module type S = sig
          (Cache.Step.Key.Verification.t, 'branches) Vector.t
          * Cache.Wrap.Key.Verification.t
     -> ?override_wrap_domain:Pickles_base.Proofs_verified.t
+    -> ?num_chunks:int
     -> public_input:
          ( 'var
          , 'value