Merge pull request #14344 from MinaProtocol/feature/recursive-chunks-too

Chunking = fin

src/lib/pickles/compile.ml

@@ -25,6 +25,9 @@ open Kimchi_backend
 module Proof_ = P.Base
 module Proof = P
 
+type chunking_data = Verify.Instance.chunking_data =
+  { num_chunks : int; domain_size : int; zk_rows : int }
+
 let pad_messages_for_next_wrap_proof
     (type local_max_proofs_verifieds max_local_max_proofs_verifieds
     max_proofs_verified )
@@ -506,7 +509,7 @@ struct
               Timer.clock __LOC__ ;
               let res =
                 Common.time "make step data" (fun () ->
-                    Step_branch_data.create ~index:!i ~feature_flags
+                    Step_branch_data.create ~index:!i ~feature_flags ~num_chunks
                       ~actual_feature_flags:rule.feature_flags
                       ~max_proofs_verified:Max_proofs_verified.n
                       ~branches:Branches.n ~self ~public_input ~auxiliary_typ
@@ -745,9 +748,7 @@ struct
                            , return_value
                            , auxiliary_value
                            , actual_wrap_domains ) =
-            step ~zk_rows:step_vk.zk_rows ~proof_cache handler
-              ~maxes:(module Maxes)
-              next_state
+            step ~proof_cache handler ~maxes:(module Maxes) next_state
           in
           let proof =
             { proof with
@@ -839,6 +840,14 @@ struct
       ; wrap_domains
       ; step_domains
       ; feature_flags
+      ; num_chunks
+      ; zk_rows =
+          ( match num_chunks with
+          | 1 ->
+              3
+          | num_chunks ->
+              let permuts = 7 in
+              ((2 * (permuts + 1) * num_chunks) - 1 + permuts) / permuts )
       }
     in
     Timer.clock __LOC__ ;
@@ -884,6 +893,8 @@ module Side_loaded = struct
       ; branches = Verification_key.Max_branches.n
       ; feature_flags =
           Plonk_types.Features.to_full ~or_:Opt.Flag.( ||| ) feature_flags
+      ; num_chunks = 1
+      ; zk_rows = 3
       }
 
   module Proof = struct
@@ -926,7 +937,7 @@ module Side_loaded = struct
                   { constraints = 0 }
               }
             in
-            Verify.Instance.T (max_proofs_verified, m, vk, x, p) )
+            Verify.Instance.T (max_proofs_verified, m, None, vk, x, p) )
         |> Verify.verify_heterogenous )
 
   let verify ~typ ts = verify_promise ~typ ts |> Promise.to_deferred
@@ -1090,6 +1101,26 @@ let compile_with_wrap_main_override_promise :
       let (Typ typ) = typ in
       fun x -> fst (typ.value_to_fields x)
   end in
+  let chunking_data =
+    match num_chunks with
+    | None ->
+        None
+    | Some num_chunks ->
+        let compiled = Types_map.lookup_compiled self.id in
+        let { h = Pow_2_roots_of_unity domain_size } =
+          compiled.step_domains
+          |> Vector.reduce_exn
+               ~f:(fun
+                    { h = Pow_2_roots_of_unity d1 }
+                    { h = Pow_2_roots_of_unity d2 }
+                  -> { h = Pow_2_roots_of_unity (Int.max d1 d2) } )
+        in
+        Some
+          { Verify.Instance.num_chunks
+          ; domain_size
+          ; zk_rows = compiled.zk_rows
+          }
+  in
   let module P = struct
     type statement = value
 
@@ -1109,7 +1140,7 @@ let compile_with_wrap_main_override_promise :
     let verification_key = wrap_vk
 
     let verify_promise ts =
-      verify_promise
+      verify_promise ?chunking_data
         ( module struct
           include Max_proofs_verified
         end )

src/lib/pickles/compile.mli

@@ -44,8 +44,12 @@ module type Proof_intf = sig
   val verify_promise : (statement * t) list -> unit Or_error.t Promise.t
 end
 
+type chunking_data = Verify.Instance.chunking_data =
+  { num_chunks : int; domain_size : int; zk_rows : int }
+
 val verify_promise :
