Skip to content

Commit b2c1745

Browse files
authored
Merge pull request #2143 from o1-labs/project/wasm-mode
2 parents 9a491fb + ee7e189 commit b2c1745

File tree

10 files changed

+116
-11
lines changed

10 files changed

+116
-11
lines changed

CHANGELOG.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
1919

2020
### Added
2121

22+
- Lazy mode for prover index computation. https://github.com/o1-labs/o1js/pull/2143
2223
- Added `ZkProgram.analyzeSingleMethod(methodName: string)` to analyze a single method of a ZkProgram. https://github.com/o1-labs/o1js/pull/2217
2324
- This is an addition to `ZkProgram.analyzeMethods()` which analyzes all methods of a ZkProgram by executing them.
2425
- Now only a single method is analyzed at a time.
@@ -28,8 +29,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
2829
### Added
2930

3031
- API support for circuit chunking. https://github.com/o1-labs/o1js/pull/1905
31-
- still requires memory optimizations to be fully functional, and
32-
- proof-systems version still needs to be updated to include [this commit](https://github.com/o1-labs/proof-systems/pull/3222/commits/8c37c293f8159eed3676964ba47fc5dc0ae6ea1e)
32+
- work in progress, still requires memory optimizations to be fully functional
3333
- Dynamic array provable type. https://github.com/o1-labs/o1js/pull/1848
3434

3535
## [2.5.0](https://github.com/o1-labs/o1js/compare/6ff7f8470a...4e23a60)

src/bindings.d.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -385,7 +385,7 @@ declare const Snarky: {
385385
/**
386386
* Generates a proving key and a verification key for the provable function `main`
387387
*/
388-
compile(main: Snarky.Main, publicInputSize: number): Snarky.Keypair;
388+
compile(main: Snarky.Main, publicInputSize: number, lazyMode: boolean): Snarky.Keypair;
389389

390390
/**
391391
* Proves a statement using the private input, public input and the keypair of the circuit.
@@ -703,6 +703,7 @@ declare const Pickles: {
703703
storable?: Pickles.Cache;
704704
overrideWrapDomain?: 0 | 1 | 2;
705705
numChunks?: number;
706+
lazyMode?: boolean;
706707
}
707708
) => {
708709
provers: MlArray<Pickles.Prover>;

src/bindings/js/web/worker-spec.js

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ function workerSpec(wasm) {
2222
undefined /* number */,
2323
// srs
2424
wasm.WasmFpSrs,
25+
// lazy_mode
26+
undefined /* boolean */,
2527
],
2628
res: wasm.WasmPastaFpPlonkIndex,
2729
},
@@ -39,6 +41,8 @@ function workerSpec(wasm) {
3941
undefined /* number */,
4042
// srs
4143
wasm.WasmFqSrs,
44+
// lazy_mode
45+
undefined /* boolean */,
4246
],
4347
res: wasm.WasmPastaFqPlonkIndex,
4448
},

src/bindings/ocaml/lib/pickles_bindings.ml

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -618,7 +618,8 @@ let pickles_compile (choices : pickles_rule_js array)
618618
; publicOutputSize : int Js.prop
619619
; storable : Cache.js_storable Js.optdef_prop
620620
; overrideWrapDomain : int Js.optdef_prop
621-
; numChunks : int Js.optdef_prop >
621+
; numChunks : int Js.optdef_prop
622+
; lazyMode : bool Js.optdef_prop >
622623
Js.t ) =
623624
(* translate number of branches and recursively verified proofs from JS *)
624625
let branches = Array.length choices in
@@ -640,6 +641,7 @@ let pickles_compile (choices : pickles_rule_js array)
640641
let num_chunks =
641642
Js.Optdef.get config##.numChunks (fun () -> 1)
642643
in
644+
let lazy_mode = Js.Optdef.get config##.lazyMode (fun () -> false) in
643645
let (Choices choices) =
644646
Choices.of_js ~public_input_size ~public_output_size choices
645647
in
@@ -662,7 +664,7 @@ let pickles_compile (choices : pickles_rule_js array)
662664
, public_input_typ public_output_size ) )
663665
~auxiliary_typ:Typ.unit
664666
~max_proofs_verified:(module Max_proofs_verified)
665-
~name ~num_chunks ~choices ()
667+
~name ~num_chunks ~lazy_mode ~choices ()
666668
in
667669

