Merge pull request #1182 from o1-labs/feat/ROT-gadget

Add rot gadget to o1js

CHANGELOG.md

@@ -31,9 +31,15 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 ### Added
 
 - `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
+
 - `Gadgets.rangeCheck64()`, new provable method to do efficient 64-bit range checks using lookup tables https://github.com/o1-labs/o1js/pull/1181
-- Added bitwise `XOR` operation support for native field elements. https://github.com/o1-labs/o1js/pull/1177
+
+- `Gadgets.rotate()`, new provable method to support bitwise rotation for native field elements. https://github.com/o1-labs/o1js/pull/1182
+
+- `Gadgets.xor()`, new provable method to support bitwise xor for native field elements. https://github.com/o1-labs/o1js/pull/1177
+
 - `Proof.dummy()` to create dummy proofs https://github.com/o1-labs/o1js/pull/1188
   - You can use this to write ZkPrograms that handle the base case and the inductive case in the same method.
 

src/examples/gadgets.ts

@@ -1,5 +1,39 @@
 import { Field, Provable, Gadgets, ZkProgram } from 'o1js';
 
+let cs = Provable.constraintSystem(() => {
+  let f = Provable.witness(Field, () => Field(12));
+
+  let res1 = Gadgets.rotate(f, 2, 'left');
+  let res2 = Gadgets.rotate(f, 2, 'right');
+
+  res1.assertEquals(Field(48));
+  res2.assertEquals(Field(3));
+
+  Provable.log(res1);
+  Provable.log(res2);
+});
+console.log('constraint system: ', cs);
+
+const ROT = ZkProgram({
+  name: 'rot-example',
+  methods: {
+    baseCase: {
+      privateInputs: [],
+      method: () => {
+        let a = Provable.witness(Field, () => Field(48));
+        let actualLeft = Gadgets.rotate(a, 2, 'left');
+        let actualRight = Gadgets.rotate(a, 2, 'right');
+
+        let expectedLeft = Field(192);
+        actualLeft.assertEquals(expectedLeft);
+
+        let expectedRight = Field(12);
+        actualRight.assertEquals(expectedRight);
+      },
+    },
+  },
+});
+
 const XOR = ZkProgram({
   name: 'xor-example',
   methods: {
@@ -19,14 +53,18 @@ const XOR = ZkProgram({
 console.log('compiling..');
 
 console.time('compile');
+await ROT.compile();
 await XOR.compile();
 console.timeEnd('compile');
 
 console.log('proving..');
 
-console.time('prove');
-let proof = await XOR.baseCase();
-console.timeEnd('prove');
+console.time('rotation prove');
+let rotProof = await ROT.baseCase();
+console.timeEnd('rotation prove');
+if (!(await ROT.verify(rotProof))) throw Error('rotate: Invalid proof');
 
+console.time('xor prove');
+let proof = await XOR.baseCase();
+console.timeEnd('xor prove');
 if (!(await XOR.verify(proof))) throw Error('Invalid proof');
-else console.log('proof valid');

src/examples/primitive_constraint_system.ts

@@ -64,6 +64,13 @@ const GroupMock = {
 };
 
 const BitwiseMock = {
+  rot() {
+    let a = Provable.witness(Field, () => new Field(12));
+    Gadgets.rotate(a, 2, 'left');
+    Gadgets.rotate(a, 2, 'right');
+    Gadgets.rotate(a, 4, 'left');
+    Gadgets.rotate(a, 4, 'right');
+  },
   xor() {
     let a = Provable.witness(Field, () => new Field(5n));
     let b = Provable.witness(Field, () => new Field(5n));

src/examples/regression_test.json

@@ -168,6 +168,10 @@
   "Bitwise Primitive": {
     "digest": "Bitwise Primitive",
     "methods": {
+      "rot": {
+        "rows": 13,
+        "digest": "2c0dadbba96fd7ddb9adb7d643425ce3"
+      },
       "xor": {
         "rows": 15,
         "digest": "b3595a9cc9562d4f4a3a397b6de44971"

src/lib/gadgets/bitwise.ts

@@ -2,8 +2,15 @@ import { Provable } from '../provable.js';
 import { Field as Fp } from '../../provable/field-bigint.js';
 import { Field } from '../field.js';
 import * as Gates from '../gates.js';
+import {
+  MAX_BITS,
+  assert,
+  witnessSlices,
+  witnessNextValue,
+  divideWithRemainder,
+} from './common.js';
 
-export { xor };
+export { xor, rotate };
 
 function xor(a: Field, b: Field, length: number) {
   // check that both input lengths are positive
@@ -111,21 +118,83 @@ function buildXor(
   zero.assertEquals(expectedOutput);
 }
 
-function assert(stmt: boolean, message?: string) {
-  if (!stmt) {
-    throw Error(message ?? 'Assertion failed');
+function rotate(
+  field: Field,
+  bits: number,
+  direction: 'left' | 'right' = 'left'
+) {
+  // Check that the rotation bits are in range
+  assert(
+    bits >= 0 && bits <= MAX_BITS,
+    `rotation: expected bits to be between 0 and 64, got ${bits}`
+  );
+
+  if (field.isConstant()) {
+    assert(
+      field.toBigInt() < 2n ** BigInt(MAX_BITS),
+      `rotation: expected field to be at most 64 bits, got ${field.toBigInt()}`
+    );
+    return new Field(Fp.rot(field.toBigInt(), bits, direction));
   }
+  const [rotated] = rot(field, bits, direction);
+  return rotated;
 }
 
-function witnessSlices(f: Field, start: number, length: number) {
-  if (length <= 0) throw Error('Length must be a positive number');
-
-  return Provable.witness(Field, () => {
-    let n = f.toBigInt();
-    return new Field((n >> BigInt(start)) & ((1n << BigInt(length)) - 1n));
-  });
-}
+function rot(
+  field: Field,
+  bits: number,
+  direction: 'left' | 'right' = 'left'
+): [Field, Field, Field] {
+  const rotationBits = direction === 'right' ? MAX_BITS - bits : bits;
+  const big2Power64 = 2n ** BigInt(MAX_BITS);
+  const big2PowerRot = 2n ** BigInt(rotationBits);
+
+  const [rotated, excess, shifted, bound] = Provable.witness(
+    Provable.Array(Field, 4),
+    () => {
+      const f = field.toBigInt();
+
+      // Obtain rotated output, excess, and shifted for the equation:
+      // f * 2^rot = excess * 2^64 + shifted
+      const { quotient: excess, remainder: shifted } = divideWithRemainder(
+        f * big2PowerRot,
+        big2Power64
+      );
+
+      // Compute rotated value as: rotated = excess + shifted
+      const rotated = shifted + excess;
+      // Compute bound to check excess < 2^rot
+      const bound = excess + big2Power64 - big2PowerRot;
+      return [rotated, excess, shifted, bound].map(Field.from);
+    }
+  );
 
-function witnessNextValue(current: Field) {
-  return Provable.witness(Field, () => new Field(current.toBigInt() >> 16n));
+  // Compute current row
+  Gates.rotate(
+    field,
+    rotated,
+    excess,
+    [
+      witnessSlices(bound, 52, 12), // bits 52-64
+      witnessSlices(bound, 40, 12), // bits 40-52
+      witnessSlices(bound, 28, 12), // bits 28-40
+      witnessSlices(bound, 16, 12), // bits 16-28
+    ],
+    [
+      witnessSlices(bound, 14, 2), // bits 14-16
+      witnessSlices(bound, 12, 2), // bits 12-14
+      witnessSlices(bound, 10, 2), // bits 10-12
+      witnessSlices(bound, 8, 2), // bits 8-10
+      witnessSlices(bound, 6, 2), // bits 6-8
+      witnessSlices(bound, 4, 2), // bits 4-6
+      witnessSlices(bound, 2, 2), // bits 2-4
+      witnessSlices(bound, 0, 2), // bits 0-2
+    ],
+    big2PowerRot
+  );
+  // Compute next row
+  Gates.rangeCheck64(shifted);
+  // Compute following row
+  Gates.rangeCheck64(excess);
+  return [rotated, excess, shifted];
 }

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

@@ -7,9 +7,10 @@ import {
   fieldWithRng,
 } from '../testing/equivalent.js';
 import { Fp, mod } from '../../bindings/crypto/finite_field.js';
-import { Field } from '../field.js';
+import { Field } from '../core.js';
 import { Gadgets } from './gadgets.js';
-import { Random } from '../testing/property.js';
+import { test, Random } from '../testing/property.js';
+import { Provable } from '../provable.js';
 
 let Bitwise = ZkProgram({
   name: 'bitwise',
@@ -21,6 +22,12 @@ let Bitwise = ZkProgram({
         return Gadgets.xor(a, b, 64);
       },
     },
+    rot: {
+      privateInputs: [Field],
+      method(a: Field) {
+        return Gadgets.rotate(a, 12, 'left');
+      },
+    },
   },
 });
 
@@ -35,14 +42,28 @@ let uint = (length: number) => fieldWithRng(Random.biguint(length));
   );
 });
 
+test(
+  Random.uint64,
+  Random.nat(64),
+  Random.boolean,
+  (x, n, direction, assert) => {
+    let z = Field(x);
+    let r1 = Fp.rot(x, n, direction ? 'left' : 'right');
+    Provable.runAndCheck(() => {
+      let f = Provable.witness(Field, () => z);
+      let r2 = Gadgets.rotate(f, n, direction ? 'left' : 'right');
+      Provable.asProver(() => assert(r1 === r2.toBigInt()));
+    });
+  }
+);
+
 let maybeUint64: Spec<bigint, Field> = {
   ...field,
   rng: Random.map(Random.oneOf(Random.uint64, Random.uint64.invalid), (x) =>
     mod(x, Field.ORDER)
   ),
 };
 
-// do a couple of proofs
 await equivalentAsync(
   { from: [maybeUint64, maybeUint64], to: field },
   { runs: 3 }
@@ -57,3 +78,37 @@ await equivalentAsync(
     return proof.publicOutput;
   }
 );
+
+await equivalentAsync({ from: [field], to: field }, { runs: 3 })(
+  (x) => {
+    if (x >= 2n ** 64n) throw Error('Does not fit into 64 bits');
+    return Fp.rot(x, 12, 'left');
+  },
+  async (x) => {
+    let proof = await Bitwise.rot(x);
+    return proof.publicOutput;
+  }
+);
+
+function testRot(
+  field: Field,
+  bits: number,
+  mode: 'left' | 'right',
+  result: Field
+) {
+  Provable.runAndCheck(() => {
+    let output = Gadgets.rotate(field, bits, mode);
+    output.assertEquals(result, `rot(${field}, ${bits}, ${mode})`);
+  });
+}
+
+testRot(Field(0), 0, 'left', Field(0));
+testRot(Field(0), 32, 'right', Field(0));
+testRot(Field(1), 1, 'left', Field(2));
+testRot(Field(1), 63, 'left', Field(9223372036854775808n));
+testRot(Field(256), 4, 'right', Field(16));
+testRot(Field(1234567890), 32, 'right', Field(5302428712241725440));
+testRot(Field(2651214356120862720), 32, 'right', Field(617283945));
+testRot(Field(1153202983878524928), 32, 'right', Field(268500993));
+testRot(Field(6510615555426900570n), 4, 'right', Field(11936128518282651045n));
+testRot(Field(6510615555426900570n), 4, 'right', Field(11936128518282651045n));

src/lib/gadgets/common.ts

@@ -0,0 +1,37 @@
+import { Provable } from '../provable.js';
+import { Field } from '../field.js';
+
+const MAX_BITS = 64 as const;
+
+export {
+  MAX_BITS,
+  assert,
+  witnessSlices,
+  witnessNextValue,
+  divideWithRemainder,
+};
+
+function assert(stmt: boolean, message?: string) {
+  if (!stmt) {
+    throw Error(message ?? 'Assertion failed');
+  }
+}
+
+function witnessSlices(f: Field, start: number, length: number) {
+  if (length <= 0) throw Error('Length must be a positive number');
+
+  return Provable.witness(Field, () => {
+    let n = f.toBigInt();
+    return new Field((n >> BigInt(start)) & ((1n << BigInt(length)) - 1n));
+  });
+}
+
+function witnessNextValue(current: Field) {
+  return Provable.witness(Field, () => new Field(current.toBigInt() >> 16n));
+}
+
+function divideWithRemainder(numerator: bigint, denominator: bigint) {
+  const quotient = numerator / denominator;
+  const remainder = numerator - denominator * quotient;
+  return { quotient, remainder };
+}