Merge pull request #1851 from o1-labs/feature/provable-proofs-3

Extract Proofs from any Provable

CHANGELOG.md

@@ -21,6 +21,12 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 - Added `VerificationKey.dummy()` method to get the dummy value of a verification key https://github.com/o1-labs/o1js/pull/1852 [@rpanic](https://github.com/rpanic)
 
+### Changed
+
+- Make `Proof` a normal provable type, that can be witnessed and composed into Structs https://github.com/o1-labs/o1js/pull/1847, https://github.com/o1-labs/o1js/pull/1851
+  - ZkProgram and SmartContract now also support private inputs that are not proofs themselves, but contain proofs nested within a Struct or array
+  - Only `SelfProof` can still not be nested because it needs special treatment
+
 ## [1.8.0](https://github.com/o1-labs/o1js/compare/5006e4f...450943) - 2024-09-18
 
 ### Added

src/lib/mina/zkapp.ts

@@ -46,8 +46,6 @@ import {
   compileProgram,
   Empty,
   getPreviousProofsForProver,
-  methodArgumentsToConstant,
-  methodArgumentTypesAndValues,
   MethodInterface,
   sortMethodArguments,
   VerificationKey,
@@ -157,7 +155,7 @@ function method<K extends string, T extends SmartContract>(
   ZkappClass._maxProofsVerified ??= 0;
   ZkappClass._maxProofsVerified = Math.max(
     ZkappClass._maxProofsVerified,
-    methodEntry.proofs.length
+    methodEntry.numberOfProofs
   ) as 0 | 1 | 2;
   let func = descriptor.value as AsyncFunction;
   descriptor.value = wrapMethod(func, ZkappClass, internalMethodEntry);
@@ -314,7 +312,7 @@ function wrapMethod(
               method.apply(
                 this,
                 actualArgs.map((a, i) => {
-                  return Provable.witness(methodIntf.args[i].type, () => a);
+                  return Provable.witness(methodIntf.args[i], () => a);
                 })
               ),
               noPromiseError
@@ -342,10 +340,7 @@ function wrapMethod(
                 methodName: methodIntf.methodName,
                 args: clonedArgs,
                 // proofs actually don't have to be cloned
-                previousProofs: getPreviousProofsForProver(
-                  actualArgs,
-                  methodIntf
-                ),
+                previousProofs: getPreviousProofsForProver(actualArgs),
                 ZkappClass,
                 memoized,
                 blindingValue,
@@ -387,7 +382,9 @@ function wrapMethod(
       let blindingValue = getBlindingValue();
 
       let runCalledContract = async () => {
-        let constantArgs = methodArgumentsToConstant(methodIntf, actualArgs);
+        let constantArgs = methodIntf.args.map((type, i) =>
+          Provable.toConstant(type, actualArgs[i])
+        );
         let constantBlindingValue = blindingValue.toConstant();
         let accountUpdate = this.self;
         accountUpdate.body.callDepth = parentAccountUpdate.body.callDepth + 1;
@@ -434,10 +431,7 @@ function wrapMethod(
             {
               methodName: methodIntf.methodName,
               args: constantArgs,
-              previousProofs: getPreviousProofsForProver(
-                constantArgs,
-                methodIntf
-              ),
+              previousProofs: getPreviousProofsForProver(constantArgs),
               ZkappClass,
               memoized,
               blindingValue: constantBlindingValue,
@@ -527,7 +521,9 @@ function computeCallData(
   blindingValue: Field
 ) {
   let { returnType, methodName } = methodIntf;
-  let args = methodArgumentTypesAndValues(methodIntf, argumentValues);
+  let args = methodIntf.args.map((type, i) => {
+    return { type: ProvableType.get(type), value: argumentValues[i] };
+  });
 
   let input: HashInput = { fields: [], packed: [] };
   for (let { type, value } of args) {

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

@@ -4,6 +4,7 @@ import { UInt64 } from '../provable/int.js';
 import {
   CompiledTag,
   Empty,
+  Void,
   ZkProgram,
   picklesRuleFromFunction,
   sortMethodArguments,
@@ -44,11 +45,8 @@ it('pickles rule creation', async () => {
 
   expect(methodIntf).toEqual({
     methodName: 'main',
-    args: [
-      { type: EmptyProof, isProof: true },
-      { type: Bool, isProof: false },
-    ],
-    proofs: [EmptyProof],
+    args: [EmptyProof, Bool],
+    numberOfProofs: 1,
   });
 
   // store compiled tag
@@ -99,6 +97,76 @@ it('pickles rule creation', async () => {
   );
 });
 
+class NestedProof extends Struct({ proof: EmptyProof, field: Field }) {}
+const NestedProof2 = Provable.Array(NestedProof, 2);
+
+// type inference
+NestedProof satisfies Provable<{ proof: Proof<Field, void>; field: Field }>;
+
+it('pickles rule creation: nested proof', async () => {
+  function main([first, _second]: [NestedProof, NestedProof]) {
+    // first proof should verify, second should not
+    first.proof.verify();
+
+    // deep type inference
+    first.proof.publicInput satisfies Field;
+    first.proof.publicOutput satisfies void;
+  }
+
+  // collect method interface
+  let methodIntf = sortMethodArguments('mock', 'main', [NestedProof2], Proof);
+
+  expect(methodIntf).toEqual({
+    methodName: 'main',
+    args: [NestedProof2],
+    numberOfProofs: 2,
+  });
+
+  // store compiled tag
+  CompiledTag.store(EmptyProgram, 'mock tag');
+
+  // create pickles rule
+  let rule: Pickles.Rule = picklesRuleFromFunction(
+    Empty as ProvablePure<any>,
+    Void as ProvablePure<any>,
+    main as AnyFunction,
+    { name: 'mock' },
+    methodIntf,
+    []
+  );
+
+  let dummy = await EmptyProof.dummy(Field(0), undefined, 0);
+  let nested1 = new NestedProof({ proof: dummy, field: Field(0) });
+  let nested2 = new NestedProof({ proof: dummy, field: Field(0) });
+  let nested = [nested1, nested2];
+
+  await Provable.runAndCheck(async () => {
+    // put witnesses in snark context
+    snarkContext.get().witnesses = [nested];
+
+    // call pickles rule
+    let {
+      shouldVerify: [, shouldVerify1, shouldVerify2],
+      previousStatements: [, ...previousStatements],
+    } = await rule.main([0]);
+
+    expect(previousStatements.length).toBe(2);
+
+    // `shouldVerify` are as expected
+    expect(Bool(shouldVerify1).isConstant()).toBe(true);
+    expect(Bool(shouldVerify2).isConstant()).toBe(true);
+    // first proof should verify, second should not
+    Bool(shouldVerify1).assertTrue();
+    Bool(shouldVerify2).assertFalse();
+  });
+});
+
+it('fails with more than two (nested) proofs', async () => {
+  expect(() => {
+    sortMethodArguments('mock', 'main', [NestedProof2, NestedProof], Proof);
+  }).toThrowError('mock.main() has more than two proof arguments');
+});
+
 // compile works with large inputs
 
 const N = 100_000;

src/lib/proof-system/proof.ts

@@ -10,12 +10,14 @@ import type { VerificationKey, JsonProof } from './zkprogram.js';
 import { Subclass } from '../util/types.js';
 import type { Provable } from '../provable/provable.js';
 import { assert } from '../util/assert.js';
+import { Unconstrained } from '../provable/types/unconstrained.js';
+import { ProvableType } from '../provable/types/provable-intf.js';
 
 // public API
 export { ProofBase, Proof, DynamicProof };
 
 // internal API
-export { dummyProof };
+export { dummyProof, extractProofs, extractProofTypes, type ProofValue };
 
 type MaxProofs = 0 | 1 | 2;
 
@@ -61,7 +63,7 @@ class ProofBase<Input = any, Output = any> {
     this.maxProofsVerified = maxProofsVerified;
   }
 
-  static get provable() {
+  static get provable(): Provable<any> {
     if (
       this.publicInputType === undefined ||
       this.publicOutputType === undefined
@@ -71,7 +73,7 @@ class ProofBase<Input = any, Output = any> {
           `class MyProof extends Proof<PublicInput, PublicOutput> { ... }`
       );
     }
-    return provableProof<any, any, any>(
+    return provableProof(
       this,
       this.publicInputType,
       this.publicOutputType,
@@ -159,6 +161,10 @@ class Proof<Input, Output> extends ProofBase<Input, Output> {
       maxProofsVerified,
     });
   }
+
+  static get provable(): ProvableProof<Proof<any, any>> {
+    return super.provable as any;
+  }
 }
 
 let sideloadedKeysCounter = 0;
@@ -299,6 +305,10 @@ class DynamicProof<Input, Output> extends ProofBase<Input, Output> {
       proof: proof.proof,
     }) as InstanceType<S>;
   }
+
+  static get provable(): ProvableProof<DynamicProof<any, any>> {
+    return super.provable as any;
+  }
 }
 
 async function dummyProof(maxProofsVerified: 0 | 1 | 2, domainLog2: number) {
@@ -308,26 +318,31 @@ async function dummyProof(maxProofsVerified: 0 | 1 | 2, domainLog2: number) {
   );
 }
 
+type ProofValue<Input, Output> = {
+  publicInput: Input;
+  publicOutput: Output;
+  proof: Pickles.Proof;
+  maxProofsVerified: 0 | 1 | 2;
+};
+
+type ProvableProof<
+  Proof extends ProofBase,
+  InputV = any,
+  OutputV = any
+> = Provable<Proof, ProofValue<InputV, OutputV>>;
+
 function provableProof<
   Class extends Subclass<typeof ProofBase<Input, Output>>,
-  Input,
-  Output,
+  Input = any,
+  Output = any,
   InputV = any,
   OutputV = any
 >(
   Class: Class,
   input: Provable<Input>,
   output: Provable<Output>,
   defaultMaxProofsVerified?: MaxProofs
-): Provable<
-  ProofBase<Input, Output>,
-  {
-    publicInput: InputV;
-    publicOutput: OutputV;
-    proof: unknown;
-    maxProofsVerified: MaxProofs;
-  }
-> {
+): Provable<ProofBase<Input, Output>, ProofValue<InputV, OutputV>> {
   return {
     sizeInFields() {
       return input.sizeInFields() + output.sizeInFields();
@@ -386,3 +401,30 @@ function provableProof<
     },
   };
 }
+
+function extractProofs(value: unknown): ProofBase[] {
+  if (value instanceof ProofBase) {
+    return [value];
+  }
+  if (value instanceof Unconstrained) return [];
+  if (value instanceof Field) return [];
+  if (value instanceof Bool) return [];
+
+  if (Array.isArray(value)) {
+    return value.flatMap((item) => extractProofs(item));
+  }
+
+  if (value === null) return [];
+  if (typeof value === 'object') {
+    return extractProofs(Object.values(value));
+  }
+
+  // functions, primitives
+  return [];
+}
+
+function extractProofTypes(type: ProvableType) {
+  let value = ProvableType.synthesize(type);
+  let proofValues = extractProofs(value);
+  return proofValues.map((proof) => proof.constructor as typeof ProofBase);
+}