Specify feature flags for Pickles.sideLoaded

CHANGELOG.md

@@ -17,6 +17,11 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased](https://github.com/o1-labs/o1js/compare/40c597775...HEAD)
 
+### Added
+
+- Added the option to specify custom feature flags for sided loaded proofs in the `DynamicProof` class.
+  - Feature flags are requires to tell Pickles what proof structure it should expect when side loading dynamic proofs and verification keys. https://github.com/o1-labs/o1js/pull/1688
+
 ## [1.3.1](https://github.com/o1-labs/o1js/compare/1ad7333e9e...40c597775) - 2024-06-11
 
 ### Breaking Changes

src/examples/zkprogram/dynamic-keys-merkletree.ts

@@ -1,5 +1,6 @@
 import {
   DynamicProof,
+  FeatureFlags,
   Field,
   MerkleTree,
   MerkleWitness,
@@ -30,13 +31,26 @@ const sideloadedProgram = ZkProgram({
         return publicInput.add(privateInput);
       },
     },
+    assertAndAdd: {
+      privateInputs: [Field],
+      async method(publicInput: Field, privateInput: Field) {
+        // this uses assert to test range check gates and their feature flags
+        publicInput.assertLessThanOrEqual(privateInput);
+        return publicInput.add(privateInput);
+      },
+    },
   },
 });
 
+const featureFlags = await FeatureFlags.fromZkProgram(sideloadedProgram);
+
 class SideloadedProgramProof extends DynamicProof<Field, Field> {
   static publicInputType = Field;
   static publicOutputType = Field;
   static maxProofsVerified = 0 as const;
+
+  // we configure this side loaded proof to only accept proofs that doesn't use any feature flags (custom gates)
+  static featureFlags = featureFlags;
 }
 
 const tree = new MerkleTree(64);
@@ -116,7 +130,7 @@ console.log('Compiling circuits...');
 const programVk = (await sideloadedProgram.compile()).verificationKey;
 const mainVk = (await mainProgram.compile()).verificationKey;
 
-console.log('Proving deployment of sideloaded key');
+console.log('Proving deployment of side-loaded key');
 const rootBefore = tree.getRoot();
 tree.setLeaf(1n, programVk.hash);
 const witness = new MerkleTreeWitness(tree.getWitness(1n));
