Merge pull request #1905 from o1-labs/querolita/chunking-mvp-bug

API chunking support

Author: querolita

CHANGELOG.md

@@ -17,6 +17,13 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased](https://github.com/o1-labs/o1js/compare/4e23a60...HEAD)
 
+### Added
+
+- [PR! 1905](https://github.com/o1-labs/o1js/pull/1905) API support for circuit chunking
+  - still requires memory optimizations to be fully functional, and
+  - proof-systems version still needs to be updated to include [this commit](https://github.com/o1-labs/proof-systems/pull/3222/commits/8c37c293f8159eed3676964ba47fc5dc0ae6ea1e) 
+    - that fixed the zero knowledge rows mismatch across Kimchi WASM bindings
+
 ## [2.5.0](https://github.com/o1-labs/o1js/compare/6ff7f8470a...4e23a60)
 
 ### Fixed

src/bindings/ocaml/lib/pickles_bindings.ml

@@ -617,7 +617,8 @@ let pickles_compile (choices : pickles_rule_js array)
       < publicInputSize : int Js.prop
       ; publicOutputSize : int Js.prop
       ; storable : Cache.js_storable Js.optdef_prop
-      ; overrideWrapDomain : int Js.optdef_prop >
+      ; overrideWrapDomain : int Js.optdef_prop 
+      ; numChunks : int Js.optdef_prop >
       Js.t ) =
   (* translate number of branches and recursively verified proofs from JS *)
   let branches = Array.length choices in
@@ -636,6 +637,9 @@ let pickles_compile (choices : pickles_rule_js array)
     Js.Optdef.to_option config##.overrideWrapDomain
     |> Option.map ~f:Pickles_base.Proofs_verified.of_int_exn
   in
+  let num_chunks =
+    Js.Optdef.get config##.numChunks (fun () -> 1)
+  in
   let (Choices choices) =
     Choices.of_js ~public_input_size ~public_output_size choices
   in
@@ -658,7 +662,7 @@ let pickles_compile (choices : pickles_rule_js array)
            , public_input_typ public_output_size ) )
       ~auxiliary_typ:Typ.unit
       ~max_proofs_verified:(module Max_proofs_verified)
-      ~name ~choices ()
+      ~name ~num_chunks ~choices ()
   in
 
   (* translate returned prover and verify functions to JS *)

src/bindings/ocaml/lib/pickles_bindings.mli

@@ -65,7 +65,8 @@ val pickles :
        -> < publicInputSize : int Js.prop
           ; publicOutputSize : int Js.prop
           ; storable : Cache.js_storable Js.optdef_prop
-          ; overrideWrapDomain : int Js.optdef_prop >
+          ; overrideWrapDomain : int Js.optdef_prop 
+          ; numChunks : int Js.optdef_prop >
           Js.t
        -> < getVerificationKey :
               (   unit

src/examples/zkprogram/program-with-chunking.ts

@@ -0,0 +1,35 @@
+import { Field, Cache, Gadgets, ZkProgram } from 'o1js';
+
+let MyProgram = ZkProgram({
+  chunks: 1,
+  overrideWrapDomain: 0,
+  name: 'example-with-chunking',
+  publicOutput: Field,
+
+  methods: {
+    baseCase: {
+      privateInputs: [Field],
+      async method(input: Field) {
+        for (let i = 0; i < 1 << 16; i++) {
+          Gadgets.rangeCheck64(Field(input).add(Field(i)));
+        }
+        // The above generates 2^16+2^15 rows which needs to be split into 2 chunks
+        return {
+          publicOutput: Field(0),
+        };
+      },
+    },
+  },
+});
+
+console.log(await MyProgram.analyzeMethods());
+
+console.log('compiling MyProgram...');
+await MyProgram.compile({ cache: Cache.None });
+
+console.log('proving base case...');
+let { proof } = await MyProgram.baseCase(Field(0));
+
+console.log('verify...');
+let ok = await MyProgram.verify(proof);
+console.log('ok?', ok);

src/lib/proof-system/zkprogram.ts

@@ -233,6 +233,7 @@ function ZkProgram<
       [I in keyof Config['methods']]: InferMethodType<Config>[I];
     };
     overrideWrapDomain?: 0 | 1 | 2;
+    numChunks?: number;
   }
 ): {
   name: string;
@@ -243,6 +244,7 @@ function ZkProgram<
     forceRecompile?: boolean;
     proofsEnabled?: boolean;
     withRuntimeTables?: boolean;
+    numChunks?: number;
   }) => Promise<{
     verificationKey: { data: string; hash: Field };
   }>;
@@ -386,6 +388,7 @@ function ZkProgram<
         cache,
         forceRecompile,
         overrideWrapDomain: config.overrideWrapDomain,
+        numChunks: config.numChunks,
         state: programState,
         withRuntimeTables,
       });
@@ -690,6 +693,7 @@ async function compileProgram({
   cache,
   forceRecompile,
   overrideWrapDomain,
+  numChunks,
   state,
   withRuntimeTables,
 }: {
@@ -703,6 +707,7 @@ async function compileProgram({
   cache: Cache;
   forceRecompile: boolean;
   overrideWrapDomain?: 0 | 1 | 2;
+  numChunks?: number;
   state?: ReturnType<typeof createProgramState>;
   withRuntimeTables?: boolean;
 }) {
@@ -760,6 +765,7 @@ If you are using a SmartContract, make sure you are using the @method decorator.
           publicOutputSize: publicOutputType.sizeInFields(),
           storable: picklesCache,
           overrideWrapDomain,
+          numChunks: numChunks ?? 1,
         });
         let { getVerificationKey, provers, verify, tag } = result;
         CompiledTag.store(proofSystemTag, tag);

src/snarky.d.ts

@@ -702,6 +702,7 @@ declare const Pickles: {
       publicOutputSize: number;
       storable?: Pickles.Cache;
       overrideWrapDomain?: 0 | 1 | 2;
+      numChunks?: number;
     }
   ) => {
     provers: MlArray<Pickles.Prover>;