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

CHANGELOG-v2.md

@@ -30,10 +30,12 @@ Showing all changes since the last release (v.1.5.0)
 
   - Removed implicit version `Provable.equal(x, y)` where you didn't have to pass in the type
 
-- `zkProgram` to support `auxiliaryOutput`
+- `ZkProgram` to support `auxiliaryOutput`
 
 ### Added
 
+-`ZkProgram` to support non-pure provable types as inputs and outputs https://github.com/o1-labs/o1js/pull/1828
+
 - New method `toCanonical()` in the `Provable<T>` interface to protect against incompleteness of certain operations on malicious witness inputs https://github.com/o1-labs/o1js/pull/1759
 - `divMod64()` division modulo 2^64 that returns the remainder and quotient of the operation
 - `addMod64()` addition modulo 2^64

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 = 'inCircuit';
+        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 === 'inCircuit');

src/lib/proof-system/zkprogram.ts

@@ -162,19 +162,36 @@ const FeatureFlags = {
 
 function createProgramState() {
   let methodCache: Map<string, unknown> = new Map();
-
   return {
+    setNonPureInput(value: any[]) {
+      methodCache.set('nonPureInput', value);
+    },
+    getNonPureInput(): any[] {
+      let entry = methodCache.get('nonPureInput');
+      if (entry === undefined) throw Error(`Non-pure input not defined`);
+      return entry as any[];
+    },
+    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);
     },
   };
 }
@@ -234,8 +251,8 @@ async function featureFlagsfromFlatMethodIntfs(
 }
 
 class ProofBase<Input, Output> {
-  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` +
@@ -558,8 +575,8 @@ let SideloadedTag = {
 
 function ZkProgram<
   Config extends {
-    publicInput?: ProvableTypePure;
-    publicOutput?: ProvableTypePure;
+    publicInput?: ProvableType;
+    publicOutput?: ProvableType;
     methods: {
       [I in string]: {
         privateInputs: Tuple<PrivateInput>;
@@ -635,10 +652,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
   );
 
@@ -778,12 +795,35 @@ 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
+        ));
+
+        // store publicInput auxilary data in programState cache
+        programState.setNonPureInput(publicInputAux);
+      } else {
+        publicInputFields = toFieldConsts(publicInputType, publicInput);
+      }
+
       let previousProofs = MlArray.to(
         getPreviousProofsForProver(args, methodIntfs[i])
       );
 
-      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);
@@ -804,8 +844,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({
@@ -1062,8 +1115,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[][];
@@ -1175,7 +1228,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
 ) {
@@ -1203,8 +1256,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 },
   {
@@ -1220,7 +1273,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: {
@@ -1266,10 +1323,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(({ proofInstance, classReference }) => {
       if (!(proofInstance instanceof DynamicProof)) {
         return;
@@ -1299,7 +1368,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)
@@ -1487,23 +1556,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 toFieldVars<T>(type: ProvablePure<T>, value: T) {
   return MlFieldArray.to(type.toFields(value));
 }
 
-function fromFieldConsts<T>(type: ProvablePure<T>, fields: MlFieldConstArray) {
-  return type.fromFields(MlFieldConstArray.from(fields));
+function fromFieldConsts<T>(
+  type: Provable<T>,
+  fields: MlFieldConstArray,
+  aux: any[] = []
+) {
+  return type.fromFields(MlFieldConstArray.from(fields), aux);
 }
-function toFieldConsts<T>(type: ProvablePure<T>, value: T) {
+
+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: {} });