668670
(* translate returned prover and verify functions to JS *)

src/bindings/ocaml/lib/pickles_bindings.mli

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,8 @@ val pickles :
6666
; publicOutputSize : int Js.prop
6767
; storable : Cache.js_storable Js.optdef_prop
6868
; overrideWrapDomain : int Js.optdef_prop
69-
; numChunks : int Js.optdef_prop >
69+
; numChunks : int Js.optdef_prop
70+
; lazyMode : bool Js.optdef_prop >
7071
Js.t
7172
-> < getVerificationKey :
7273
( unit

src/bindings/ocaml/lib/snarky_bindings.ml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -384,11 +384,11 @@ module Circuit = struct
384384
main'
385385
end
386386

387-
let compile main public_input_size =
387+
let compile main public_input_size lazy_mode =
388388
let input_typ = typ public_input_size in
389389
let return_typ = Impl.Typ.unit in
390390
let cs = Impl.constraint_system ~input_typ ~return_typ (Main.of_js main) in
391-
Impl.Keypair.generate ~prev_challenges:0 cs
391+
Impl.Keypair.generate ~lazy_mode ~prev_challenges:0 cs
392392

393393
let prove main public_input_size public_input keypair =
394394
let pk = Impl.Keypair.pk keypair in

src/bindings/ocaml/lib/snarky_bindings.mli

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -249,7 +249,8 @@ val snarky :
249249
Js.t
250250
Js.readonly_prop
251251
; circuit :
252-
< compile : ((Field.t array -> unit) -> int -> Impl.Keypair.t) Js.meth
252+
< compile :
253+
((Field.t array -> unit) -> int -> bool -> Impl.Keypair.t) Js.meth
253254
; keypair :
254255
< getConstraintSystemJSON : (Impl.Keypair.t -> 'a) Js.meth
255256
; getVerificationKey :

src/lib/proof-system/circuit.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,13 +22,13 @@ class Circuit {
2222
* const keypair = await MyCircuit.generateKeypair();
2323
* ```
2424
*/
25-
static async generateKeypair() {
25+
static async generateKeypair(lazyMode: boolean = false) {
2626
let main = mainFromCircuitData(this._main);
2727
let publicInputSize = this._main.publicInputType.sizeInFields();
2828
await initializeBindings();
2929
return prettifyStacktracePromise(
3030
withThreadPool(async () => {
31-
let keypair = Snarky.circuit.compile(main, publicInputSize);
31+
let keypair = Snarky.circuit.compile(main, publicInputSize, lazyMode);
3232
return new Keypair(keypair);
3333
})
3434
);
Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
import { Field } from '../provable/field.js';
2+
import { Cache } from '../proof-system/cache.js';
3+
import { ZkProgram } from '../proof-system/zkprogram.js';
4+
import { Spec, fieldWithRng } from '../testing/equivalent.js';
5+
import { Random } from '../testing/property.js';
6+
import { assert } from '../provable/gadgets/common.js';
7+
import { Gadgets } from '../provable/gadgets/gadgets.js';
8+
import { wasm } from '../../bindings.js';
9+
import { spawn } from 'child_process';
10+
import { fileURLToPath } from 'url';
11+
import { dirname, join, basename } from 'path';
12+
13+
// Path resolution for subprocess execution
14+
const __filename = fileURLToPath(import.meta.url);
15+
const __dirname = dirname(__filename);
16+
const scriptPath = join(__dirname, basename(__filename));
17+
18+
function getMemory() {
19+
return {
20+
wasm: (wasm as any).__wasm.memory.buffer.byteLength / (1024 * 1024),
21+
js: process.memoryUsage().heapTotal / (1024 * 1024),
22+
};
23+
}
24+
25+
let uint = (n: number | bigint): Spec<bigint, Field> => {
26+
return fieldWithRng(Random.bignat((1n << BigInt(n)) - 1n));
27+
};
28+
29+
// Dummy circuit
30+
let LazyMode = ZkProgram({
31+
name: 'lazy-mode',
32+
methods: {
33+
baseCase: {
34+
privateInputs: [Field],
35+
async method(v: Field) {
36+
for (let i = 0; i < 1 << 15; i++) {
37+
let w = v.add(new Field(i));
38+
Gadgets.rangeCheck64(w);
39+
}
40+
},
41+
},
42+
},
43+
});
44+
45+
export async function testLazyMode(lazyMode: boolean) {
46+
console.log(`(${lazyMode ? 'Lazy' : 'Eager'}) Memory before compilation`, getMemory());
47+
48+
await LazyMode.compile({
49+
lazyMode,
50+
cache: Cache.None,
51+
forceRecompile: true,
52+
});
53+
54+
console.log(`(${lazyMode ? 'Lazy' : 'Eager'}) Memory after compilation`, getMemory());
55+
console.log(' ');
56+
console.log(`(${lazyMode ? 'Lazy' : 'Eager'}) Memory before proof`, getMemory());
57+
58+
let { proof } = await LazyMode.baseCase(new Field(1n));
59+
60+
console.log(`(${lazyMode ? 'Lazy' : 'Eager'}) Memory after proof`, getMemory());
61+
console.log(' ');
62+
63+
let isValid = await LazyMode.verify(proof);
64+
65+
console.log(`(${lazyMode ? 'Lazy' : 'Eager'}) Memory after verify`, getMemory());
66+
67+
assert(isValid);
68+
}
69+
70+
// If launched as subprocess: run the actual test
71+
if (process.argv[2] === 'true' || process.argv[2] === 'false') {
72+
const lazyMode = process.argv[2] === 'true';
73+
await testLazyMode(lazyMode);
74+
process.exit(0);
75+
}
76+
77+
// Parent process logic: spawn two subprocesses
78+
function runSubprocess(lazyMode: boolean): Promise<void> {
79+
return new Promise((resolve, reject) => {
80+
const child = spawn('node', [scriptPath, String(lazyMode)], {
81+
stdio: 'inherit',
82+
});
83+
child.on('exit', (code) => {
84+
console.log(`(Parent) Process lazyMode=${lazyMode} exited with code ${code}`);
85+
code === 0 ? resolve() : reject(new Error(`Test failed for lazyMode=${lazyMode}`));
86+
});
87+
});
88+
}
89+
90+
await Promise.all([runSubprocess(true), runSubprocess(false)]);

src/lib/proof-system/zkprogram.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,7 @@ function ZkProgram<
245245
proofsEnabled?: boolean;
246246
withRuntimeTables?: boolean;
247247
numChunks?: number;
248+
lazyMode?: boolean;
248249
}) => Promise<{
249250
verificationKey: { data: string; hash: Field };
250251
}>;
@@ -382,6 +383,7 @@ function ZkProgram<
382383
forceRecompile = false,
383384
proofsEnabled = undefined as boolean | undefined,
384385
withRuntimeTables = false,
386+
lazyMode = false,
385387
} = {}) {
386388
doProving = proofsEnabled ?? doProving;
387389

@@ -405,6 +407,7 @@ function ZkProgram<
405407
numChunks: config.numChunks,
406408
state: programState,
407409
withRuntimeTables,
410+
lazyMode,
408411
});
409412

410413
compileOutput = { provers, verify, maxProofsVerified };
@@ -712,6 +715,7 @@ async function compileProgram({
712715
numChunks,
713716
state,
714717
withRuntimeTables,
718+
lazyMode,
715719
}: {
716720
publicInputType: Provable<any>;
717721
publicOutputType: Provable<any>;
@@ -726,6 +730,7 @@ async function compileProgram({
726730
numChunks?: number;
727731
state?: ReturnType<typeof createProgramState>;
728732
withRuntimeTables?: boolean;
733+
lazyMode?: boolean;
729734
}) {
730735
await initializeBindings();
731736
if (methodIntfs.length === 0)
@@ -782,6 +787,7 @@ If you are using a SmartContract, make sure you are using the @method decorator.
782787
storable: picklesCache,
783788
overrideWrapDomain,
784789
numChunks: numChunks ?? 1,
790+
lazyMode: lazyMode ?? false,
785791
});
786792
let { getVerificationKey, provers, verify, tag } = result;
787793
CompiledTag.store(proofSystemTag, tag);

0 commit comments

Comments
 (0)