Merge pull request o1-labs#1847 from o1-labs/feature/provable-proofs-2

Proof as Provable

src/lib/mina/test/test-contract.ts

@@ -170,7 +170,7 @@ async function runInstruction(
     let feepayer = instruction.sender ?? sender;
     let signers = [feepayer.key, ...(instruction.signers ?? [])];
     let tx = await Mina.transaction(feepayer, instruction.callback);
-    await assertionWithTrace(instruction.trace, async () => {
+    await assertionWithTracePreferInner(instruction.trace, async () => {
       // console.log(instruction.label, tx.toPretty());
       await tx.sign(signers).prove();
       await tx.send();
@@ -311,9 +311,28 @@ async function assertionWithTrace(trace: string | undefined, fn: () => any) {
   try {
     await fn();
   } catch (err: any) {
-    if (trace !== undefined) {
-      err.message += `\n\nAssertion was created here:${trace}\n\nError was thrown from here:`;
-    }
-    throw Error(err.message);
+    if (trace === undefined) throw err;
+
+    let message = err.message;
+    throw Error(
+      `${message}\n\nAssertion was created here:${trace}\n\nError was thrown from here:`
+    );
+  }
+}
+
+async function assertionWithTracePreferInner(
+  trace: string | undefined,
+  fn: () => any
+) {
+  try {
+    await fn();
+  } catch (err: any) {
+    if (trace === undefined) throw err;
+
+    // note: overwriting the message doesn't work with all errors and I don't know why.
+    // we call this function, rather than `assertionWithInner()`, when we're ok with not having the added message,
+    // and prefer preserving the original stack trace at all costs
+    err.message = `${err.message}\n\nError was thrown from here:${trace}\n\nAssertion was created here:`;
+    throw err;
   }
 }

src/lib/mina/zkapp.ts

@@ -156,7 +156,7 @@ function method<K extends string, T extends SmartContract>(
   ZkappClass._maxProofsVerified ??= 0;
   ZkappClass._maxProofsVerified = Math.max(
     ZkappClass._maxProofsVerified,
-    methodEntry.proofArgs.length
+    methodEntry.proofs.length
   ) as 0 | 1 | 2;
   let func = descriptor.value as AsyncFunction;
   descriptor.value = wrapMethod(func, ZkappClass, internalMethodEntry);
@@ -313,12 +313,7 @@ function wrapMethod(
               method.apply(
                 this,
                 actualArgs.map((a, i) => {
-                  let arg = methodIntf.allArgs[i];
-                  if (arg.type === 'witness') {
-                    let type = methodIntf.witnessArgs[arg.index];
-                    return Provable.witness(type, () => a);
-                  }
-                  return a;
+                  return Provable.witness(methodIntf.args[i].type, () => a);
                 })
               ),
               noPromiseError

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

@@ -44,12 +44,11 @@ it('pickles rule creation', async () => {
 
   expect(methodIntf).toEqual({
     methodName: 'main',
-    witnessArgs: [Bool],
-    proofArgs: [EmptyProof],
-    allArgs: [
-      { type: 'proof', index: 0 },
-      { type: 'witness', index: 0 },
+    args: [
+      { type: EmptyProof, isProof: true },
+      { type: Bool, isProof: false },
     ],
+    proofs: [EmptyProof],
   });
 
   // store compiled tag

src/lib/proof-system/proof.ts

@@ -1,22 +1,25 @@
 import { initializeBindings, withThreadPool } from '../../snarky.js';
 import { Pickles } from '../../snarky.js';
 import { Field, Bool } from '../provable/wrapped.js';
-import {
+import type {
   FlexibleProvablePure,
   InferProvable,
 } from '../provable/types/struct.js';
-import { ProvablePure } from '../provable/types/provable-intf.js';
 import { FeatureFlags } from './feature-flags.js';
 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';
 
 // public API
 export { ProofBase, Proof, DynamicProof };
 
 // internal API
-export { dummyProof, getStatementType };
+export { dummyProof };
 
-class ProofBase<Input, Output> {
+type MaxProofs = 0 | 1 | 2;
+
+class ProofBase<Input = any, Output = any> {
   static publicInputType: FlexibleProvablePure<any> = undefined as any;
   static publicOutputType: FlexibleProvablePure<any> = undefined as any;
   static tag: () => { name: string } = () => {
@@ -32,10 +35,10 @@ class ProofBase<Input, Output> {
   shouldVerify = Bool(false);
 
   toJSON(): JsonProof {
-    let type = getStatementType(this.constructor as any);
+    let fields = this.publicFields();
     return {
-      publicInput: type.input.toFields(this.publicInput).map(String),
-      publicOutput: type.output.toFields(this.publicOutput).map(String),
+      publicInput: fields.input.map(String),
+      publicOutput: fields.output.map(String),
       maxProofsVerified: this.maxProofsVerified,
       proof: Pickles.proofToBase64([this.maxProofsVerified, this.proof]),
     };
@@ -57,6 +60,36 @@ class ProofBase<Input, Output> {
     this.proof = proof; // TODO optionally convert from string?
     this.maxProofsVerified = maxProofsVerified;
   }
+
+  static get provable() {
+    if (
+      this.publicInputType === undefined ||
+      this.publicOutputType === undefined
+    ) {
+      throw Error(
+        `You cannot use the \`Proof\` class directly. Instead, define a subclass:\n` +
+          `class MyProof extends Proof<PublicInput, PublicOutput> { ... }`
+      );
+    }
+    return provableProof<any, any, any>(
+      this,
+      this.publicInputType,
+      this.publicOutputType,
+      (this as any).maxProofsVerified
+    );
+  }
+
+  static publicFields(value: ProofBase) {
+    let fields = this.provable.toFields(value);
+    let inputSize = this.publicInputType.sizeInFields();
+    return {
+      input: fields.slice(0, inputSize),
+      output: fields.slice(inputSize),
+    };
+  }
+  publicFields() {
+    return (this.constructor as typeof ProofBase).publicFields(this);
+  }
 }
 
 class Proof<Input, Output> extends ProofBase<Input, Output> {
@@ -83,15 +116,12 @@ class Proof<Input, Output> extends ProofBase<Input, Output> {
   > {
     await initializeBindings();
     let [, proof] = Pickles.proofOfBase64(proofString, maxProofsVerified);
-    let type = getStatementType(this);
-    let publicInput = type.input.fromFields(publicInputJson.map(Field));
-    let publicOutput = type.output.fromFields(publicOutputJson.map(Field));
-    return new this({
-      publicInput,
-      publicOutput,
-      proof,
-      maxProofsVerified,
-    }) as any;
+    let fields = publicInputJson.map(Field).concat(publicOutputJson.map(Field));
+    return this.provable.fromFields(fields, [
+      [],
+      [],
+      [proof, maxProofsVerified],
+    ]) as any;
   }
 
   /**
@@ -131,7 +161,7 @@ class Proof<Input, Output> extends ProofBase<Input, Output> {
   }
 }
 
-var sideloadedKeysCounter = 0;
+let sideloadedKeysCounter = 0;
 
 /**
  * The `DynamicProof` class enables circuits to verify proofs using in-ciruit verfication keys.
@@ -181,7 +211,7 @@ class DynamicProof<Input, Output> extends ProofBase<Input, Output> {
    * If you want to verify _any_ proof, no matter what custom gates it uses, you can use {@link FeatureFlags.allMaybe}. Please note that this might incur a significant overhead.
    *
    * You can also toggle specific feature flags manually by specifying them here.
-   * Alternatively, you can use {@FeatureFlags.fromZkProgram} to compute the set of feature flags that are compatible with a given program.
+   * Alternatively, you can use {@link FeatureFlags.fromZkProgram} to compute the set of feature flags that are compatible with a given program.
    */
   static featureFlags: FeatureFlags = FeatureFlags.allNone;
 
@@ -227,15 +257,12 @@ class DynamicProof<Input, Output> extends ProofBase<Input, Output> {
   > {
     await initializeBindings();
     let [, proof] = Pickles.proofOfBase64(proofString, maxProofsVerified);
-    let type = getStatementType(this);
-    let publicInput = type.input.fromFields(publicInputJson.map(Field));
-    let publicOutput = type.output.fromFields(publicOutputJson.map(Field));
-    return new this({
-      publicInput,
-      publicOutput,
-      proof,
-      maxProofsVerified,
-    }) as any;
+    let fields = publicInputJson.map(Field).concat(publicOutputJson.map(Field));
+    return this.provable.fromFields(fields, [
+      [],
+      [],
+      [proof, maxProofsVerified],
+    ]) as any;
   }
 
   static async dummy<S extends Subclass<typeof DynamicProof>>(
@@ -281,24 +308,81 @@ async function dummyProof(maxProofsVerified: 0 | 1 | 2, domainLog2: number) {
   );
 }
 
-// helper
-
-function getStatementType<
-  T,
-  O,
-  P extends Subclass<typeof ProofBase> = typeof ProofBase
->(Proof: P): { input: ProvablePure<T>; output: ProvablePure<O> } {
-  if (
-    Proof.publicInputType === undefined ||
-    Proof.publicOutputType === undefined
-  ) {
-    throw Error(
-      `You cannot use the \`Proof\` class directly. Instead, define a subclass:\n` +
-        `class MyProof extends Proof<PublicInput, PublicOutput> { ... }`
-    );
+function provableProof<
+  Class extends Subclass<typeof ProofBase<Input, Output>>,
+  Input,
+  Output,
+  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;
   }
+> {
   return {
-    input: Proof.publicInputType as any,
-    output: Proof.publicOutputType as any,
+    sizeInFields() {
+      return input.sizeInFields() + output.sizeInFields();
+    },
+    toFields(value) {
+      return input
+        .toFields(value.publicInput)
+        .concat(output.toFields(value.publicOutput));
+    },
+    toAuxiliary(value) {
+      let inputAux = input.toAuxiliary(value?.publicInput);
+      let outputAux = output.toAuxiliary(value?.publicOutput);
+      let proofAux = [
+        value?.proof ?? undefined,
+        value?.maxProofsVerified ?? defaultMaxProofsVerified ?? 0,
+      ];
+      return [inputAux, outputAux, proofAux];
+    },
+    fromFields(fields, aux) {
+      let inputFields = fields.slice(0, input.sizeInFields());
+      let outputFields = fields.slice(input.sizeInFields());
+      assert(outputFields.length === output.sizeInFields());
+      let [inputAux, outputAux, [proof, maxProofsVerified]] = aux;
+      let publicInput = input.fromFields(inputFields, inputAux);
+      let publicOutput = output.fromFields(outputFields, outputAux);
+      return new Class({
+        publicInput,
+        publicOutput,
+        proof,
+        maxProofsVerified,
+      });
+    },
+    check(value) {
+      input.check(value.publicInput);
+      output.check(value.publicOutput);
+    },
+    toValue(value) {
+      let inputV = input.toValue(value.publicInput);
+      let outputV = output.toValue(value.publicOutput);
+      return {
+        publicInput: inputV,
+        publicOutput: outputV,
+        proof: value.proof,
+        maxProofsVerified: value.maxProofsVerified,
+      };
+    },
+    fromValue(value) {
+      let inputT = input.fromValue(value.publicInput);
+      let outputT = output.fromValue(value.publicOutput);
+      return new Class({
+        publicInput: inputT,
+        publicOutput: outputT,
+        proof: value.proof,
+        maxProofsVerified: value.maxProofsVerified,
+      });
+    },
   };
 }

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

@@ -0,0 +1,46 @@
+import { expect } from 'expect';
+import { Provable } from '../provable/provable.js';
+import { Field } from '../provable/wrapped.js';
+import { Proof } from './proof.js';
+import { Void } from './zkprogram.js';
+import test from 'node:test';
+
+test('Proof provable', async () => {
+  class MyProof extends Proof<Field, Void> {
+    static publicInputType = Field;
+    static publicOutputType = Void;
+  }
+
+  expect(MyProof.provable.sizeInFields()).toEqual(1);
+
+  let proof = await MyProof.dummy(Field(1n), undefined, 0);
+
+  expect(MyProof.provable.toFields(proof)).toEqual([Field(1n)]);
+  expect(MyProof.provable.toAuxiliary(proof)).toEqual([
+    [],
+    [],
+    [proof.proof, 0],
+  ]);
+
+  async function circuit() {
+    let publicInput = Provable.witness(Field, () => 5n);
+
+    let proof = await Provable.witnessAsync(MyProof, () => {
+      return MyProof.dummy(publicInput, undefined, 0);
+    });
+
+    expect(proof).toBeInstanceOf(MyProof);
+    expect(proof.publicOutput).toBeUndefined();
+    expect(proof.maxProofsVerified).toEqual(0);
+
+    Provable.asProver(() => {
+      expect(proof.publicInput.toBigInt()).toEqual(5n);
+
+      let value = MyProof.provable.toValue(proof);
+      expect(value.publicInput).toEqual(5n);
+      expect(value.proof).toBe(proof.proof);
+    });
+  }
+
+  await Provable.runAndCheck(circuit);
+});