Activate custom gates in JS and add rangeCheck64 gadget

CHANGELOG.md

@@ -27,6 +27,8 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 - `Lightnet` namespace to interact with the account manager provided by the [lightnet Mina network](https://hub.docker.com/r/o1labs/mina-local-network). https://github.com/o1-labs/o1js/pull/1167
 
+- Internal support for several custom gates (range check, bitwise operations, foreign field operations) and lookup tables https://github.com/o1-labs/o1js/pull/1176
+
 ## [0.13.1](https://github.com/o1-labs/o1js/compare/c2f392fe5...045faa7)
 
 ### Breaking changes

src/examples/ex02_root.ts

@@ -1,6 +1,4 @@
-import { Field, Circuit, circuitMain, public_, isReady } from 'o1js';
-
-await isReady;
+import { Field, Circuit, circuitMain, public_, UInt64 } from 'o1js';
 
 /* Exercise 2:
 
@@ -11,9 +9,9 @@ Prove:
 
 class Main extends Circuit {
   @circuitMain
-  static main(y: Field, @public_ x: Field) {
+  static main(@public_ y: Field, x: UInt64) {
     let y3 = y.square().mul(y);
-    y3.assertEquals(x);
+    y3.assertEquals(x.value);
   }
 }
 
@@ -24,15 +22,15 @@ console.timeEnd('generating keypair...');
 
 console.log('prove...');
 console.time('prove...');
-const x = new Field(8);
+const x = UInt64.from(8);
 const y = new Field(2);
-const proof = await Main.prove([y], [x], kp);
+const proof = await Main.prove([x], [y], kp);
 console.timeEnd('prove...');
 
 console.log('verify...');
 console.time('verify...');
 let vk = kp.verificationKey();
-let ok = await Main.verify([x], vk, proof);
+let ok = await Main.verify([y], vk, proof);
 console.timeEnd('verify...');
 
 console.log('ok?', ok);

src/examples/ex02_root_program.ts

@@ -0,0 +1,42 @@
+import {
+  Field,
+  Circuit,
+  circuitMain,
+  public_,
+  UInt64,
+  Experimental,
+} from 'o1js';
+
+let { ZkProgram } = Experimental;
+
+const Main = ZkProgram({
+  publicInput: Field,
+  methods: {
+    main: {
+      privateInputs: [UInt64],
+      method(y: Field, x: UInt64) {
+        let y3 = y.square().mul(y);
+        y3.assertEquals(x.value);
+      },
+    },
+  },
+});
+
+console.log('generating keypair...');
+console.time('generating keypair...');
+const kp = await Main.compile();
+console.timeEnd('generating keypair...');
+
+console.log('prove...');
+console.time('prove...');
+const x = UInt64.from(8);
+const y = new Field(2);
+const proof = await Main.main(y, x);
+console.timeEnd('prove...');
+
+console.log('verify...');
+console.time('verify...');
+let ok = await Main.verify(proof);
+console.timeEnd('verify...');
+
+console.log('ok?', ok);

src/lib/gadgets/range-check.ts

@@ -0,0 +1,14 @@
+import { type Field } from '../field.js';
+import * as Gates from '../gates.js';
+
+export { rangeCheck64 };
+
+function rangeCheck64(x: Field) {
+  if (x.isConstant()) {
+    if (x.toBigInt() >= 1n << 64n) {
+      throw Error(`rangeCheck64: expected field to fit in 64 bits, got ${x}`);
+    }
+  } else {
+    Gates.rangeCheck64(x);
+  }
+}

src/lib/gates.ts

@@ -0,0 +1,49 @@
+import { Snarky } from '../snarky.js';
+import { FieldVar, FieldConst, type Field } from './field.js';
+
+export { rangeCheck64 };
+
+/**
+ * Asserts that x is at most 64 bits
+ */
+function rangeCheck64(x: Field) {
+  let [, x0, x2, x4, x6, x8, x10, x12, x14] = Snarky.exists(8, () => {
+    let xx = x.toBigInt();
+    // crumbs (2-bit limbs)
+    return [
+      0,
+      getBits(xx, 0, 2),
+      getBits(xx, 2, 2),
+      getBits(xx, 4, 2),
+      getBits(xx, 6, 2),
+      getBits(xx, 8, 2),
+      getBits(xx, 10, 2),
+      getBits(xx, 12, 2),
+      getBits(xx, 14, 2),
+    ];
+  });
+  // 12-bit limbs
+  let [, x16, x28, x40, x52] = Snarky.exists(4, () => {
+    let xx = x.toBigInt();
+    return [
+      0,
+      getBits(xx, 16, 12),
+      getBits(xx, 28, 12),
+      getBits(xx, 40, 12),
+      getBits(xx, 52, 12),
+    ];
+  });
+  Snarky.gates.rangeCheck0(
+    x.value,
+    [0, FieldVar[0], FieldVar[0], x52, x40, x28, x16],
+    [0, x14, x12, x10, x8, x6, x4, x2, x0],
+    // not using compact mode
+    FieldConst[0]
+  );
+}
+
+function getBits(x: bigint, start: number, length: number) {
+  return FieldConst.fromBigint(
+    (x >> BigInt(start)) & ((1n << BigInt(length)) - 1n)
+  );
+}

src/lib/int.ts

@@ -69,9 +69,11 @@ class UInt64 extends CircuitValue {
     let actual = x.value.rangeCheckHelper(64);
     actual.assertEquals(x.value);
   }
+
   static toInput(x: UInt64): HashInput {
     return { packed: [[x.value, 64]] };
   }
+
   /**
    * Encodes this structure into a JSON-like object.
    */

src/lib/ml/base.ts

@@ -77,5 +77,11 @@ const MlOption = Object.assign(
       if (option === undefined) return 0;
       return [0, map(option)];
     },
+    isNone(option: MlOption<unknown>): option is 0 {
+      return option === 0;
+    },
+    isSome<T>(option: MlOption<T>): option is [0, T] {
+      return option !== 0;
+    },
   }
 );

