Allow non-pure provable types as inputs/outputs in zkprogram

CHANGELOG.md

@@ -19,6 +19,8 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Added
 
+-`ZkProgram` to support non-pure provable types as inputs and outputs https://github.com/o1-labs/o1js/pull/1828
+
 - Support secp256r1 in elliptic curve and ECDSA gadgets https://github.com/o1-labs/o1js/pull/1885
 
 ### Fixed

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

@@ -0,0 +1,39 @@
+import { Field, Provable, Struct, ZkProgram, assert } from 'o1js';
+
+class MyStruct extends Struct({
+  label: String,
+  value: Field,
+}) {}
+
+let MyProgram = ZkProgram({
+  name: 'example-with-non-pure-inputs',
+  publicInput: MyStruct,
+  publicOutput: MyStruct,
+
+  methods: {
+    baseCase: {
+      privateInputs: [],
+      async method(input: MyStruct) {
+        //update input in circuit
+        input.label = 'in-circuit';
+        return {
+          publicOutput: input,
+        };
+      },
+    },
+  },
+});
+
+//
+
+console.log('compiling MyProgram...');
+await MyProgram.compile();
+console.log('compile done');
+
+let input = new MyStruct({ label: 'input', value: Field(5) });
+
+let { proof } = await MyProgram.baseCase(input);
+let ok = await MyProgram.verify(proof);
+
+assert(ok, 'proof not valid!');
+assert(proof.publicOutput.label === 'in-circuit');

src/lib/proof-system/proof.ts

@@ -2,7 +2,7 @@ import { initializeBindings, withThreadPool } from '../../snarky.js';
 import { Pickles } from '../../snarky.js';
 import { Field, Bool } from '../provable/wrapped.js';
 import type {
-  FlexibleProvablePure,
+  FlexibleProvable,
   InferProvable,
 } from '../provable/types/struct.js';
 import { FeatureFlags } from './feature-flags.js';
@@ -22,8 +22,8 @@ export { dummyProof, extractProofs, extractProofTypes, type ProofValue };
 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 publicInputType: FlexibleProvable<any> = undefined as any;