-     (module Nat.Intf with type n = 'n)
+     ?chunking_data:chunking_data
+  -> (module Nat.Intf with type n = 'n)
   -> (module Statement_value_intf with type t = 'a)
   -> Verification_key.t
   -> ('a * ('n, 'n) Proof.t) list

src/lib/pickles/per_proof_witness.ml

@@ -134,7 +134,7 @@ module Constant = struct
   end
 end
 
-let typ (type n avar aval) ~feature_flags
+let typ (type n avar aval) ~feature_flags ~num_chunks
     (statement : (avar, aval) Impls.Step.Typ.t) (max_proofs_verified : n Nat.t)
     =
   let module Sc = Scalar_challenge in
@@ -156,7 +156,7 @@ let typ (type n avar aval) ~feature_flags
         (Branch_data.typ
            (module Impl)
            ~assert_16_bits:(Step_verifier.assert_n_bits ~n:16) )
-    ; Plonk_types.All_evals.typ
+    ; Plonk_types.All_evals.typ ~num_chunks
         (module Impl)
         (* Assume we have lookup iff we have runtime tables *)
         feature_flags

src/lib/pickles/per_proof_witness.mli

@@ -113,6 +113,7 @@ end
 
 val typ :
      feature_flags:Opt.Flag.t Plonk_types.Features.Full.t
+  -> num_chunks:int
   -> ('avar, 'aval) Impl.Typ.t
   -> 'n Pickles_types.Nat.t
   -> (('avar, 'n, _) t, ('aval, 'n) Constant.t) Impl.Typ.t

src/lib/pickles/pickles.ml

@@ -47,6 +47,9 @@ module Make_str (_ : Wire_types.Concrete) = struct
 
   exception Return_digest = Compile.Return_digest
 
+  type chunking_data = Verify.Instance.chunking_data =
+    { num_chunks : int; domain_size : int; zk_rows : int }
+
   let verify_promise = Verify.verify
 
   let verify max_proofs_verified statement key proofs =
@@ -248,6 +251,8 @@ module Make_str (_ : Wire_types.Concrete) = struct
         ; branches = Verification_key.Max_branches.n
         ; feature_flags =
             Plonk_types.(Features.to_full ~or_:Opt.Flag.( ||| ) feature_flags)
+        ; num_chunks = 1
+        ; zk_rows = 3
         }
 
     module Proof = struct
@@ -291,7 +296,7 @@ module Make_str (_ : Wire_types.Concrete) = struct
                     { constraints = 0 }
                 }
               in
-              Verify.Instance.T (max_proofs_verified, m, vk, x, p) )
+              Verify.Instance.T (max_proofs_verified, m, None, vk, x, p) )
           |> Verify.verify_heterogenous )
 
     let verify ~typ ts = verify_promise ~typ ts |> Promise.to_deferred
@@ -1146,7 +1151,7 @@ module Make_str (_ : Wire_types.Concrete) = struct
           end in
           let proofs_verifieds = Vector.singleton 2 in
           let (T inner_step_data as step_data) =
-            Step_branch_data.create ~index:0 ~feature_flags
+            Step_branch_data.create ~index:0 ~feature_flags ~num_chunks:1
               ~actual_feature_flags ~max_proofs_verified:Max_proofs_verified.n
               ~branches:Branches.n ~self ~public_input:(Input typ)
               ~auxiliary_typ:typ A.to_field_elements A_value.to_field_elements
@@ -1305,8 +1310,7 @@ module Make_str (_ : Wire_types.Concrete) = struct
               let wrap =
                 let wrap_vk = Lazy.force wrap_vk in
                 let%bind.Promise proof, (), (), _ =
-                  step ~zk_rows:pairing_vk.zk_rows ~proof_cache:None
-                    ~maxes:(module Maxes)
+                  step ~proof_cache:None ~maxes:(module Maxes)
                 in
                 let proof =
                   { proof with
@@ -1827,6 +1831,8 @@ module Make_str (_ : Wire_types.Concrete) = struct
             ; wrap_vk = Lazy.map wrap_vk ~f:Verification_key.index
             ; wrap_domains
             ; step_domains
+            ; num_chunks = 1
+            ; zk_rows = 3
             }
           in
           Types_map.add_exn self data ;

src/lib/pickles/pickles_intf.mli

@@ -238,8 +238,12 @@ module type S = sig
       }
   end
 
+  type chunking_data = Verify.Instance.chunking_data =
+    { num_chunks : int; domain_size : int; zk_rows : int }
+
   val verify_promise :
