Foreign fields 2: Test DSL to detect broken gate chains

CHANGELOG.md

@@ -19,6 +19,14 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased](https://github.com/o1-labs/o1js/compare/26363465d...HEAD)
 
+### Added
+
+- Comprehensive internal testing of constraint system layouts generated by new gadgets https://github.com/o1-labs/o1js/pull/1241
+
+### Changed
+
+- Expose raw provable methods of a `ZkProgram` on `zkProgram.rawMethods` https://github.com/o1-labs/o1js/pull/1241
+
 ### Fixed
 
 - Add a parameter to `checkZkappTransaction` for block length to check for transaction inclusion. This fixes a case where `Transaction.wait()` only checked the latest block, which led to an error once the transaction was included in a block that was not the latest. https://github.com/o1-labs/o1js/pull/1239
@@ -27,9 +35,9 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Added
 
-- `Gadgets.not()`, new provable method to support bitwise shifting for native field elements. https://github.com/o1-labs/o1js/pull/1198
-- `Gadgets.leftShift() / Gadgets.rightShift()`, new provable method to support bitwise shifting for native field elements. https://github.com/o1-labs/o1js/pull/1194
-- `Gadgets.and()`, new provable method to support bitwise and for native field elements. https://github.com/o1-labs/o1js/pull/1193
+- `Gadgets.not()`, new provable method to support bitwise not. https://github.com/o1-labs/o1js/pull/1198
+- `Gadgets.leftShift() / Gadgets.rightShift()`, new provable methods to support bitwise shifting. https://github.com/o1-labs/o1js/pull/1194
+- `Gadgets.and()`, new provable method to support bitwise and. https://github.com/o1-labs/o1js/pull/1193
 - `Gadgets.multiRangeCheck()` and `Gadgets.compactMultiRangeCheck()`, two building blocks for non-native arithmetic with bigints of size up to 264 bits. https://github.com/o1-labs/o1js/pull/1216
 
 ## [0.14.0](https://github.com/o1-labs/o1js/compare/045faa7...e8e7510e1)

src/lib/field.ts

@@ -78,13 +78,22 @@ const FieldVar = {
   isConstant(x: FieldVar): x is ConstantFieldVar {
     return x[0] === FieldType.Constant;
   },
-  // TODO: handle (special) constants
   add(x: FieldVar, y: FieldVar): FieldVar {
+    if (FieldVar.isConstant(x) && x[1][1] === 0n) return y;
+    if (FieldVar.isConstant(y) && y[1][1] === 0n) return x;
+    if (FieldVar.isConstant(x) && FieldVar.isConstant(y)) {
+      return FieldVar.constant(Fp.add(x[1][1], y[1][1]));
+    }
     return [FieldType.Add, x, y];
   },
-  // TODO: handle (special) constants
-  scale(c: FieldConst, x: FieldVar): FieldVar {
-    return [FieldType.Scale, c, x];
+  scale(c: bigint | FieldConst, x: FieldVar): FieldVar {
+    let c0 = typeof c === 'bigint' ? FieldConst.fromBigint(c) : c;
+    if (c0[1] === 0n) return FieldVar.constant(0n);
+    if (c0[1] === 1n) return x;
+    if (FieldVar.isConstant(x)) {
+      return FieldVar.constant(Fp.mul(c0[1], x[1][1]));
+    }
+    return [FieldType.Scale, c0, x];
   },
   [0]: [FieldType.Constant, FieldConst[0]] satisfies ConstantFieldVar,
   [1]: [FieldType.Constant, FieldConst[1]] satisfies ConstantFieldVar,

src/lib/gadgets/bitwise.ts

@@ -8,6 +8,7 @@ import {
   witnessSlice,
   witnessNextValue,
   divideWithRemainder,
+  toVar,
 } from './common.js';
 import { rangeCheck64 } from './range-check.js';
 
@@ -37,18 +38,12 @@ function not(a: Field, length: number, checked: boolean = false) {
   }
 
   // create a bitmask with all ones
-  let allOnesF = new Field(2n ** BigInt(length) - 1n);
-
-  let allOnes = Provable.witness(Field, () => {
-    return allOnesF;
-  });
-
-  allOnesF.assertEquals(allOnes);
+  let allOnes = new Field(2n ** BigInt(length) - 1n);
 
   if (checked) {
     return xor(a, allOnes, length);
   } else {
-    return allOnes.sub(a);
+    return allOnes.sub(a).seal();
   }
 }
 
@@ -252,6 +247,10 @@ function rot(
     }
   );
 
+  // flush zero var to prevent broken gate chain (zero is used in rangeCheck64)
+  // TODO this is an abstraction leak, but not clear to me how to improve
+  toVar(0n);
+
   // Compute current row
   Gates.rotate(
     field,

src/lib/gadgets/bitwise.unit-test.ts

@@ -9,6 +9,16 @@ import { Fp, mod } from '../../bindings/crypto/finite_field.js';
 import { Field } from '../core.js';
 import { Gadgets } from './gadgets.js';
 import { Random } from '../testing/property.js';
+import {
+  constraintSystem,
+  contains,
+  equals,
+  ifNotAllConstant,
+  repeat,
+  and,
+  withoutGenerics,
+} from '../testing/constraint-system.js';
+import { GateType } from '../../snarky.js';
 
 const maybeField = {
   ...field,
@@ -183,3 +193,63 @@ await equivalentAsync({ from: [field], to: field }, { runs: 3 })(
     return proof.publicOutput;
   }
 );
+
+// check that gate chains stay intact
+
+function xorChain(bits: number) {
+  return repeat(Math.ceil(bits / 16), 'Xor16').concat('Zero');
+}
+
+constraintSystem(
+  'xor',
+  { from: [Field, Field] },
+  Bitwise.rawMethods.xor,
+  ifNotAllConstant(contains(xorChain(254)))
+);
+
+constraintSystem(
+  'not checked',
+  { from: [Field] },
+  Bitwise.rawMethods.notChecked,
+  ifNotAllConstant(contains(xorChain(254)))
+);
+
+constraintSystem(
+  'not unchecked',
+  { from: [Field] },
+  Bitwise.rawMethods.notUnchecked,
+  ifNotAllConstant(contains('Generic'))
+);
+
+constraintSystem(
+  'and',
+  { from: [Field, Field] },
+  Bitwise.rawMethods.and,
+  ifNotAllConstant(contains(xorChain(64)))
+);
+
+let rotChain: GateType[] = ['Rot64', 'RangeCheck0', 'RangeCheck0'];
+let isJustRotate = ifNotAllConstant(
+  and(contains(rotChain), withoutGenerics(equals(rotChain)))
+);
+
+constraintSystem(
+  'rotate',
+  { from: [Field] },
+  Bitwise.rawMethods.rot,
+  isJustRotate
+);
+
+constraintSystem(
+  'left shift',
+  { from: [Field] },
+  Bitwise.rawMethods.leftShift,
+  isJustRotate
+);
+
+constraintSystem(
+  'right shift',
+  { from: [Field] },
+  Bitwise.rawMethods.rightShift,
+  isJustRotate
+);

src/lib/gadgets/common.ts

@@ -1,6 +1,6 @@
 import { Provable } from '../provable.js';
-import { Field, FieldConst } from '../field.js';
-import { TupleN } from '../util/types.js';
+import { Field, FieldConst, FieldVar, FieldType } from '../field.js';
+import { Tuple, TupleN } from '../util/types.js';
 import { Snarky } from '../../snarky.js';
 import { MlArray } from '../ml/base.js';
 
@@ -9,13 +9,21 @@ const MAX_BITS = 64 as const;
 export {
   MAX_BITS,
   exists,
+  existsOne,
+  toVars,
+  toVar,
   assert,
   bitSlice,
   witnessSlice,
   witnessNextValue,
   divideWithRemainder,
 };
 
+function existsOne(compute: () => bigint) {
+  let varMl = Snarky.existsVar(() => FieldConst.fromBigint(compute()));
+  return new Field(varMl);
+}
+
 function exists<N extends number, C extends () => TupleN<bigint, N>>(
   n: N,
   compute: C
@@ -27,6 +35,31 @@ function exists<N extends number, C extends () => TupleN<bigint, N>>(
   return TupleN.fromArray(n, vars);
 }
 
+/**
+ * Given a Field, collapse its AST to a pure Var. See {@link FieldVar}.
+ *
+ * This is useful to prevent rogue Generic gates added in the middle of gate chains,
+ * which are caused by snarky auto-resolving constants, adds and scales to vars.
+ *
+ * Same as `Field.seal()` with the difference that `seal()` leaves constants as is.
+ */
+function toVar(x: Field | bigint) {
+  // don't change existing vars
+  if (x instanceof Field && x.value[1] === FieldType.Var) return x;
+  let xVar = existsOne(() => Field.from(x).toBigInt());
+  xVar.assertEquals(x);
+  return xVar;
+}
+
+/**
+ * Apply {@link toVar} to each element of a tuple.
+ */
+function toVars<T extends Tuple<Field | bigint>>(
+  fields: T
+): { [k in keyof T]: Field } {
+  return Tuple.map(fields, toVar);
+}
+
 function assert(stmt: boolean, message?: string) {
   if (!stmt) {
     throw Error(message ?? 'Assertion failed');

src/lib/gadgets/range-check.ts

@@ -1,6 +1,6 @@
 import { Field } from '../field.js';
 import * as Gates from '../gates.js';
-import { bitSlice, exists } from './common.js';
+import { bitSlice, exists, toVar, toVars } from './common.js';
 
 export { rangeCheck64, multiRangeCheck, compactMultiRangeCheck, L };
 
@@ -64,10 +64,13 @@ function multiRangeCheck([x, y, z]: [Field, Field, Field]) {
     }
     return;
   }
+  // ensure we are using pure variables
+  [x, y, z] = toVars([x, y, z]);
+  let zero = toVar(0n);
 
   let [x64, x76] = rangeCheck0Helper(x);
   let [y64, y76] = rangeCheck0Helper(y);
-  rangeCheck1Helper({ x64, x76, y64, y76, z, yz: new Field(0) });
+  rangeCheck1Helper({ x64, x76, y64, y76, z, yz: zero });
 }
 
 /**
@@ -88,6 +91,8 @@ function compactMultiRangeCheck(xy: Field, z: Field): [Field, Field, Field] {
     let [x, y] = splitCompactLimb(xy.toBigInt());
     return [new Field(x), new Field(y), z];
   }
+  // ensure we are using pure variables
+  [xy, z] = toVars([xy, z]);
 
   let [x, y] = exists(2, () => splitCompactLimb(xy.toBigInt()));
 

src/lib/gadgets/range-check.unit-test.ts

@@ -1,19 +1,23 @@
-import type { Gate } from '../../snarky.js';
 import { mod } from '../../bindings/crypto/finite_field.js';
 import { Field } from '../../lib/core.js';
 import { ZkProgram } from '../proof_system.js';
-import { Provable } from '../provable.js';
 import {
   Spec,
   boolean,
   equivalentAsync,
   fieldWithRng,
 } from '../testing/equivalent.js';
 import { Random } from '../testing/property.js';
-import { assert, exists } from './common.js';
+import { assert } from './common.js';
 import { Gadgets } from './gadgets.js';
 import { L } from './range-check.js';
-import { expect } from 'expect';
+import {
+  constraintSystem,
+  contains,
+  equals,
+  ifNotAllConstant,
+  withoutGenerics,
+} from '../testing/constraint-system.js';
 
 let uint = (n: number | bigint): Spec<bigint, Field> => {
   let uint = Random.bignat((1n << BigInt(n)) - 1n);
@@ -29,29 +33,30 @@ let maybeUint = (n: number | bigint): Spec<bigint, Field> => {
 
 // constraint system sanity check
 
-function csWithoutGenerics(gates: Gate[]) {
-  return gates.map((g) => g.type).filter((type) => type !== 'Generic');
-}
-
-let check64 = Provable.constraintSystem(() => {
-  let [x] = exists(1, () => [0n]);
-  Gadgets.rangeCheck64(x);
-});
-let multi = Provable.constraintSystem(() => {
-  let x = exists(3, () => [0n, 0n, 0n]);
-  Gadgets.multiRangeCheck(x);
-});
-let compact = Provable.constraintSystem(() => {
-  let [xy, z] = exists(2, () => [0n, 0n]);
-  Gadgets.compactMultiRangeCheck(xy, z);
-});
+constraintSystem(
+  'range check 64',
+  { from: [Field] },
+  Gadgets.rangeCheck64,
+  ifNotAllConstant(withoutGenerics(equals(['RangeCheck0'])))
+);
 
-let expectedLayout64 = ['RangeCheck0'];
-let expectedLayoutMulti = ['RangeCheck0', 'RangeCheck0', 'RangeCheck1', 'Zero'];
+constraintSystem(
+  'multi-range check',
+  { from: [Field, Field, Field] },
+  (x, y, z) => Gadgets.multiRangeCheck([x, y, z]),
+  ifNotAllConstant(
+    contains(['RangeCheck0', 'RangeCheck0', 'RangeCheck1', 'Zero'])
+  )
+);
 
-expect(csWithoutGenerics(check64.gates)).toEqual(expectedLayout64);
-expect(csWithoutGenerics(multi.gates)).toEqual(expectedLayoutMulti);
-expect(csWithoutGenerics(compact.gates)).toEqual(expectedLayoutMulti);
+constraintSystem(
+  'compact multi-range check',
+  { from: [Field, Field] },
+  Gadgets.compactMultiRangeCheck,
+  ifNotAllConstant(
+    contains(['RangeCheck0', 'RangeCheck0', 'RangeCheck1', 'Zero'])
+  )
+);
 
 // TODO: make a ZkFunction or something that doesn't go through Pickles
 // --------------------------

src/lib/proof_system.ts

@@ -271,6 +271,13 @@ function ZkProgram<
   analyzeMethods: () => ReturnType<typeof analyzeMethod>[];
   publicInputType: ProvableOrUndefined<Get<StatementType, 'publicInput'>>;
   publicOutputType: ProvableOrVoid<Get<StatementType, 'publicOutput'>>;
+  rawMethods: {
+    [I in keyof Types]: Method<
+      InferProvableOrUndefined<Get<StatementType, 'publicInput'>>,
+      InferProvableOrVoid<Get<StatementType, 'publicOutput'>>,
+      Types[I]
+    >['method'];
+  };
 } & {
   [I in keyof Types]: Prover<
     InferProvableOrUndefined<Get<StatementType, 'publicInput'>>,
@@ -427,6 +434,9 @@ function ZkProgram<
         Get<StatementType, 'publicOutput'>
       >,
       analyzeMethods,
+      rawMethods: Object.fromEntries(
+        Object.entries(methods).map(([k, v]) => [k, v.method])
+      ) as any,
     },
     provers
   );