@@ -144,3 +158,18 @@ const proof2 = await mainProgram.validateUsingTree(
 
 const validProof2 = await verify(proof2, mainVk);
 console.log('ok?', validProof2);
+
+console.log('Proving different method of child program');
+const childProof2 = await sideloadedProgram.assertAndAdd(Field(0), Field(10));
+
+console.log('Proving verification inside main program');
+const proof3 = await mainProgram.validateUsingTree(
+  proof1.publicOutput,
+  proof1,
+  programVk,
+  witness,
+  SideloadedProgramProof.fromProof(childProof)
+);
+
+const validProof3 = await verify(proof2, mainVk);
+console.log('ok?', validProof2);

src/index.ts

@@ -83,6 +83,7 @@ export {
   Undefined,
   Void,
   VerificationKey,
+  FeatureFlags,
 } from './lib/proof-system/zkprogram.js';
 export { Cache, CacheHeader } from './lib/proof-system/cache.js';
 

src/lib/ml/base.ts

@@ -14,6 +14,7 @@ export {
   MlUnit,
   MlString,
   MlTuple,
+  MlArrayOptionalElements,
 };
 
 // ocaml types
@@ -26,6 +27,11 @@ type MlBool = 0 | 1;
 type MlResult<T, E> = [0, T] | [1, E];
 type MlUnit = 0;
 
+// custom types
+type MlArrayOptionalElements<MlArray extends any[]> = {
+  [K in keyof MlArray]: MlArray[K] extends 0 ? 0 : MlOption<MlArray[K]>;
+};
+
 /**
  * js_of_ocaml representation of a byte array,
  * see https://github.com/ocsigen/js_of_ocaml/blob/master/runtime/mlBytes.js

src/lib/proof-system/zkprogram.ts

@@ -4,13 +4,7 @@ import {
   EmptyVoid,
 } from '../../bindings/lib/generic.js';
 import { Snarky, initializeBindings, withThreadPool } from '../../snarky.js';
-import {
-  Pickles,
-  FeatureFlags,
-  MlFeatureFlags,
-  Gate,
-  GateType,
-} from '../../snarky.js';
+import { Pickles, MlFeatureFlags, Gate, GateType } from '../../snarky.js';
 import { Field, Bool } from '../provable/wrapped.js';
 import {
   FlexibleProvable,
@@ -24,7 +18,14 @@ import { Provable } from '../provable/provable.js';
 import { assert, prettifyStacktracePromise } from '../util/errors.js';
 import { snarkContext } from '../provable/core/provable-context.js';
 import { hashConstant } from '../provable/crypto/poseidon.js';
-import { MlArray, MlBool, MlResult, MlPair } from '../ml/base.js';
+import {
+  MlArray,
+  MlBool,
+  MlResult,
+  MlPair,
+  MlOption,
+  MlArrayOptionalElements,
+} from '../ml/base.js';
 import { MlFieldArray, MlFieldConstArray } from '../ml/fields.js';
 import { FieldVar, FieldConst } from '../provable/core/fieldvar.js';
 import { Cache, readCache, writeCache } from './cache.js';
@@ -54,6 +55,7 @@ export {
   Undefined,
   Void,
   VerificationKey,
+  FeatureFlags,
 };
 
 // internal API
@@ -83,6 +85,101 @@ const Empty = Undefined;
 type Void = undefined;
 const Void: ProvablePureExtended<void, void, null> = EmptyVoid<Field>();
 
+type MaybeFeatureFlag = boolean | undefined;
+type FeatureFlags = {
+  rangeCheck0: MaybeFeatureFlag;
+  rangeCheck1: MaybeFeatureFlag;
+  foreignFieldAdd: MaybeFeatureFlag;
+  foreignFieldMul: MaybeFeatureFlag;
+  xor: MaybeFeatureFlag;
+  rot: MaybeFeatureFlag;
+  lookup: MaybeFeatureFlag;
+  runtimeTables: MaybeFeatureFlag;
+};
+const FeatureFlags = {
+  /**
+   * Returns a feature flag configuration where all flags are set to false.
+   */
+  allNone: {
+    rangeCheck0: false,
+    rangeCheck1: false,
+    foreignFieldAdd: false,
+    foreignFieldMul: false,
+    xor: false,
+    rot: false,
+    lookup: false,
+    runtimeTables: false,
+  },
+  /**
+   * Returns a feature flag configuration where all flags are optional.
+   */
+  allMaybe: {
+    rangeCheck0: undefined,
+    rangeCheck1: undefined,
+    foreignFieldAdd: undefined,
+    foreignFieldMul: undefined,
+    xor: undefined,
+    rot: undefined,
+    lookup: undefined,
+    runtimeTables: undefined,
+  },
+
+  /**
+   * Given a list of gates, returns the feature flag configuration that the gates use.
+   */
+  fromGates: (gates: Gate[]) => featureFlagsFromGates(gates),
+
+  fromZkProgram: fromZkProgram,
+};
+
+async function fromZkProgram<
+  P extends {
+    analyzeMethods: () => Promise<{
+      [I in keyof any]: UnwrapPromise<ReturnType<typeof analyzeMethod>>;
+    }>;
+  }
+>(program: P): Promise<FeatureFlags> {
+  let methodIntfs = await program.analyzeMethods();
+
+  let flags = Object.entries(methodIntfs).map(([_, { gates }]) => {
+    return featureFlagsFromGates(gates);
+  });
+
+  if (flags.length === 0)
+    throw Error(
+      'The ZkProgram has no methods, in order to calculate feature flags, please attach a method to your ZkProgram.'
+    );
+
+  // initialize feature flags to all false
+  let globalFlags: Record<string, boolean | undefined> = {
+    rangeCheck0: false,
+    rangeCheck1: false,
+    foreignFieldAdd: false,
+    foreignFieldMul: false,
+    xor: false,
+    rot: false,
+    lookup: false,
+    runtimeTables: false,
+  };
+
+  // only one method defines the feature flags
+  if (flags.length === 1) return flags[0];
+
+  // calculating the crossover between all methods, compute the shared feature flag set
+  flags.forEach((featureFlags, i) => {
+    for (const [flagType, currentFlag] of Object.entries(featureFlags)) {
+      if (i === 0) {
+        // initialize first iteration of flags freely
+        globalFlags[flagType] = currentFlag;
+      } else if (globalFlags[flagType] != currentFlag) {
+        // if flags conflict, set them to undefined to account for both cases (true and false) ^= maybe
+        globalFlags[flagType] = undefined;
+      }
+    }
+  });
+  return globalFlags as FeatureFlags;
+}
+
 class ProofBase<Input, Output> {
   static publicInputType: FlexibleProvablePure<any> = undefined as any;
   static publicOutputType: FlexibleProvablePure<any> = undefined as any;
@@ -227,7 +324,7 @@ var sideloadedKeysCounter = 0;
  * This pattern differs a lot from the usage of normal `Proof`, where the verification key is baked into the compiled circuit.
  * @see {@link src/examples/zkprogram/dynamic-keys-merkletree.ts} for an example of how this can be done using merkle trees
  *
- * Assertions generally only happen using the vk hash that is part of the `VerificationKey` struct along with the raw vk data as auxilary data.
+ * Assertions generally only happen using the vk hash that is part of the `VerificationKey` struct along with the raw vk data as auxiliary data.
  * When using verify() on a `DynamicProof`, Pickles makes sure that the verification key data matches the hash.
  * Therefore all manual assertions have to be made on the vk's hash and it can be assumed that the vk's data is checked to match the hash if it is used with verify().
  */
@@ -236,6 +333,21 @@ class DynamicProof<Input, Output> extends ProofBase<Input, Output> {
 
   private static memoizedCounter: number | undefined;
 
+  /**
+   * As the name indicates, feature flags are features of the proof system.
+   *
+   * If we want to side load proofs and verification keys, we first have to tell Pickles what _shape_ of proofs it should expect.
+   *
+   * For example, if we want to side load proofs that use foreign field arithmetic custom gates, we have to make Pickles aware of that by defining
+   * these custom gates.
+   *
+   * _Note:_ Only proofs use exactly the same composition of custom gates as was expected by Pickles can be verified.
+   * If you want to verify _any_ proofs, no matter what custom gates it uses, you can use {@link FeatureFlags.allMaybe}. Please note that this might incur a small overhead.
+   *
+   * You can also toggle specific feature flags manually by specifying them here.
+   */
+  static featureFlags: FeatureFlags = FeatureFlags.allMaybe;
+
   static tag() {
     let counter: number;
     if (this.memoizedCounter !== undefined) {
@@ -1014,11 +1126,19 @@ function picklesRuleFromFunction(
       let computedTag: unknown;
       // Only create the tag if it hasn't already been created for this specific Proof class
       if (SideloadedTag.get(tag.name) === undefined) {
+        let featureFlags = [
+          0,
+          ...Object.entries(Proof.featureFlags).map(([_, flag]) =>
+            MlOption.mapTo(flag, MlBool)
+          ),
+        ] as MlArrayOptionalElements<MlFeatureFlags>;
+
         computedTag = Pickles.sideLoaded.create(
           tag.name,
           Proof.maxProofsVerified,
           Proof.publicInputType?.sizeInFields() ?? 0,
-          Proof.publicOutputType?.sizeInFields() ?? 0
+          Proof.publicOutputType?.sizeInFields() ?? 0,
+          featureFlags
         );
         SideloadedTag.store(tag.name, computedTag);
       } else {
@@ -1037,7 +1157,7 @@ function picklesRuleFromFunction(
     }
   });
 
-  let featureFlags = computeFeatureFlags(gates);
+  let featureFlags = featureFlagsToMl(featureFlagsFromGates(gates));
 
   return {
     identifier: methodName,
@@ -1215,7 +1335,7 @@ const gateToFlag: Partial<Record<GateType, keyof FeatureFlags>> = {
   Lookup: 'lookup',
 };
 
-function computeFeatureFlags(gates: Gate[]): MlFeatureFlags {
+function featureFlagsFromGates(gates: Gate[]): FeatureFlags {
   let flags: FeatureFlags = {
     rangeCheck0: false,
     rangeCheck1: false,
@@ -1230,17 +1350,16 @@ function computeFeatureFlags(gates: Gate[]): MlFeatureFlags {
     let flag = gateToFlag[gate.type];
     if (flag !== undefined) flags[flag] = true;
   }
+  return flags;
+}
+
+function featureFlagsToMl(flags: FeatureFlags): MlFeatureFlags {
   return [
     0,
-    MlBool(flags.rangeCheck0),
-    MlBool(flags.rangeCheck1),
-    MlBool(flags.foreignFieldAdd),
-    MlBool(flags.foreignFieldMul),
-    MlBool(flags.xor),
-    MlBool(flags.rot),
-    MlBool(flags.lookup),
-    MlBool(flags.runtimeTables),
-  ];
+    ...Object.entries(flags).map(([_, flag]) =>
+      MlBool(typeof flag === 'boolean' ? flag : false)
+    ),
+  ] as MlFeatureFlags;
 }
 
 // helpers for circuit context

src/snarky.d.ts

@@ -18,6 +18,7 @@ import type {
   MlUnit,
   MlString,
   MlTuple,
+  MlArrayOptionalElements,
 } from './lib/ml/base.js';
 import type { MlHashInput } from './lib/ml/conversion.js';
 import type {
@@ -45,7 +46,6 @@ export {
   JsonGate,
   MlPublicKey,
   MlPublicKeyVar,
-  FeatureFlags,
   MlFeatureFlags,
 };
 
@@ -598,17 +598,6 @@ type Test = {
   };
 };
 
-type FeatureFlags = {
-  rangeCheck0: boolean;
-  rangeCheck1: boolean;
-  foreignFieldAdd: boolean;
-  foreignFieldMul: boolean;
-  xor: boolean;
-  rot: boolean;
-  lookup: boolean;
-  runtimeTables: boolean;
-};
-
 type MlFeatureFlags = [
   _: 0,
   rangeCheck0: MlBool,
@@ -747,7 +736,8 @@ declare const Pickles: {
       name: string,
       numProofsVerified: 0 | 1 | 2,
       publicInputLength: number,
-      publicOutputLength: number
+      publicOutputLength: number,
+      featureFlags: MlArrayOptionalElements<MlFeatureFlags>
     ) => unknown /* tag */;
     // Instantiate the verification key inside the circuit (required).
     inCircuit: (tag: unknown, verificationKey: unknown) => undefined;