Clean up previous proof arguments handling in Pickles/SnarkyJS

CHANGELOG.md

@@ -17,7 +17,12 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased](https://github.com/o1-labs/snarkyjs/compare/bcc666f2...HEAD)
 
-> No unreleased changes yet
+### Changes
+
+- Allow ZkPrograms to return their public output https://github.com/o1-labs/snarkyjs/pull/874 https://github.com/o1-labs/snarkyjs/pull/876
+  - new option `ZkProgram({ publicOutput?: Provable<any>, ... })`; `publicOutput` has to match the _return type_ of all ZkProgram methods.
+  - the `publicInput` option becomes optional; if not provided, methods no longer expect the public input as first argument
+  - full usage example: https://github.com/o1-labs/snarkyjs/blob/f95cf2903e97292df9e703b74ee1fc3825df826d/src/examples/program.ts
 
 ## [0.10.0](https://github.com/o1-labs/snarkyjs/compare/97e393ed...bcc666f2)
 

src/examples/program.ts

@@ -7,6 +7,7 @@ import {
   Proof,
   JsonProof,
   Provable,
+  Empty,
 } from 'snarkyjs';
 
 await isReady;
@@ -24,15 +25,15 @@ let MyProgram = Experimental.ZkProgram({
 
     inductiveCase: {
       privateInputs: [SelfProof],
-      method(earlierProof: SelfProof<undefined, Field>) {
+      method(earlierProof: SelfProof<Empty, Field>) {
         earlierProof.verify();
         return earlierProof.publicOutput.add(1);
       },
     },
   },
 });
 // type sanity checks
-MyProgram.publicInputType satisfies Provable<undefined>;
+MyProgram.publicInputType satisfies Provable<Empty>;
 MyProgram.publicOutputType satisfies typeof Field;
 
 let MyProof = Experimental.ZkProgram.Proof(MyProgram);

src/examples/zkapps/simple_zkapp_with_proof.ts

@@ -11,6 +11,7 @@ import {
   ZkappPublicInput,
   SelfProof,
   verify,
+  Empty,
 } from 'snarkyjs';
 
 await isReady;
@@ -32,7 +33,7 @@ class NotSoSimpleZkapp extends SmartContract {
 
   @method update(
     y: Field,
-    oldProof: SelfProof<ZkappPublicInput, null>,
+    oldProof: SelfProof<ZkappPublicInput, Empty>,
     trivialProof: TrivialProof
   ) {
     oldProof.verify();

src/index.ts

@@ -37,6 +37,7 @@ export {
   SelfProof,
   verify,
   JsonProof,
+  Empty,
   Undefined,
   Void,
 } from './lib/proof_system.js';

src/lib/account_update.ts

@@ -16,6 +16,7 @@ import { SmartContract } from './zkapp.js';
 import * as Precondition from './precondition.js';
 import {
   dummyBase64Proof,
+  Empty,
   inCheckedComputation,
   Proof,
   Prover,
@@ -592,11 +593,7 @@ type LazyProof = {
   kind: 'lazy-proof';
   methodName: string;
   args: any[];
-  previousProofs: {
-    publicInput: Field[];
-    publicOutput: Field[];
-    proof: Pickles.Proof;
-  }[];
+  previousProofs: Pickles.Proof[];
   ZkappClass: typeof SmartContract;
   memoized: { fields: Field[]; aux: any[] }[];
   blindingValue: Field;
@@ -1949,7 +1946,7 @@ async function addMissingProofs(
   { proofsEnabled = true }
 ): Promise<{
   zkappCommand: ZkappCommandProved;
-  proofs: (Proof<ZkappPublicInput, null> | undefined)[];
+  proofs: (Proof<ZkappPublicInput, Empty> | undefined)[];
 }> {
   type AccountUpdateProved = AccountUpdate & {
     lazyAuthorization?: LazySignature;
@@ -2021,7 +2018,7 @@ async function addMissingProofs(
       accountUpdateProved: accountUpdate as AccountUpdateProved,
       proof: new Proof({
         publicInput,
-        publicOutput: null,
+        publicOutput: undefined,
         proof,
         maxProofsVerified,
       }),
@@ -2032,7 +2029,7 @@ async function addMissingProofs(
   // compute proofs serially. in parallel would clash with our global variable
   // hacks
   let accountUpdatesProved: AccountUpdateProved[] = [];
-  let proofs: (Proof<ZkappPublicInput, null> | undefined)[] = [];
+  let proofs: (Proof<ZkappPublicInput, Empty> | undefined)[] = [];
   for (let i = 0; i < accountUpdates.length; i++) {
     let { accountUpdateProved, proof } = await addProof(i, accountUpdates[i]);
     accountUpdatesProved.push(accountUpdateProved);

src/lib/mina.ts

@@ -19,7 +19,7 @@ import {
 import * as Fetch from './fetch.js';
 import { assertPreconditionInvariants, NetworkValue } from './precondition.js';
 import { cloneCircuitValue, toConstant } from './circuit_value.js';
-import { Proof, snarkContext, verify } from './proof_system.js';
+import { Empty, Proof, verify } from './proof_system.js';
 import { Context } from './global-context.js';
 import { SmartContract } from './zkapp.js';
 import { invalidTransactionError } from './errors.js';
@@ -94,7 +94,7 @@ type Transaction = {
    *
    * This can take some time.
    */
-  prove(): Promise<(Proof<ZkappPublicInput, null> | undefined)[]>;
+  prove(): Promise<(Proof<ZkappPublicInput, Empty> | undefined)[]>;
   /**
    * Sends the {@link Transaction} to the network.
    */

src/lib/proof_system.ts

@@ -1,4 +1,8 @@
-import { Empty, EmptyUndefined, EmptyVoid } from '../bindings/lib/generic.js';
+import {
+  EmptyNull,
+  EmptyUndefined,
+  EmptyVoid,
+} from '../bindings/lib/generic.js';
 import { withThreadPool } from '../bindings/js/wrapper.js';
 import {
   Bool,
@@ -20,7 +24,16 @@ import {
 import { Context } from './global-context.js';
 
 // public API
-export { Proof, SelfProof, JsonProof, ZkProgram, verify };
+export {
+  Proof,
+  SelfProof,
+  JsonProof,
+  ZkProgram,
+  verify,
+  Empty,
+  Undefined,
+  Void,
+};
 
 // internal API
 export {
@@ -46,8 +59,6 @@ export {
   inCheckedComputation,
   inCompileMode,
   dummyBase64Proof,
-  Undefined,
-  Void,
 };
 
 // global circuit-related context
@@ -66,6 +77,8 @@ let snarkContext = Context.create<SnarkContext>({ default: {} });
 type Undefined = undefined;
 const Undefined: ProvablePureExtended<undefined, null> =
   EmptyUndefined<Field>();
+type Empty = Undefined;
+const Empty = Undefined;
 type Void = undefined;
 const Void: ProvablePureExtended<void, null> = EmptyVoid<Field>();
 
@@ -80,7 +93,7 @@ class Proof<Input, Output> {
   };
   publicInput: Input;
   publicOutput: Output;
-  proof: RawProof;
+  proof: Pickles.Proof;
   maxProofsVerified: 0 | 1 | 2;
   shouldVerify = Bool(false);
 
@@ -129,7 +142,7 @@ class Proof<Input, Output> {
     publicOutput,
     maxProofsVerified,
   }: {
-    proof: RawProof;
+    proof: Pickles.Proof;
     publicInput: Input;
     publicOutput: Output;
     maxProofsVerified: 0 | 1 | 2;
@@ -145,7 +158,7 @@ async function verify(
   proof: Proof<any, any> | JsonProof,
   verificationKey: string
 ) {
-  let picklesProof: unknown;
+  let picklesProof: Pickles.Proof;
   let statement: Pickles.Statement;
   if (typeof proof.proof === 'string') {
     // json proof
@@ -171,7 +184,6 @@ async function verify(
   );
 }
 
-type RawProof = unknown;
 type JsonProof = {
   publicInput: string[];
   publicOutput: string[];
@@ -260,7 +272,7 @@ function ZkProgram<
         provers: Pickles.Prover[];
         verify: (
           statement: Pickles.Statement,
-          proof: unknown
+          proof: Pickles.Proof
         ) => Promise<boolean>;
       }
     | undefined;
@@ -475,22 +487,13 @@ function isGeneric(type: unknown): type is typeof GenericArgument {
 
 function getPreviousProofsForProver(
   methodArgs: any[],
-  { allArgs, proofArgs }: MethodInterface
+  { allArgs }: MethodInterface
 ) {
-  let previousProofs: Pickles.ProofWithStatement[] = [];
+  let previousProofs: Pickles.Proof[] = [];
   for (let i = 0; i < allArgs.length; i++) {
     let arg = allArgs[i];
     if (arg.type === 'proof') {
-      let { proof, publicInput, publicOutput } = methodArgs[i] as Proof<
-        any,
-        any
-      >;
-      let type = getStatementType(proofArgs[arg.index]);
-      previousProofs[arg.index] = {
-        publicInput: type.input.toFields(publicInput),
-        publicOutput: type.output.toFields(publicOutput),
-        proof,
-      };
+      previousProofs[arg.index] = (methodArgs[i] as Proof<any, any>).proof;
     }
   }
   return previousProofs;
@@ -539,10 +542,10 @@ async function compileProgram(
   );
   // wrap provers
   let wrappedProvers = provers.map(
-    (prover) =>
+    (prover): Pickles.Prover =>
       async function picklesProver(
         publicInput: Field[],
-        previousProofs: Pickles.ProofWithStatement[]
+        previousProofs: Pickles.Proof[]
       ) {
         return withThreadPool(() => prover(publicInput, previousProofs));
       }
@@ -583,38 +586,36 @@ function picklesRuleFromFunction(
   proofSystemTag: { name: string },
   { methodName, witnessArgs, proofArgs, allArgs }: MethodInterface
 ): Pickles.Rule {
-  function main(
-    publicInput: Field[],
-    previousInputsAndOutputs: Field[][]
-  ): ReturnType<Pickles.Rule['main']> {
+  function main(publicInput: Field[]): ReturnType<Pickles.Rule['main']> {
     let { witnesses: argsWithoutPublicInput } = snarkContext.get();
     let finalArgs = [];
     let proofs: Proof<any, any>[] = [];
+    let previousStatements: Pickles.Statement[] = [];
     for (let i = 0; i < allArgs.length; i++) {
       let arg = allArgs[i];
       if (arg.type === 'witness') {
         let type = witnessArgs[arg.index];
         finalArgs[i] = argsWithoutPublicInput
-          ? Circuit.witness(type, () => argsWithoutPublicInput![i])
+          ? Circuit.witness(type, () => argsWithoutPublicInput?.[i])
           : emptyWitness(type);
       } else if (arg.type === 'proof') {
         let Proof = proofArgs[arg.index];
-        // split in input & output
-        // TODO: would be nice to either make Pickles return the right split or completely move
-        // getting the public input & outputs to the JS side, from `snarkContext`, like private inputs
         let type = getStatementType(Proof);
-        let previousStatement = previousInputsAndOutputs[arg.index];
-        let inputFields = previousStatement.slice(0, type.input.sizeInFields());
-        let outputFields = previousStatement.slice(type.input.sizeInFields());
-        let publicInput = type.input.fromFields(inputFields);
-        let publicOutput = type.output.fromFields(outputFields);
-        let proof: unknown;
-        if (argsWithoutPublicInput) {
-          ({ proof } = argsWithoutPublicInput[i] as any);
-        }
+        let proof_ = (argsWithoutPublicInput?.[i] as Proof<any, any>) ?? {
+          proof: undefined,
+          publicInput: emptyValue(type.input),
+          publicOutput: emptyValue(type.output),
+        };
+        let { proof, publicInput, publicOutput } = proof_;
+        publicInput = Circuit.witness(type.input, () => publicInput);
+        publicOutput = Circuit.witness(type.output, () => publicOutput);
         let proofInstance = new Proof({ publicInput, publicOutput, proof });
         finalArgs[i] = proofInstance;
         proofs.push(proofInstance);
+        previousStatements.push({
+          input: type.input.toFields(publicInput),
+          output: type.output.toFields(publicOutput),
+        });
       } else if (arg.type === 'generic') {
         finalArgs[i] = argsWithoutPublicInput?.[i] ?? emptyGeneric();
       }
@@ -630,6 +631,7 @@ function picklesRuleFromFunction(
     let publicOutput = hasPublicOutput ? publicOutputType.toFields(result) : [];
     return {
       publicOutput,
+      previousStatements,
       shouldVerify: proofs.map((proof) => proof.shouldVerify),
     };
   }
@@ -704,7 +706,7 @@ function methodArgumentsToConstant(
   return constArgs;
 }
 
-let Generic = Empty<Field>();
+let Generic = EmptyNull<Field>();
 
 type TypeAndValue<T> = { type: Provable<T>; value: T };
 

src/lib/zkapp.ts

@@ -52,6 +52,7 @@ import {
 import {
   analyzeMethod,
   compileProgram,
+  Empty,
   emptyValue,
   GenericArgument,
   getPreviousProofsForProver,
@@ -68,7 +69,6 @@ import {
 } from './proof_system.js';
 import { PrivateKey, PublicKey } from './signature.js';
 import { assertStatePrecondition, cleanStatePrecondition } from './state.js';
-import { Empty } from '../bindings/lib/generic.js';
 
 // external API
 export {
@@ -120,9 +120,9 @@ function method<T extends SmartContract>(
     methodName
   );
 
-  class SelfProof extends Proof<ZkappPublicInput, null> {
+  class SelfProof extends Proof<ZkappPublicInput, Empty> {
     static publicInputType = ZkappPublicInput;
-    static publicOutputType = Empty<Field>();
+    static publicOutputType = Empty;
     static tag = () => ZkappClass;
   }
   let internalMethodEntry = sortMethodArguments(
@@ -628,9 +628,9 @@ class SmartContract {
    */
   static Proof() {
     let Contract = this;
-    return class extends Proof<ZkappPublicInput, null> {
+    return class extends Proof<ZkappPublicInput, Empty> {
       static publicInputType = ZkappPublicInput;
-      static publicOutputType = Empty<Field>();
+      static publicOutputType = Empty;
       static tag = () => Contract;
     };
   }
@@ -683,7 +683,7 @@ class SmartContract {
       verify,
     } = await compileProgram(
       ZkappPublicInput,
-      Empty(),
+      Empty,
       methodIntfs,
       methods,
       this