Merge pull request #1933 from o1-labs/feature/zkprogram-qol

ZkPrograms QoL improvements

CHANGELOG.md

@@ -30,6 +30,11 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 - Method for optional types to assert none https://github.com/o1-labs/o1js/pull/1922
 - Increased maximum supported amount of methods in a `SmartContract` or `ZkProgram` to 30. https://github.com/o1-labs/o1js/pull/1918
 - Expose low-level conversion methods `Proof.{_proofToBase64,_proofFromBase64}` https://github.com/o1-labs/o1js/pull/1928
+- Expore `maxProofsVerified()` and a `Proof` class directly on ZkPrograms https://github.com/o1-labs/o1js/pull/1933
+
+### Changed
+
+- Changed an internal type to improve IntelliSense on ZkProgram methods https://github.com/o1-labs/o1js/pull/1933
 
 ### Fixed
 

src/examples/zkprogram/mututal-recursion.ts

@@ -47,15 +47,13 @@ const add = ZkProgram({
   },
 });
 
-const AddProof = ZkProgram.Proof(add);
-
 const multiply = ZkProgram({
   name: 'multiply',
   publicInput: Undefined,
   publicOutput: Field,
   methods: {
     performMultiplication: {
-      privateInputs: [Field, AddProof],
+      privateInputs: [Field, add.Proof],
       async method(field: Field, addProof: Proof<Undefined, Field>) {
         addProof.verify();
         const multiplicationResult = addProof.publicOutput.mul(field);

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

@@ -35,8 +35,6 @@ MyProgram.publicOutputType satisfies Provable<void>;
 MyProgram.privateInputTypes;
 MyProgram.auxiliaryOutputTypes;
 
-let MyProof = ZkProgram.Proof(MyProgram);
-
 console.log('program digest', await MyProgram.digest());
 
 console.log('compiling MyProgram...');
@@ -45,7 +43,7 @@ console.log('verification key', verificationKey.data.slice(0, 10) + '..');
 
 console.log('proving base case...');
 let { proof } = await MyProgram.baseCase(Field(0));
-proof = await testJsonRoundtrip(MyProof, proof);
+proof = await testJsonRoundtrip(MyProgram.Proof, proof);
 
 // type sanity check
 proof satisfies Proof<Field, void>;
@@ -60,7 +58,7 @@ console.log('ok (alternative)?', ok);
 
 console.log('proving step 1...');
 let { proof: proof1 } = await MyProgram.inductiveCase(Field(1), proof);
-proof1 = await testJsonRoundtrip(MyProof, proof1);
+proof1 = await testJsonRoundtrip(MyProgram.Proof, proof1);
 
 console.log('verify...');
 ok = await verify(proof1, verificationKey);
@@ -72,7 +70,7 @@ console.log('ok (alternative)?', ok);
 
 console.log('proving step 2...');
 let { proof: proof2 } = await MyProgram.inductiveCase(Field(2), proof1);
-proof2 = await testJsonRoundtrip(MyProof, proof2);
+proof2 = await testJsonRoundtrip(MyProgram.Proof, proof2);
 
 console.log('verify...');
 ok = await verify(proof2.toJSON(), verificationKey);

src/lib/mina/actions/offchain-state-rollup.ts

@@ -239,7 +239,7 @@ function OffchainStateRollup({
     },
   });
 
-  let RollupProof = ZkProgram.Proof(offchainStateRollup);
+  let RollupProof = offchainStateRollup.Proof;
 
   let isCompiled = false;
 

src/lib/proof-system/proof-system.unit-test.ts

@@ -26,7 +26,7 @@ const EmptyProgram = ZkProgram({
   methods: { run: { privateInputs: [], async method(_) {} } },
 });
 
-class EmptyProof extends ZkProgram.Proof(EmptyProgram) {}
+class EmptyProof extends EmptyProgram.Proof {}
 
 // unit-test zkprogram creation helpers:
 // -) sortMethodArguments

src/lib/proof-system/zkprogram.ts

@@ -224,6 +224,8 @@ function ZkProgram<
   }
 ): {
   name: string;
+  maxProofsVerified(): Promise<0 | 1 | 2>;
+
   compile: (options?: {
     cache?: Cache;
     forceRecompile?: boolean;
@@ -243,13 +245,20 @@ function ZkProgram<
       ReturnType<typeof analyzeMethod>
     >;
   }>;
+
   publicInputType: ProvableOrUndefined<Get<Config, 'publicInput'>>;
   publicOutputType: ProvableOrVoid<Get<Config, 'publicOutput'>>;
   privateInputTypes: PrivateInputs;
   auxiliaryOutputTypes: AuxiliaryOutputs;
   rawMethods: {
     [I in keyof Config['methods']]: Methods[I]['method'];
   };
+
+  Proof: typeof Proof<
+    InferProvableOrUndefined<Get<Config, 'publicInput'>>,
+    InferProvableOrVoid<Get<Config, 'publicOutput'>>
+  >;
+
   proofsEnabled: boolean;
   setProofsEnabled(proofsEnabled: boolean): void;
 } & {
@@ -522,10 +531,13 @@ function ZkProgram<
   const program = Object.assign(
     selfTag,
     {
+      maxProofsVerified: getMaxProofsVerified,
+
       compile,
       verify,
       digest,
       analyzeMethods,
+
       publicInputType: publicInputType as ProvableOrUndefined<
         Get<Config, 'publicInput'>
       >,
@@ -541,6 +553,9 @@ function ZkProgram<
       rawMethods: Object.fromEntries(
         methodKeys.map((key) => [key, methods[key].method])
       ) as any,
+
+      Proof: SelfProof,
+
       proofsEnabled: doProving,
       setProofsEnabled(proofsEnabled: boolean) {
         doProving = proofsEnabled;
@@ -1127,7 +1142,7 @@ type Infer<T> = T extends Subclass<typeof ProofBase>
 
 type TupleToInstances<T> = {
   [I in keyof T]: Infer<T[I]>;
-} & any[];
+};
 
 type PrivateInput = ProvableType | Subclass<typeof ProofBase>;
 

src/lib/provable/test/custom-gates-recursion.unit-test.ts

@@ -29,7 +29,7 @@ let emptyProgram = ZkProgram({
   name: 'empty',
   methods: { run: { privateInputs: [], async method() {} } },
 });
-class EmptyProof extends ZkProgram.Proof(emptyProgram) {}
+class EmptyProof extends emptyProgram.Proof {}
 
 let program = ZkProgram({
   name: 'ecdsa',

src/tests/fake-proof.ts

@@ -42,7 +42,7 @@ const FakeProgram = ZkProgram({
   },
 });
 
-class RealProof extends ZkProgram.Proof(RealProgram) {}
+class RealProof extends RealProgram.Proof {}
 class Nested extends Struct({ inner: RealProof }) {}
 
 const RecursiveProgram = ZkProgram({

src/tests/inductive-proofs.ts

@@ -1,5 +1,5 @@
-import { SelfProof, Field, ZkProgram, Proof } from 'o1js';
-import { tic, toc } from '../examples/utils/tic-toc.node.js';
+import { SelfProof, Field, ZkProgram, Proof, JsonProof } from 'o1js';
+import { tic, toc } from '../examples/utils/tic-toc.js';
 
 let MaxProofsVerifiedZero = ZkProgram({
   name: 'no-recursion',
@@ -100,12 +100,12 @@ async function testRecursion(
 ) {
   console.log(`testing maxProofsVerified = ${maxProofsVerified}`);
 
-  let ProofClass = ZkProgram.Proof(Program);
+  class ProofClass extends Program.Proof {}
 
   tic('executing base case');
-  let initialProof = await Program.baseCase(Field(0));
+  let { proof: initialProof } = await Program.baseCase(Field(0));
   toc();
-  initialProof = testJsonRoundtrip(ProofClass, initialProof);
+  initialProof = await testJsonRoundtrip(ProofClass, initialProof);
   initialProof.verify();
   initialProof.publicInput.assertEquals(Field(0));
 
@@ -115,13 +115,13 @@ async function testRecursion(
     );
   }
 
-  let p1, p2;
+  let p1: Proof<any, any>, p2: Proof<any, any>;
   if (initialProof.maxProofsVerified === 0) return;
 
   tic('executing mergeOne');
-  p1 = await Program.mergeOne(Field(1), initialProof);
+  p1 = (await Program.mergeOne(Field(1), initialProof)).proof;
   toc();
-  p1 = testJsonRoundtrip(ProofClass, p1);
+  p1 = await testJsonRoundtrip(ProofClass, p1);
   p1.verify();
   p1.publicInput.assertEquals(Field(1));
   if (p1.maxProofsVerified != maxProofsVerified) {
@@ -132,9 +132,9 @@ async function testRecursion(
 
   if (initialProof.maxProofsVerified === 1) return;
   tic('executing mergeTwo');
-  p2 = await Program.mergeTwo(Field(2), initialProof, p1);
+  p2 = (await Program.mergeTwo(Field(2), initialProof, p1)).proof;
   toc();
-  p2 = testJsonRoundtrip(ProofClass, p2);
+  p2 = await testJsonRoundtrip(ProofClass, p2);
   p2.verify();
   p2.publicInput.assertEquals(Field(2));
   if (p2.maxProofsVerified != maxProofsVerified) {
@@ -144,7 +144,10 @@ async function testRecursion(
   }
 }
 
-function testJsonRoundtrip(ProofClass: any, proof: Proof<Field, void>) {
+function testJsonRoundtrip(
+  ProofClass: { fromJSON: (p: JsonProof) => Promise<Proof<any, any>> },
+  proof: Proof<Field, void>
+) {
   let jsonProof = proof.toJSON();
   console.log(
     'json roundtrip',