src/lib/proof_system.ts

@@ -4,7 +4,14 @@ import {
   EmptyVoid,
 } from '../bindings/lib/generic.js';
 import { withThreadPool } from '../bindings/js/wrapper.js';
-import { ProvablePure, Pickles } from '../snarky.js';
+import {
+  ProvablePure,
+  Pickles,
+  FeatureFlags,
+  MlFeatureFlags,
+  Gate,
+  GateType,
+} from '../snarky.js';
 import { Field, Bool } from './core.js';
 import {
   FlexibleProvable,
@@ -18,7 +25,7 @@ import { Provable } from './provable.js';
 import { assert, prettifyStacktracePromise } from './errors.js';
 import { snarkContext } from './provable-context.js';
 import { hashConstant } from './hash.js';
-import { MlArray, MlTuple } from './ml/base.js';
+import { MlArray, MlBool, MlTuple } from './ml/base.js';
 import { MlFieldArray, MlFieldConstArray } from './ml/fields.js';
 import { FieldConst, FieldVar } from './field.js';
 
@@ -249,6 +256,12 @@ function ZkProgram<
   let methodFunctions = keys.map((key) => methods[key].method);
   let maxProofsVerified = getMaxProofsVerified(methodIntfs);
 
+  function analyzeMethods() {
+    return methodIntfs.map((methodEntry, i) =>
+      analyzeMethod(publicInputType, methodEntry, methodFunctions[i])
+    );
+  }
+
   let compileOutput:
     | {
         provers: Pickles.Prover[];
@@ -260,14 +273,17 @@ function ZkProgram<
     | undefined;
 
   async function compile() {
-    let { provers, verify, verificationKey } = await compileProgram(
+    let methodsMeta = analyzeMethods();
+    let gates = methodsMeta.map((m) => m.gates);
+    let { provers, verify, verificationKey } = await compileProgram({
       publicInputType,
       publicOutputType,
       methodIntfs,
-      methodFunctions,
-      selfTag,
-      config.overrideWrapDomain
-    );
+      methods: methodFunctions,
+      gates,
+      proofSystemTag: selfTag,
+      overrideWrapDomain: config.overrideWrapDomain,
+    });
     compileOutput = { provers, verify };
     return { verificationKey: verificationKey.data };
   }
@@ -352,12 +368,6 @@ function ZkProgram<
     return hash.toBigInt().toString(16);
   }
 
-  function analyzeMethods() {
-    return methodIntfs.map((methodEntry, i) =>
-      analyzeMethod(publicInputType, methodEntry, methodFunctions[i])
-    );
-  }
-
   return Object.assign(
     selfTag,
     {
@@ -500,21 +510,31 @@ type MethodInterface = {
 // reasonable default choice for `overrideWrapDomain`
 const maxProofsToWrapDomain = { 0: 0, 1: 1, 2: 1 } as const;
 
-async function compileProgram(
-  publicInputType: ProvablePure<any>,
-  publicOutputType: ProvablePure<any>,
-  methodIntfs: MethodInterface[],
-  methods: ((...args: any) => void)[],
-  proofSystemTag: { name: string },
-  overrideWrapDomain?: 0 | 1 | 2
-) {
+async function compileProgram({
+  publicInputType,
+  publicOutputType,
+  methodIntfs,
+  methods,
+  gates,
+  proofSystemTag,
+  overrideWrapDomain,
+}: {
+  publicInputType: ProvablePure<any>;
+  publicOutputType: ProvablePure<any>;
+  methodIntfs: MethodInterface[];
+  methods: ((...args: any) => void)[];
+  gates: Gate[][];
+  proofSystemTag: { name: string };
+  overrideWrapDomain?: 0 | 1 | 2;
+}) {
   let rules = methodIntfs.map((methodEntry, i) =>
     picklesRuleFromFunction(
       publicInputType,
       publicOutputType,
       methods[i],
       proofSystemTag,
-      methodEntry
+      methodEntry,
+      gates[i]
     )
   );
   let maxProofs = getMaxProofsVerified(methodIntfs);
@@ -589,7 +609,8 @@ function picklesRuleFromFunction(
   publicOutputType: ProvablePure<unknown>,
   func: (...args: unknown[]) => any,
   proofSystemTag: { name: string },
-  { methodName, witnessArgs, proofArgs, allArgs }: MethodInterface
+  { methodName, witnessArgs, proofArgs, allArgs }: MethodInterface,
+  gates: Gate[]
 ): Pickles.Rule {
   function main(publicInput: MlFieldArray): ReturnType<Pickles.Rule['main']> {
     let { witnesses: argsWithoutPublicInput, inProver } = snarkContext.get();
@@ -664,9 +685,13 @@ function picklesRuleFromFunction(
       return { isSelf: false, tag: compiledTag };
     }
   });
+
+  let featureFlags = computeFeatureFlags(gates);
+
   return {
     identifier: methodName,
     main,
+    featureFlags,
     proofsToVerify: MlArray.to(proofsToVerify),
   };
 }
@@ -822,6 +847,46 @@ function dummyBase64Proof() {
   return withThreadPool(async () => Pickles.dummyBase64Proof());
 }
 
+// what feature flags to set to enable certain gate types
+
+const gateToFlag: Partial<Record<GateType, keyof FeatureFlags>> = {
+  RangeCheck0: 'rangeCheck0',
+  RangeCheck1: 'rangeCheck1',
+  ForeignFieldAdd: 'foreignFieldAdd',
+  ForeignFieldMul: 'foreignFieldMul',
+  Xor16: 'xor',
+  Rot64: 'rot',
+  Lookup: 'lookup',
+};
+
+function computeFeatureFlags(gates: Gate[]): MlFeatureFlags {
+  let flags: FeatureFlags = {
+    rangeCheck0: false,
+    rangeCheck1: false,
+    foreignFieldAdd: false,
+    foreignFieldMul: false,
+    xor: false,
+    rot: false,
+    lookup: false,
+    runtimeTables: false,
+  };
+  for (let gate of gates) {
+    let flag = gateToFlag[gate.type];
+    if (flag !== undefined) flags[flag] = true;
+  }
+  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),
+  ];
+}
+
 // helpers for circuit context
 
 function Prover<ProverData>() {

src/lib/zkapp.ts

@@ -675,18 +675,20 @@ class SmartContract {
       };
     });
     // run methods once to get information that we need already at compile time
-    this.analyzeMethods();
+    let methodsMeta = this.analyzeMethods();
+    let gates = methodIntfs.map((intf) => methodsMeta[intf.methodName].gates);
     let {
       verificationKey: verificationKey_,
       provers,
       verify,
-    } = await compileProgram(
-      ZkappPublicInput,
-      Empty,
+    } = await compileProgram({
+      publicInputType: ZkappPublicInput,
+      publicOutputType: Empty,
       methodIntfs,
       methods,
-      this
-    );
+      gates,
+      proofSystemTag: this,
+    });
     let verificationKey = {
       data: verificationKey_.data,
       hash: Field(verificationKey_.hash),