-       (module Nat.Intf with type n = 'n)
+       ?chunking_data:chunking_data
+    -> (module Nat.Intf with type n = 'n)
     -> (module Statement_value_intf with type t = 'a)
     -> Verification_key.t
     -> ('a * ('n, 'n) Proof.t) list

src/lib/pickles/step.ml

@@ -59,7 +59,7 @@ struct
           with type length = Max_proofs_verified.n
            and type ns = max_local_max_proof_verifieds )
       ~(prevs_length : (prev_vars, prevs_length) Length.t) ~self ~step_domains
-      ~zk_rows ~feature_flags ~self_dlog_plonk_index
+      ~feature_flags ~self_dlog_plonk_index
       ~(public_input :
          ( var
          , value
@@ -204,7 +204,7 @@ struct
           Plonk_checks.scalars_env
             (module Env_bool)
             (module Env_field)
-            ~srs_length_log2:Common.Max_degree.step_log2 ~zk_rows:3
+            ~srs_length_log2:Common.Max_degree.step_log2 ~zk_rows:data.zk_rows
             ~endo:Endo.Step_inner_curve.base ~mds:Tick_field_sponge.params.mds
             ~field_of_hex:(fun s ->
               Kimchi_pasta.Pasta.Bigint256.of_hex_string s
@@ -237,7 +237,7 @@ struct
         Wrap_deferred_values.expand_deferred ~evals:t.prev_evals
           ~old_bulletproof_challenges:
             statement.messages_for_next_step_proof.old_bulletproof_challenges
-          ~zk_rows ~proof_state:statement.proof_state
+          ~zk_rows:data.zk_rows ~proof_state:statement.proof_state
       in
       let prev_statement_with_hashes :
           ( _

src/lib/pickles/step.mli

@@ -31,7 +31,6 @@ module Make
     -> prevs_length:('prev_vars, 'prevs_length) Pickles_types.Hlist.Length.t
     -> self:('a, 'b, 'c, 'd) Tag.t
     -> step_domains:(Import.Domains.t, 'self_branches) Pickles_types.Vector.t
-    -> zk_rows:int
     -> feature_flags:Opt.Flag.t Plonk_types.Features.Full.t
     -> self_dlog_plonk_index:
          Backend.Tick.Inner_curve.Affine.t array

src/lib/pickles/step_branch_data.ml

@@ -73,7 +73,7 @@ let create
     (type branches max_proofs_verified var value a_var a_value ret_var ret_value)
     ~index ~(self : (var, value, max_proofs_verified, branches) Tag.t)
     ~wrap_domains ~(feature_flags : Opt.Flag.t Plonk_types.Features.Full.t)
-    ~(actual_feature_flags : bool Plonk_types.Features.t)
+    ~num_chunks ~(actual_feature_flags : bool Plonk_types.Features.t)
     ~(max_proofs_verified : max_proofs_verified Nat.t)
     ~(proofs_verifieds : (int, branches) Vector.t) ~(branches : branches Nat.t)
     ~(public_input :
@@ -141,6 +141,14 @@ let create
         ; wrap_domains
         ; step_domains
         ; feature_flags
+        ; num_chunks
+        ; zk_rows =
+            ( match num_chunks with
+            | 1 ->
+                3
+            | num_chunks ->
+                let permuts = 7 in
+                ((2 * (permuts + 1) * num_chunks) - 1 + permuts) / permuts )
         }
       ~public_input ~auxiliary_typ ~self_branches:branches ~proofs_verified
       ~local_signature:widths ~local_signature_length ~local_branches:heights

src/lib/pickles/step_branch_data.mli

@@ -84,6 +84,7 @@ val create :
   -> self:('var, 'value, 'max_proofs_verified, 'branches) Tag.t
   -> wrap_domains:Import.Domains.t
   -> feature_flags:Opt.Flag.t Plonk_types.Features.Full.t
+  -> num_chunks:int
   -> actual_feature_flags:bool Plonk_types.Features.t
   -> max_proofs_verified:'max_proofs_verified Pickles_types.Nat.t
   -> proofs_verifieds:(int, 'branches) Pickles_types.Vector.t

src/lib/pickles/step_main.ml

@@ -42,7 +42,8 @@ let verify_one ~srs
         in
         (* TODO: Refactor args into an "unfinalized proof" struct *)
         finalize_other_proof d.max_proofs_verified ~step_domains:d.step_domains
-          ~sponge ~prev_challenges deferred_values prev_proof_evals )
+          ~zk_rows:d.zk_rows ~sponge ~prev_challenges deferred_values
+          prev_proof_evals )
   in
   let branch_data = deferred_values.branch_data in
   let sponge_after_index, hash_messages_for_next_step_proof =
@@ -193,22 +194,23 @@ let step_main :
         Per_proof_witness.Constant.No_app_state.t )
       Typ.t
   end in
-  let feature_flags (d : _ Tag.t) =
-    if Type_equal.Id.same self.id d.id then basic.feature_flags
-    else Types_map.feature_flags d
+  let feature_flags_and_num_chunks (d : _ Tag.t) =
+    if Type_equal.Id.same self.id d.id then
+      (basic.feature_flags, basic.num_chunks)
+    else (Types_map.feature_flags d, Types_map.num_chunks d)
   in
-  let feature_flags =
+  let feature_flags_and_num_chunks =
     let rec go :
         type pvars pvals ns1 ns2 br.
            (pvars, pvals, ns1, ns2) H4.T(Tag).t
         -> (pvars, br) Length.t
-        -> (Opt.Flag.t Plonk_types.Features.Full.t, br) Vector.t =
+        -> (Opt.Flag.t Plonk_types.Features.Full.t * int, br) Vector.t =
      fun ds ld ->
       match[@warning "-4"] (ds, ld) with
       | [], Z ->
           []
       | d :: ds, S ld ->
-          feature_flags d :: go ds ld
+          feature_flags_and_num_chunks d :: go ds ld
       | [], _ ->
           .
       | _ :: _, _ ->
@@ -225,10 +227,12 @@ let step_main :
         -> (pvars, br) Length.t
         -> (ns1, br) Length.t
         -> (ns2, br) Length.t
-        -> (Opt.Flag.t Plonk_types.Features.Full.t, br) Vector.t
+        -> (Opt.Flag.t Plonk_types.Features.Full.t * int, br) Vector.t
         -> (pvars, pvals, ns1, ns2) H4.T(Typ_with_max_proofs_verified).t =
-     fun ds ns1 ns2 ld ln1 ln2 feature_flagss ->
-      match[@warning "-4"] (ds, ns1, ns2, ld, ln1, ln2, feature_flagss) with
+     fun ds ns1 ns2 ld ln1 ln2 feature_flags_and_num_chunkss ->
+      match[@warning "-4"]
+        (ds, ns1, ns2, ld, ln1, ln2, feature_flags_and_num_chunkss)
+      with
       | [], [], [], Z, Z, Z, [] ->
           []
       | ( _d :: ds
@@ -237,16 +241,18 @@ let step_main :
         , S ld
         , S ln1
         , S ln2
-        , feature_flags :: feature_flagss ) ->
-          let t = Per_proof_witness.typ Typ.unit n1 ~feature_flags in
-          t :: join ds ns1 ns2 ld ln1 ln2 feature_flagss
+        , (feature_flags, num_chunks) :: feature_flags_and_num_chunkss ) ->
+          let t =
+            Per_proof_witness.typ Typ.unit n1 ~feature_flags ~num_chunks
+          in
+          t :: join ds ns1 ns2 ld ln1 ln2 feature_flags_and_num_chunkss
       | [], _, _, _, _, _, _ ->
           .
       | _ :: _, _, _, _, _, _, _ ->
           .
     in
     join rule.prevs local_signature local_branches proofs_verified
-      local_signature_length local_branches_length feature_flags
+      local_signature_length local_branches_length feature_flags_and_num_chunks
   in
   let module Prev_typ =
     H4.Typ (Impls.Step) (Typ_with_max_proofs_verified)
@@ -353,7 +359,7 @@ let step_main :
                (Vector.map
                   ~f:(fun _feature_flags ->
                     Unfinalized.typ ~wrap_rounds:Backend.Tock.Rounds.n )
-                  feature_flags ) )
+                  feature_flags_and_num_chunks ) )
             ~request:(fun () -> Req.Unfinalized_proofs)
         and messages_for_next_wrap_proof =
           exists (Vector.typ Digest.typ Max_proofs_verified.n)
@@ -485,6 +491,8 @@ let step_main :
                     ; step_domains = `Known basic.step_domains
                     ; wrap_key = dlog_plonk_index
                     ; feature_flags = basic.feature_flags
+                    ; num_chunks = basic.num_chunks
+                    ; zk_rows = basic.zk_rows
                     }
                   in
                   let module M =

src/lib/pickles/step_verifier.ml

@@ -854,7 +854,7 @@ struct
   let finalize_other_proof (type b branches)
       (module Proofs_verified : Nat.Add.Intf with type n = b)
       ~(step_domains :
-         [ `Known of (Domains.t, branches) Vector.t | `Side_loaded ] )
+         [ `Known of (Domains.t, branches) Vector.t | `Side_loaded ] ) ~zk_rows
       ~(* TODO: Add "actual proofs verified" so that proofs don't
           carry around dummy "old bulletproof challenges" *)
        sponge ~(prev_challenges : (_, b) Vector.t)
@@ -1016,7 +1016,7 @@ struct
           Plonk_checks.scalars_env
             (module Env_bool)
             (module Env_field)
-            ~srs_length_log2:Common.Max_degree.step_log2 ~zk_rows:(* TODO *) 3
+            ~srs_length_log2:Common.Max_degree.step_log2 ~zk_rows
             ~endo:(Impl.Field.constant Endo.Step_inner_curve.base)
             ~mds:sponge_params.mds
             ~field_of_hex:(fun s ->