pass chunks to compile promise in pickles bindings

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

@@ -1,4 +1,4 @@
-import { Field, ZkProgram, verify } from 'o1js';
+import { Field, Cache, Gadgets, ZkProgram } from 'o1js';
 
 let MyProgram = ZkProgram({
   chunks: 1,
@@ -8,20 +8,28 @@ let MyProgram = ZkProgram({
 
   methods: {
     baseCase: {
-      privateInputs: [],
-      async method() {
-        return Field(0);
+      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...');
-let { verificationKey } = await MyProgram.compile();
+await MyProgram.compile({ cache: Cache.None });
 
 console.log('proving base case...');
-let proof = await MyProgram.baseCase();
+let { proof } = await MyProgram.baseCase(Field(0));
 
 console.log('verify...');
-let ok = await verify(proof.toJSON(), verificationKey);
+let ok = await MyProgram.verify(proof);
 console.log('ok?', ok);

src/lib/proof-system/zkprogram.ts

@@ -233,7 +233,7 @@ function ZkProgram<
       [I in keyof Config['methods']]: InferMethodType<Config>[I];
     };
     overrideWrapDomain?: 0 | 1 | 2;
-    chunks?: number;
+    numChunks?: number;
   }
 ): {
   name: string;
@@ -244,6 +244,7 @@ function ZkProgram<
     forceRecompile?: boolean;
     proofsEnabled?: boolean;
     withRuntimeTables?: boolean;
+    numChunks?: number;
   }) => Promise<{
     verificationKey: { data: string; hash: Field };
   }>;
@@ -387,9 +388,9 @@ function ZkProgram<
         cache,
         forceRecompile,
         overrideWrapDomain: config.overrideWrapDomain,
+        numChunks: config.numChunks,
         state: programState,
         withRuntimeTables,
-        chunks: config.chunks,
       });
 
       compileOutput = { provers, verify, maxProofsVerified };
@@ -692,9 +693,9 @@ async function compileProgram({
   cache,
   forceRecompile,
   overrideWrapDomain,
+  numChunks,
   state,
   withRuntimeTables,
-  chunks,
 }: {
   publicInputType: Provable<any>;
   publicOutputType: Provable<any>;
@@ -706,9 +707,9 @@ async function compileProgram({
   cache: Cache;
   forceRecompile: boolean;
   overrideWrapDomain?: 0 | 1 | 2;
+  numChunks?: number;
   state?: ReturnType<typeof createProgramState>;
   withRuntimeTables?: boolean;
-  chunks?: number;
 }) {
   await initializeBindings();
   if (methodIntfs.length === 0)
@@ -753,36 +754,35 @@ If you are using a SmartContract, make sure you are using the @method decorator.
     MlBool(cache.canWrite),
   ];
 
-  let { verificationKey, provers, verify, tag } =
-    await prettifyStacktracePromise(
-      withThreadPool(async () => {
-        let result: ReturnType<typeof Pickles.compile>;
-        let id = snarkContext.enter({ inCompile: true });
-        setSrsCache(cache);
-        try {
-          result = Pickles.compile(MlArray.to(rules), {
-            publicInputSize: publicInputType.sizeInFields(),
-            publicOutputSize: publicOutputType.sizeInFields(),
-            storable: picklesCache,
-            overrideWrapDomain,
-            chunks: chunks ?? 1,
-          });
-          let { getVerificationKey, provers, verify, tag } = result;
-          CompiledTag.store(proofSystemTag, tag);
-          let [, data, hash] = await getVerificationKey();
-          let verificationKey = { data, hash: Field(hash) };
-          return {
-            verificationKey,
-            provers: MlArray.from(provers),
-            verify,
-            tag,
-          };
-        } finally {
-          snarkContext.leave(id);
-          unsetSrsCache();
-        }
-      })
-    );
+  let { verificationKey, provers, verify, tag } = await prettifyStacktracePromise(
+    withThreadPool(async () => {
+      let result: ReturnType<typeof Pickles.compile>;
+      let id = snarkContext.enter({ inCompile: true });
+      setSrsCache(cache);
+      try {
+        result = Pickles.compile(MlArray.to(rules), {
+          publicInputSize: publicInputType.sizeInFields(),
+          publicOutputSize: publicOutputType.sizeInFields(),
+          storable: picklesCache,
+          overrideWrapDomain,
+          numChunks: numChunks ?? 1,
+        });
+        let { getVerificationKey, provers, verify, tag } = result;
+        CompiledTag.store(proofSystemTag, tag);
+        let [, data, hash] = await getVerificationKey();
+        let verificationKey = { data, hash: Field(hash) };
+        return {
+          verificationKey,
+          provers: MlArray.from(provers),
+          verify,
+          tag,
+        };
+      } finally {
+        snarkContext.leave(id);
+        unsetSrsCache();
+      }
+    })
+  );
   // wrap provers
   let wrappedProvers = provers.map(
     (prover): Pickles.Prover =>

src/snarky.d.ts

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