Merge pull request #2458 from o1-labs/shigoto/performance-regression-migrate-tests

Update ZkProgram tests to utilize the new performance (compilation, proving, verification) regression framework

Author: Shigoto-dev19

src/examples/zkprogram/gadgets.ts

@@ -1,4 +1,5 @@
-import { Field, Provable, Gadgets, ZkProgram } from 'o1js';
+import { Field, Gadgets, Provable, ZkProgram } from 'o1js';
+import { Performance } from '../../lib/testing/perf-regression.js';
 
 let cs = await Provable.constraintSystem(() => {
   let f = Provable.witness(Field, () => 12);
@@ -54,25 +55,40 @@ const BitwiseProver = ZkProgram({
   },
 });
 
-console.log('compiling..');
+const csBitwise = await BitwiseProver.analyzeMethods();
+const perfBitwise = Performance.create(BitwiseProver.name, csBitwise);
 
-console.time('compile');
+console.log('\ncompiling..');
+
+perfBitwise.start('compile');
 await BitwiseProver.compile();
-console.timeEnd('compile');
+perfBitwise.end();
 
-console.log('proving..');
+console.log('\nproving..');
 
-console.time('rotation prove');
+perfBitwise.start('prove', 'rot');
 let { proof: rotProof } = await BitwiseProver.rot();
-console.timeEnd('rotation prove');
-if (!(await BitwiseProver.verify(rotProof))) throw Error('rot: Invalid proof');
+perfBitwise.end();
+
+perfBitwise.start('verify', 'rot');
+const isValidRot = await BitwiseProver.verify(rotProof);
+perfBitwise.end();
+if (!isValidRot) throw Error('rot: Invalid proof');
 
-console.time('xor prove');
+perfBitwise.start('prove', 'xor');
 let { proof: xorProof } = await BitwiseProver.xor();
-console.timeEnd('xor prove');
-if (!(await BitwiseProver.verify(xorProof))) throw Error('xor: Invalid proof');
+perfBitwise.end();
 
-console.time('and prove');
+perfBitwise.start('verify', 'xor');
+const isValidXor = await BitwiseProver.verify(xorProof);
+perfBitwise.end();
+if (!isValidXor) throw Error('xor: Invalid proof');
+
+perfBitwise.start('prove', 'and');
 let { proof: andProof } = await BitwiseProver.and();
-console.timeEnd('and prove');
-if (!(await BitwiseProver.verify(andProof))) throw Error('and: Invalid proof');
+perfBitwise.end();
+
+perfBitwise.start('verify', 'and');
+const isValidAnd = await BitwiseProver.verify(andProof);
+perfBitwise.end();
+if (!isValidAnd) throw Error('and: Invalid proof');

src/examples/zkprogram/hash-chain.ts

@@ -5,6 +5,7 @@
  * We implement this as a self-recursive ZkProgram, using `proveRecursivelyIf()`
  */
 import { assert, Bool, Experimental, Field, Poseidon, Provable, Struct, ZkProgram } from 'o1js';
+import { Performance } from '../../lib/testing/perf-regression.js';
 
 const HASHES_PER_PROOF = 30;
 
@@ -47,14 +48,24 @@ const hashChain = ZkProgram({
 });
 let hashChainRecursive = Experimental.Recursive(hashChain);
 
+const cs = await hashChain.analyzeMethods();
+const perfHashChain = Performance.create(hashChain.name, cs);
+
+perfHashChain.start('compile');
 await hashChain.compile();
+perfHashChain.end();
 
 let n = 100;
 let x = Field.random();
 
+perfHashChain.start('prove', 'chain');
 let { proof } = await hashChain.chain({ x, n });
+perfHashChain.end();
 
-assert(await hashChain.verify(proof), 'Proof invalid');
+perfHashChain.start('verify', 'chain');
+const isValid = await hashChain.verify(proof);
+perfHashChain.end();
+assert(isValid, 'Proof invalid');
 
 // check that the output is correct
 let z = Array.from({ length: n }, () => 0).reduce((y) => Poseidon.hash([y]), x);

src/examples/zkprogram/mutual-recursion.ts

@@ -1,4 +1,5 @@
-import { ZkProgram, Field, DynamicProof, Proof, VerificationKey, Undefined, verify } from 'o1js';
+import { DynamicProof, Field, Proof, Undefined, VerificationKey, ZkProgram, verify } from 'o1js';
+import { Performance } from '../../lib/testing/perf-regression.js';
 
 /**
  * This example showcases mutual recursion (A -> B -> A) through two circuits that respectively
@@ -51,21 +52,38 @@ const multiply = ZkProgram({
   },
 });
 
-console.log('Compiling circuits...');
+const csAdd = await add.analyzeMethods();
+const csMultiply = await multiply.analyzeMethods();
+
+const perfAdd = Performance.create(add.name, csAdd);
+const perfMultiply = Performance.create(multiply.name, csMultiply);
+
+perfAdd.start('compile');
 const addVk = (await add.compile()).verificationKey;
+perfAdd.end();
+
+perfMultiply.start('compile');
 const multiplyVk = (await multiply.compile()).verificationKey;
+perfMultiply.end();
 
-console.log('Proving basecase');
 const dummyProof = await DynamicMultiplyProof.dummy(undefined, Field(0), 1);
+
+perfAdd.start('prove', 'performAddition');
 const { proof: baseCase } = await add.performAddition(Field(5), dummyProof, multiplyVk);
+perfAdd.end();
 
+perfAdd.start('verify', 'performAddition');
 const validBaseCase = await verify(baseCase, addVk);
+perfAdd.end();
 console.log('ok?', validBaseCase);
 
-console.log('Proving first multiplication');
+perfMultiply.start('prove', 'performMultiplication');
 const { proof: multiply1 } = await multiply.performMultiplication(Field(3), baseCase);
+perfMultiply.end();
 
+perfMultiply.start('verify', 'performMultiplication');
 const validMultiplication = await verify(multiply1, multiplyVk);
+perfMultiply.end();
 console.log('ok?', validMultiplication);
 
 console.log('Proving second (recursive) addition');

src/examples/zkprogram/program-small-big.ts

@@ -1,4 +1,5 @@
 import { Bytes, Field, Hash, Poseidon, UInt8, ZkProgram, verify } from 'o1js';
+import { Performance } from '../../lib/testing/perf-regression.js';
 
 const SmallProgram = ZkProgram({
   name: 'small-program',
@@ -52,24 +53,27 @@ console.log('small program rows: ', csSmall.poseidonHash.rows);
 const csBig = await BigProgram.analyzeMethods();
 console.log('big program rows: ', csBig.combinedHash.rows, '\n');
 
-console.time('compile small');
+const perfSmall = Performance.create(SmallProgram.name, csSmall);
+const perfBig = Performance.create(BigProgram.name, csBig);
+
+perfSmall.start('compile');
 await SmallProgram.compile();
-console.timeEnd('compile small');
+perfSmall.end();
 
-console.time('compile big');
+perfBig.start('compile');
 const { verificationKey: verificationKeyBig } = await BigProgram.compile();
-console.timeEnd('compile big');
+perfBig.end();
 
-console.time('prove small');
+perfSmall.start('prove', 'poseidonHash');
 const proofSmall = await SmallProgram.poseidonHash(Field.random());
-console.timeEnd('prove small');
+perfSmall.end();
 
-console.time('prove big');
+perfBig.start('prove', 'combinedHash');
 const { proof: proofBig } = await BigProgram.combinedHash(proofSmall.proof);
-console.timeEnd('prove big');
+perfBig.end();
 
-console.time('verify big');
+perfBig.start('verify', 'combinedHash');
 await verify(proofBig, verificationKeyBig);
-console.timeEnd('verify big');
+perfBig.end();
 
 console.log('Final Digest: ', proofBig.publicOutput.toHex());

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

@@ -1,4 +1,5 @@
-import { Field, Cache, Gadgets, ZkProgram } from 'o1js';
+import { Cache, Field, Gadgets, ZkProgram } from 'o1js';
+import { Performance } from '../../lib/testing/perf-regression.js';
 
 let MyProgram = ZkProgram({
   chunks: 1,
@@ -22,14 +23,21 @@ let MyProgram = ZkProgram({
   },
 });
 
-console.log(await MyProgram.analyzeMethods());
+const cs = await MyProgram.analyzeMethods();
+const perf = Performance.create(MyProgram.name, cs);
 
-console.log('compiling MyProgram...');
+console.log('MyProgram baseCase method rows: ', cs.baseCase.rows);
+
+perf.start('compile');
 await MyProgram.compile({ cache: Cache.None });
+perf.end();
 
-console.log('proving base case...');
+perf.start('prove', 'baseCase');
 let { proof } = await MyProgram.baseCase(Field(0));
+perf.end();
 
-console.log('verify...');
+perf.start('verify', 'baseCase');
 let ok = await MyProgram.verify(proof);
+perf.end();
+
 console.log('ok?', ok);

src/examples/zkprogram/runtime-table/run.ts

@@ -1,7 +1,13 @@
 import { Field, verify } from 'o1js';
+import { Performance } from '../../../lib/testing/perf-regression.js';
 import { PayrollRuntimeTableZkProgram } from './payroll.js';
 
+const cs = await PayrollRuntimeTableZkProgram.analyzeMethods();
+const perfPayroll = Performance.create(PayrollRuntimeTableZkProgram.name, cs);
+
+perfPayroll.start('compile');
 let { verificationKey } = await PayrollRuntimeTableZkProgram.compile({ withRuntimeTables: true });
+perfPayroll.end();
 
 export const examplePayrollPublicInput = Field(1600_00); // total withheld (cents)
 export const examplePayrollPrivateInputs = [
@@ -13,9 +19,13 @@ export const examplePayrollPrivateInputs = [
   Field(3_000), // Charlie rate
 ] as const;
 
+perfPayroll.start('prove', 'verifyPayroll');
 let { proof } = await PayrollRuntimeTableZkProgram.verifyPayroll(
   examplePayrollPublicInput,
   ...examplePayrollPrivateInputs
 );
+perfPayroll.end();
 
-verify(proof, verificationKey);
+perfPayroll.start('verify', 'verifyPayroll');
+await verify(proof, verificationKey);
+perfPayroll.end();

src/examples/zkprogram/side-loading/run.ts

@@ -1,23 +1,35 @@
 import { Field, MerkleTree, verify } from 'o1js';
+import { Performance } from '../../../lib/testing/perf-regression.js';
 import {
-  sideloadedProgram,
-  mainProgram,
-  SideloadedProgramProof,
-  MerkleTreeWitness,
   MainProgramState,
+  MerkleTreeWitness,
+  SideloadedProgramProof,
+  mainProgram,
+  sideloadedProgram,
 } from './dynamic-keys-merkletree.js';
 
-console.log('Compiling circuits...');
+const csSide = await sideloadedProgram.analyzeMethods();
+const csMain = await mainProgram.analyzeMethods();
+
+const perfSide = Performance.create(sideloadedProgram.name, csSide);
+const perfMain = Performance.create(mainProgram.name, csMain);
+
+perfSide.start('compile');
 const sideVk = (await sideloadedProgram.compile()).verificationKey;
+perfSide.end();
+
+perfMain.start('compile');
 const mainVk = (await mainProgram.compile()).verificationKey;
+perfMain.end();
 
 const tree = new MerkleTree(64);
 
-console.log('Proving deployment of side-loaded key');
+console.log('\nProving deployment of side-loaded key');
 const rootBefore = tree.getRoot();
 tree.setLeaf(1n, sideVk.hash);
 const witness = new MerkleTreeWitness(tree.getWitness(1n));
 
+perfMain.start('prove', 'addSideloadedProgram');
 const { proof: proof1 } = await mainProgram.addSideloadedProgram(
   new MainProgramState({
     treeRoot: rootBefore,
@@ -26,26 +38,35 @@ const { proof: proof1 } = await mainProgram.addSideloadedProgram(
   sideVk,
   witness
 );
+perfMain.end();
 
-console.log('Proving child program execution');
+console.log('\nProving child program execution');
+perfSide.start('prove', 'compute');
 const { proof: childProof1 } = await sideloadedProgram.compute(Field(0), Field(10));
+perfSide.end();
 
-console.log('Proving verification inside main program');
+console.log('\nProving verification inside main program');
+perfMain.start('prove', 'validateUsingTree');
 const { proof: proof2 } = await mainProgram.validateUsingTree(
   proof1.publicOutput,
   proof1,
   sideVk,
   witness,
   SideloadedProgramProof.fromProof(childProof1)
 );
+perfMain.end();
 
+perfMain.start('verify', 'validateUsingTree');
 const validProof2 = await verify(proof2, mainVk);
+perfMain.end();
 console.log('ok?', validProof2);
 
-console.log('Proving different method of child program');
+console.log('\nProving different method of child program');
+perfSide.start('prove', 'assertAndAdd');
 const { proof: childProof2 } = await sideloadedProgram.assertAndAdd(Field(0), Field(10));
+perfSide.end();
 
-console.log('Proving verification inside main program');
+console.log('\nProving verification inside main program');
 const { proof: proof3 } = await mainProgram.validateUsingTree(
   proof1.publicOutput,
   proof1,

tests/perf-regression/perf-regression.sh

@@ -15,5 +15,12 @@ MODE=$1
 ./run src/examples/crypto/blake2b/run.ts --bundle "$MODE"
 ./run src/examples/crypto/rsa/run.ts --bundle "$MODE"
 
+./run src/examples/zkprogram/mutual-recursion.ts --bundle "$MODE"
+./run src/examples/zkprogram/hash-chain.ts --bundle "$MODE"
+./run src/examples/zkprogram/gadgets.ts --bundle "$MODE"
+./run src/examples/zkprogram/side-loading/run.ts --bundle "$MODE"
+./run src/examples/zkprogram/runtime-table/run.ts --bundle "$MODE"
+./run src/examples/zkprogram/program-small-big.ts --bundle "$MODE"
+
 # Run CS + zkApps performance regression tests
 ./run tests/perf-regression/perf-regression.ts --bundle "$MODE"

tests/vk-regression/diverse-zk-program-run.ts

@@ -1,11 +1,29 @@
 import assert from 'node:assert';
-import { diverse, Bytes128 } from './diverse-zk-program.js';
+import { verify } from 'o1js';
+import { Performance } from '../../src/lib/testing/perf-regression.js';
+import { Bytes128, diverse } from './diverse-zk-program.js';
 
-console.log('testing proof generation for diverse program');
-await diverse.compile();
+const cs = await diverse.analyzeMethods();
+const perfDiverse = Performance.create(diverse.name, cs);
 
+perfDiverse.start('compile');
+const { verificationKey } = await diverse.compile();
+perfDiverse.end();
+
+perfDiverse.start('prove', 'sha3');
 let { proof: proof1 } = await diverse.sha3(Bytes128.fromString('hello'));
-assert(await diverse.verify(proof1), 'verifies');
+perfDiverse.end();
+
+perfDiverse.start('verify', 'sha3');
+const isValid1 = await verify(proof1, verificationKey);
+perfDiverse.end();
+assert(isValid1, 'proof1 verification failed!');
 
+perfDiverse.start('prove', 'recursive');
 let { proof: proof2 } = await diverse.recursive(proof1);
-assert(await diverse.verify(proof2), 'verifies');
+perfDiverse.end();
+
+perfDiverse.start('verify', 'recursive');
+const isValid2 = await verify(proof2, verificationKey);
+perfDiverse.end();
+assert(isValid2, 'proof2 verification failed!');