@@ -3,7 +3,7 @@ import { Snarky, initializeBindings, withThreadPool } from '../../snarky.js';
33import { Pickles , Gate } from '../../snarky.js' ;
44import { Field } from '../provable/wrapped.js' ;
55import {
6- FlexibleProvablePure ,
6+ FlexibleProvable ,
77 InferProvable ,
88 ProvablePureExtended ,
99 Struct ,
@@ -89,8 +89,15 @@ const Void: ProvablePureExtended<void, void, null> = EmptyVoid<Field>();
8989
9090function createProgramState ( ) {
9191 let methodCache : Map < string , unknown > = new Map ( ) ;
92-
9392 return {
93+ setNonPureOutput ( value : any [ ] ) {
94+ methodCache . set ( '__nonPureOutput__' , value ) ;
95+ } ,
96+ getNonPureOutput ( ) : any [ ] {
97+ let entry = methodCache . get ( '__nonPureOutput__' ) ;
98+ if ( entry === undefined ) return [ ] ;
99+ return entry as any [ ] ;
100+ } ,
94101 setAuxiliaryOutput ( value : unknown , methodName : string ) {
95102 methodCache . set ( methodName , value ) ;
96103 } ,
@@ -100,8 +107,8 @@ function createProgramState() {
100107 throw Error ( `Auxiliary value for method ${ methodName } not defined` ) ;
101108 return entry ;
102109 } ,
103- reset ( methodName : string ) {
104- methodCache . delete ( methodName ) ;
110+ reset ( key : string ) {
111+ methodCache . delete ( key ) ;
105112 } ,
106113 } ;
107114}
@@ -173,8 +180,8 @@ let SideloadedTag = {
173180
174181function ZkProgram <
175182 Config extends {
176- publicInput ?: ProvableTypePure ;
177- publicOutput ?: ProvableTypePure ;
183+ publicInput ?: ProvableType ;
184+ publicOutput ?: ProvableType ;
178185 methods : {
179186 [ I in string ] : {
180187 privateInputs : Tuple < PrivateInput > ;
@@ -250,10 +257,10 @@ function ZkProgram<
250257 let doProving = true ;
251258
252259 let methods = config . methods ;
253- let publicInputType : ProvablePure < any > = ProvableType . get (
260+ let publicInputType : Provable < any > = ProvableType . get (
254261 config . publicInput ?? Undefined
255262 ) ;
256- let publicOutputType : ProvablePure < any > = ProvableType . get (
263+ let publicOutputType : Provable < any > = ProvableType . get (
257264 config . publicOutput ?? Void
258265 ) ;
259266
@@ -391,10 +398,20 @@ function ZkProgram<
391398 `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.`
392399 ) ;
393400 }
394- let publicInputFields = toFieldConsts ( publicInputType , publicInput ) ;
401+
402+ let { publicInputFields, publicInputAux } = toFieldAndAuxConsts (
403+ publicInputType ,
404+ publicInput
405+ ) ;
406+
395407 let previousProofs = MlArray . to ( getPreviousProofsForProver ( args ) ) ;
396408
397- let id = snarkContext . enter ( { witnesses : args , inProver : true } ) ;
409+ let id = snarkContext . enter ( {
410+ witnesses : args ,
411+ inProver : true ,
412+ auxInputData : publicInputAux ,
413+ } ) ;
414+
398415 let result : UnwrapPromise < ReturnType < typeof picklesProver > > ;
399416 try {
400417 result = await picklesProver ( publicInputFields , previousProofs ) ;
@@ -416,7 +433,16 @@ function ZkProgram<
416433 }
417434
418435 let [ publicOutputFields , proof ] = MlPair . from ( result ) ;
419- let publicOutput = fromFieldConsts ( publicOutputType , publicOutputFields ) ;
436+
437+ let nonPureOutput = programState . getNonPureOutput ( ) ;
438+
439+ let publicOutput = fromFieldConsts (
440+ publicOutputType ,
441+ publicOutputFields ,
442+ nonPureOutput
443+ ) ;
444+
445+ programState . reset ( '__nonPureOutput__' ) ;
420446
421447 return {
422448 proof : new ProgramProof ( {
@@ -649,8 +675,8 @@ async function compileProgram({
649675 overrideWrapDomain,
650676 state,
651677} : {
652- publicInputType : ProvablePure < any > ;
653- publicOutputType : ProvablePure < any > ;
678+ publicInputType : Provable < any > ;
679+ publicOutputType : Provable < any > ;
654680 methodIntfs : MethodInterface [ ] ;
655681 methods : ( ( ...args : any ) => unknown ) [ ] ;
656682 gates : Gate [ ] [ ] ;
@@ -762,7 +788,7 @@ If you are using a SmartContract, make sure you are using the @method decorator.
762788}
763789
764790function analyzeMethod (
765- publicInputType : ProvablePure < any > ,
791+ publicInputType : Provable < any > ,
766792 methodIntf : MethodInterface ,
767793 method : ( ...args : any ) => unknown
768794) {
@@ -790,8 +816,8 @@ function inCircuitVkHash(inCircuitVk: unknown): Field {
790816}
791817
792818function picklesRuleFromFunction (
793- publicInputType : ProvablePure < unknown > ,
794- publicOutputType : ProvablePure < unknown > ,
819+ publicInputType : Provable < unknown > ,
820+ publicOutputType : Provable < unknown > ,
795821 func : ( ...args : unknown [ ] ) => unknown ,
796822 proofSystemTag : { name : string } ,
797823 { methodName, args, auxiliaryType } : MethodInterface ,
@@ -801,7 +827,11 @@ function picklesRuleFromFunction(
801827 async function main (
802828 publicInput : MlFieldArray
803829 ) : ReturnType < Pickles . Rule [ 'main' ] > {
804- let { witnesses : argsWithoutPublicInput , inProver } = snarkContext . get ( ) ;
830+ let {
831+ witnesses : argsWithoutPublicInput ,
832+ inProver,
833+ auxInputData,
834+ } = snarkContext . get ( ) ;
805835 assert ( ! ( inProver && argsWithoutPublicInput === undefined ) ) ;
806836 let finalArgs = [ ] ;
807837 let proofs : {
@@ -837,10 +867,16 @@ function picklesRuleFromFunction(
837867 if ( publicInputType === Undefined || publicInputType === Void ) {
838868 result = ( await func ( ...finalArgs ) ) as any ;
839869 } else {
840- let input = fromFieldVars ( publicInputType , publicInput ) ;
870+ let input = fromFieldVars ( publicInputType , publicInput , auxInputData ) ;
841871 result = ( await func ( input , ...finalArgs ) ) as any ;
842872 }
843873
874+ if ( result ?. publicOutput ) {
875+ // store the nonPure auxiliary data in program state cache if it exists
876+ let nonPureOutput = publicOutputType . toAuxiliary ( result . publicOutput ) ;
877+ state ?. setNonPureOutput ( nonPureOutput ) ;
878+ }
879+
844880 proofs . forEach ( ( { Proof, proof } ) => {
845881 if ( ! ( proof instanceof DynamicProof ) ) return ;
846882
@@ -869,7 +905,7 @@ function picklesRuleFromFunction(
869905 Pickles . sideLoaded . inCircuit ( computedTag , circuitVk ) ;
870906 } ) ;
871907
872- // if the public output is empty, we don't evaluate `toFields(result)` to allow the function to return something else in that case
908+ // if the output is empty, we don't evaluate `toFields(result)` to allow the function to return something else in that case
873909 let hasPublicOutput = publicOutputType . sizeInFields ( ) !== 0 ;
874910 let publicOutput = hasPublicOutput
875911 ? publicOutputType . toFields ( result . publicOutput )
@@ -957,20 +993,36 @@ function getMaxProofsVerified(methodIntfs: MethodInterface[]) {
957993 ) as any as 0 | 1 | 2 ;
958994}
959995
960- function fromFieldVars < T > ( type : ProvablePure < T > , fields : MlFieldArray ) {
961- return type . fromFields ( MlFieldArray . from ( fields ) ) ;
996+ function fromFieldVars < T > (
997+ type : Provable < T > ,
998+ fields : MlFieldArray ,
999+ auxData : any [ ] = [ ]
1000+ ) {
1001+ return type . fromFields ( MlFieldArray . from ( fields ) , auxData ) ;
9621002}
9631003
964- function fromFieldConsts < T > ( type : ProvablePure < T > , fields : MlFieldConstArray ) {
965- return type . fromFields ( MlFieldConstArray . from ( fields ) ) ;
1004+ function fromFieldConsts < T > (
1005+ type : Provable < T > ,
1006+ fields : MlFieldConstArray ,
1007+ aux : any [ ] = [ ]
1008+ ) {
1009+ return type . fromFields ( MlFieldConstArray . from ( fields ) , aux ) ;
9661010}
967- function toFieldConsts < T > ( type : ProvablePure < T > , value : T ) {
1011+
1012+ function toFieldConsts < T > ( type : Provable < T > , value : T ) {
9681013 return MlFieldConstArray . to ( type . toFields ( value ) ) ;
9691014}
9701015
1016+ function toFieldAndAuxConsts < T > ( type : Provable < T > , value : T ) {
1017+ return {
1018+ publicInputFields : MlFieldConstArray . to ( type . toFields ( value ) ) ,
1019+ publicInputAux : type . toAuxiliary ( value ) ,
1020+ } ;
1021+ }
1022+
9711023ZkProgram . Proof = function <
972- PublicInputType extends FlexibleProvablePure < any > ,
973- PublicOutputType extends FlexibleProvablePure < any >
1024+ PublicInputType extends FlexibleProvable < any > ,
1025+ PublicOutputType extends FlexibleProvable < any >
9741026> ( program : {
9751027 name : string ;
9761028 publicInputType : PublicInputType ;
0 commit comments