+  static publicOutputType: FlexibleProvable<any> = undefined as any;
   static tag: () => { name: string } = () => {
     throw Error(
       `You cannot use the \`Proof\` class directly. Instead, define a subclass:\n` +

src/lib/proof-system/zkprogram.ts

@@ -3,7 +3,7 @@ import { Snarky, initializeBindings, withThreadPool } from '../../snarky.js';
 import { Pickles, Gate } from '../../snarky.js';
 import { Field } from '../provable/wrapped.js';
 import {
-  FlexibleProvablePure,
+  FlexibleProvable,
   InferProvable,
   ProvablePureExtended,
   Struct,
@@ -89,19 +89,28 @@ const Void: ProvablePureExtended<void, void, null> = EmptyVoid<Field>();
 
 function createProgramState() {
   let methodCache: Map<string, unknown> = new Map();
-
   return {
+    setNonPureOutput(value: any[]) {
+      methodCache.set('__nonPureOutput__', value);
+    },
+    getNonPureOutput(): any[] {
+      let entry = methodCache.get('__nonPureOutput__');
+      if (entry === undefined) throw Error(`Non-pure output not defined`);
+      return entry as any[];
+    },
+
     setAuxiliaryOutput(value: unknown, methodName: string) {
       methodCache.set(methodName, value);
     },
+
     getAuxiliaryOutput(methodName: string): unknown {
       let entry = methodCache.get(methodName);
       if (entry === undefined)
         throw Error(`Auxiliary value for method ${methodName} not defined`);
       return entry;
     },
-    reset(methodName: string) {
-      methodCache.delete(methodName);
+    reset(key: string) {
+      methodCache.delete(key);
     },
   };
 }
@@ -173,8 +182,8 @@ let SideloadedTag = {
 
 function ZkProgram<
   Config extends {
-    publicInput?: ProvableTypePure;
-    publicOutput?: ProvableTypePure;
+    publicInput?: ProvableType;
+    publicOutput?: ProvableType;
     methods: {
       [I in string]: {
         privateInputs: Tuple<PrivateInput>;
@@ -250,10 +259,10 @@ function ZkProgram<
   let doProving = true;
 
   let methods = config.methods;
-  let publicInputType: ProvablePure<any> = ProvableType.get(
+  let publicInputType: Provable<any> = ProvableType.get(
     config.publicInput ?? Undefined
   );
-  let publicOutputType: ProvablePure<any> = ProvableType.get(
+  let publicOutputType: Provable<any> = ProvableType.get(
     config.publicOutput ?? Void
   );
 
@@ -391,10 +400,30 @@ function ZkProgram<
             `Try calling \`await program.compile()\` first, this will cache provers in the background.\nIf you compiled your zkProgram with proofs disabled (\`proofsEnabled = false\`), you have to compile it with proofs enabled first.`
         );
       }
-      let publicInputFields = toFieldConsts(publicInputType, publicInput);
+
+      let nonPureInputExists =
+        publicInputType.toAuxiliary(publicInput).length !== 0;
+      console.log('nonPure Input Exists', nonPureInputExists);
+      let publicInputFields, publicInputAux;
+      if (nonPureInputExists) {
+        // serialize publicInput into pure provable field elements and auxilary data
+        ({ publicInputFields, publicInputAux } = toFieldAndAuxConsts(
+          publicInputType,
+          publicInput
+        ));
+      } else {
+        publicInputFields = toFieldConsts(publicInputType, publicInput);
+      }
+
       let previousProofs = MlArray.to(getPreviousProofsForProver(args));
 
-      let id = snarkContext.enter({ witnesses: args, inProver: true });
+      console.log('auxdata before entering snarkContext ', publicInputAux);
+      let id = snarkContext.enter({
+        witnesses: args,
+        inProver: true,
+        auxInputData: publicInputAux,
+      });
+
       let result: UnwrapPromise<ReturnType<typeof picklesProver>>;
       try {
         result = await picklesProver(publicInputFields, previousProofs);
@@ -415,8 +444,21 @@ function ZkProgram<
         programState.reset(methodIntfs[i].methodName);
       }
 
+      let publicOutput;
       let [publicOutputFields, proof] = MlPair.from(result);
-      let publicOutput = fromFieldConsts(publicOutputType, publicOutputFields);
+      if (nonPureInputExists) {
+        let nonPureOutput = programState.getNonPureOutput();
+
+        publicOutput = fromFieldConsts(
+          publicOutputType,
+          publicOutputFields,
+          nonPureOutput
+        );
+
+        programState.reset('nonPureOutput');
+      } else {
+        publicOutput = fromFieldConsts(publicOutputType, publicOutputFields);
+      }
 
       return {
         proof: new ProgramProof({
@@ -649,8 +691,8 @@ async function compileProgram({
   overrideWrapDomain,
   state,
 }: {
-  publicInputType: ProvablePure<any>;
-  publicOutputType: ProvablePure<any>;
+  publicInputType: Provable<any>;
+  publicOutputType: Provable<any>;
   methodIntfs: MethodInterface[];
   methods: ((...args: any) => unknown)[];
   gates: Gate[][];
@@ -762,7 +804,7 @@ If you are using a SmartContract, make sure you are using the @method decorator.
 }
 
 function analyzeMethod(
-  publicInputType: ProvablePure<any>,
+  publicInputType: Provable<any>,
   methodIntf: MethodInterface,
   method: (...args: any) => unknown
 ) {
@@ -790,8 +832,8 @@ function inCircuitVkHash(inCircuitVk: unknown): Field {
 }
 
 function picklesRuleFromFunction(
-  publicInputType: ProvablePure<unknown>,
-  publicOutputType: ProvablePure<unknown>,
+  publicInputType: Provable<unknown>,
+  publicOutputType: Provable<unknown>,
   func: (...args: unknown[]) => unknown,
   proofSystemTag: { name: string },
   { methodName, args, auxiliaryType }: MethodInterface,
@@ -801,7 +843,11 @@ function picklesRuleFromFunction(
   async function main(
     publicInput: MlFieldArray
   ): ReturnType<Pickles.Rule['main']> {
-    let { witnesses: argsWithoutPublicInput, inProver } = snarkContext.get();
+    let {
+      witnesses: argsWithoutPublicInput,
+      inProver,
+      auxInputData,
+    } = snarkContext.get();
     assert(!(inProver && argsWithoutPublicInput === undefined));
     let finalArgs = [];
     let proofs: {
@@ -837,10 +883,22 @@ function picklesRuleFromFunction(
     if (publicInputType === Undefined || publicInputType === Void) {
       result = (await func(...finalArgs)) as any;
     } else {
-      let input = fromFieldVars(publicInputType, publicInput);
+      console.log('auxData before input', auxInputData);
+      let input = fromFieldVars(publicInputType, publicInput, auxInputData);
       result = (await func(input, ...finalArgs)) as any;
     }
 
+    console.log('result input', result);
+    if (result?.publicOutput) {
+      // store the nonPure auxiliary data in program state cache if it exists
+      let nonPureOutput = publicOutputType.toAuxiliary(result.publicOutput);
+      let nonPureOutputExists = nonPureOutput.length !== 0;
+
+      if (state !== undefined && nonPureOutputExists) {
+        state.setNonPureOutput(nonPureOutput);
+      }
+    }
+
     proofs.forEach(({ Proof, proof }) => {
       if (!(proof instanceof DynamicProof)) return;
 
@@ -869,7 +927,7 @@ function picklesRuleFromFunction(
       Pickles.sideLoaded.inCircuit(computedTag, circuitVk);
     });
 
-    // if the public output is empty, we don't evaluate `toFields(result)` to allow the function to return something else in that case
+    // if the output is empty, we don't evaluate `toFields(result)` to allow the function to return something else in that case
     let hasPublicOutput = publicOutputType.sizeInFields() !== 0;
     let publicOutput = hasPublicOutput
       ? publicOutputType.toFields(result.publicOutput)
@@ -957,20 +1015,40 @@ function getMaxProofsVerified(methodIntfs: MethodInterface[]) {
   ) as any as 0 | 1 | 2;
 }
 
-function fromFieldVars<T>(type: ProvablePure<T>, fields: MlFieldArray) {
-  return type.fromFields(MlFieldArray.from(fields));
+function fromFieldVars<T>(
+  type: Provable<T>,
+  fields: MlFieldArray,
+  auxData: any[] = []
+) {
+  return type.fromFields(MlFieldArray.from(fields), auxData);
 }
 
-function fromFieldConsts<T>(type: ProvablePure<T>, fields: MlFieldConstArray) {
-  return type.fromFields(MlFieldConstArray.from(fields));
+function toFieldVars<T>(type: ProvablePure<T>, value: T) {
+  return MlFieldArray.to(type.toFields(value));
 }
-function toFieldConsts<T>(type: ProvablePure<T>, value: T) {
+
+function fromFieldConsts<T>(
+  type: Provable<T>,
+  fields: MlFieldConstArray,
+  aux: any[] = []
+) {
+  return type.fromFields(MlFieldConstArray.from(fields), aux);
+}
+
+function toFieldConsts<T>(type: Provable<T>, value: T) {
   return MlFieldConstArray.to(type.toFields(value));
 }
 
+function toFieldAndAuxConsts<T>(type: Provable<T>, value: T) {
+  return {
+    publicInputFields: MlFieldConstArray.to(type.toFields(value)),
+    publicInputAux: type.toAuxiliary(value),
+  };
+}
+
 ZkProgram.Proof = function <
-  PublicInputType extends FlexibleProvablePure<any>,
-  PublicOutputType extends FlexibleProvablePure<any>
+  PublicInputType extends FlexibleProvable<any>,
+  PublicOutputType extends FlexibleProvable<any>
 >(program: {
   name: string;
   publicInputType: PublicInputType;

src/lib/provable/core/provable-context.ts

@@ -40,6 +40,7 @@ type SnarkContext = {
   inCheckedComputation?: boolean;
   inAnalyze?: boolean;
   inWitnessBlock?: boolean;
+  auxInputData?: any[];
 };
 let snarkContext = Context.create<SnarkContext>({ default: {} });