Merge pull request #1828 from o1-labs/feature/non-pure-zkprogram-inputs

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

Author: ymekuria

CHANGELOG.md

@@ -17,24 +17,21 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased](https://github.com/o1-labs/o1js/compare/e1bac02...HEAD)
 
-
 ### Added
 
+-`ZkProgram` to support non-pure provable types as inputs and outputs https://github.com/o1-labs/o1js/pull/1828
+
 - Add `enforceTransactionLimits` parameter on Network https://github.com/o1-labs/o1js/issues/1910
 
+- Method for optional types to assert none https://github.com/o1-labs/o1js/pull/1922
+
 ### Fixed
 
 - Compiling stuck in the browser for recursive zkprograms https://github.com/o1-labs/o1js/pull/1906
 - Error message in `rangeCheck16` gadget https://github.com/o1-labs/o1js/pull/1920
 
-### Added
-
-- Method for optional types to assert none https://github.com/o1-labs/o1js/pull/1922
-
 ## [2.1.0](https://github.com/o1-labs/o1js/compare/b04520d...e1bac02) - 2024-11-13
 
-### Added
-
 - 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,62 @@
+import { Field, Struct, ZkProgram, assert } from 'o1js';
+
+class MyStruct extends Struct({
+  label: String,
+  value: Field,
+}) {}
+
+let NonPureIOprogram = ZkProgram({
+  name: 'example-with-non-pure-io',
+  publicInput: MyStruct,
+  publicOutput: MyStruct,
+
+  methods: {
+    baseCase: {
+      privateInputs: [],
+      async method(input: MyStruct) {
+        //update input in circuit
+        input.label = 'in-circuit';
+        return {
+          publicOutput: input,
+        };
+      },
+    },
+  },
+});
+
+let NonPureOutputProgram = ZkProgram({
+  name: 'example-with-non-pure-output',
+  publicOutput: MyStruct,
+
+  methods: {
+    baseCase: {
+      privateInputs: [],
+      async method() {
+        return {
+          publicOutput: new MyStruct({ label: 'output', value: Field(5) }),
+        };
+      },
+    },
+  },
+});
+
+console.log('compiling NonPureIOprogram...');
+await NonPureIOprogram.compile();
+console.log('compile done');
+let input = new MyStruct({ label: 'input', value: Field(5) });
+let proof;
+({ proof } = await NonPureIOprogram.baseCase(input));
+let isProof1Valid = await NonPureIOprogram.verify(proof);
+assert(isProof1Valid, 'proof not valid!');
+assert(proof.publicOutput.label === 'in-circuit');
+console.log('i/o proof', proof);
+
+console.log('compiling NonPureOutputProgram...');
+await NonPureOutputProgram.compile();
+console.log('compile done');
+
+({ proof } = await NonPureOutputProgram.baseCase());
+let isProof2Valid = await NonPureOutputProgram.verify(proof);
+assert(isProof2Valid, 'proof not valid!');
+assert(proof.publicOutput.label === 'output');
+console.log('output proof', proof);

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,8 +89,15 @@ 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) return [];
+      return entry as any[];
+    },
     setAuxiliaryOutput(value: unknown, methodName: string) {
       methodCache.set(methodName, value);
     },
@@ -100,8 +107,8 @@ function createProgramState() {
         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 +180,8 @@ let SideloadedTag = {
 
 function ZkProgram<
   Config extends {
-    publicInput?: ProvableTypePure;
-    publicOutput?: ProvableTypePure;
+    publicInput?: ProvableType;
+    publicOutput?: ProvableType;
     methods: {
       [I in string]: {
         privateInputs: Tuple<PrivateInput>;
@@ -250,10 +257,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 +398,20 @@ 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 { publicInputFields, publicInputAux } = toFieldAndAuxConsts(
+        publicInputType,
+        publicInput
+      );
+
       let previousProofs = MlArray.to(getPreviousProofsForProver(args));
 
-      let id = snarkContext.enter({ witnesses: args, inProver: true });
+      let id = snarkContext.enter({
+        witnesses: args,
+        inProver: true,
+        auxInputData: publicInputAux,
+      });
+
       let result: UnwrapPromise<ReturnType<typeof picklesProver>>;
       try {
         result = await picklesProver(publicInputFields, previousProofs);
@@ -416,7 +433,16 @@ function ZkProgram<
       }
 
       let [publicOutputFields, proof] = MlPair.from(result);
-      let publicOutput = fromFieldConsts(publicOutputType, publicOutputFields);
+
+      let nonPureOutput = programState.getNonPureOutput();
+
+      let publicOutput = fromFieldConsts(
+        publicOutputType,
+        publicOutputFields,
+        nonPureOutput
+      );
+
+      programState.reset('__nonPureOutput__');
 
       return {
         proof: new ProgramProof({
@@ -649,8 +675,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 +788,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 +816,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 +827,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 +867,16 @@ function picklesRuleFromFunction(
     if (publicInputType === Undefined || publicInputType === Void) {
       result = (await func(...finalArgs)) as any;
     } else {
-      let input = fromFieldVars(publicInputType, publicInput);
+      let input = fromFieldVars(publicInputType, publicInput, auxInputData);
       result = (await func(input, ...finalArgs)) as any;
     }
 
+    if (result?.publicOutput) {
+      // store the nonPure auxiliary data in program state cache if it exists
+      let nonPureOutput = publicOutputType.toAuxiliary(result.publicOutput);
+      state?.setNonPureOutput(nonPureOutput);
+    }
+
     proofs.forEach(({ Proof, proof }) => {
       if (!(proof instanceof DynamicProof)) return;
 
@@ -869,7 +905,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 +993,36 @@ 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 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: {} });