o1js/src/lib/provable/test/bitwise.unit-test.ts

TypeScript framework for zk-SNARKs and zkApps

import { ZkProgram } from '../../proof-system/zkprogram.js';
import {
  equivalentProvable as equivalent,
  equivalentAsync,
  field,
  fieldWithRng,
} from '../../testing/equivalent.js';
import { Fp, mod } from '../../../bindings/crypto/finite-field.js';
import { Field } from '../wrapped.js';
import { Gadgets } from '../gadgets/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,
  rng: Random.map(Random.oneOf(Random.field, Random.field.invalid), (x) =>
    mod(x, Field.ORDER)
  ),
};

let uint = (length: number) => fieldWithRng(Random.biguint(length));

let Bitwise = ZkProgram({
  name: 'bitwise',
  publicOutput: Field,
  methods: {
    xor: {
      privateInputs: [Field, Field],
      async method(a: Field, b: Field) {
        return Gadgets.xor(a, b, 254);
      },
    },
    notUnchecked: {
      privateInputs: [Field],
      async method(a: Field) {
        return Gadgets.not(a, 254, false);
      },
    },
    notChecked: {
      privateInputs: [Field],
      async method(a: Field) {
        return Gadgets.not(a, 254, true);
      },
    },
    and: {
      privateInputs: [Field, Field],
      async method(a: Field, b: Field) {
        return Gadgets.and(a, b, 64);
      },
    },
    rot32: {
      privateInputs: [Field],
      async method(a: Field) {
        return Gadgets.rotate32(a, 12, 'left');
      },
    },
    rot64: {
      privateInputs: [Field],
      async method(a: Field) {
        return Gadgets.rotate64(a, 12, 'left');
      },
    },
    leftShift64: {
      privateInputs: [Field],
      async method(a: Field) {
        return Gadgets.leftShift64(a, 12);
      },
    },
    leftShift32: {
      privateInputs: [Field],
      async method(a: Field) {
        Gadgets.rangeCheck32(a);
        return Gadgets.leftShift32(a, 12);
      },
    },
    rightShift64: {
      privateInputs: [Field],
      async method(a: Field) {
        return Gadgets.rightShift64(a, 12);
      },
    },
  },
});
await Bitwise.compile();

[2, 4, 8, 16, 32, 64, 128].forEach((length) => {
  equivalent({ from: [uint(length), uint(length)], to: field })(
    (x, y) => x ^ y,
    (x, y) => Gadgets.xor(x, y, length)
  );
  equivalent({ from: [uint(length), uint(length)], to: field })(
    (x, y) => x & y,
    (x, y) => Gadgets.and(x, y, length)
  );
  // NOT unchecked
  equivalent({ from: [uint(length)], to: field })(
    (x) => Fp.not(x, length),
    (x) => Gadgets.not(x, length, false)
  );
  // NOT checked
  equivalent({ from: [uint(length)], to: field })(
    (x) => Fp.not(x, length),
    (x) => Gadgets.not(x, length, true)
  );
});

[2, 4, 8, 16, 32, 64].forEach((length) => {
  equivalent({ from: [uint(length)], to: field })(
    (x) => Fp.rot(x, 12n, 'left'),
    (x) => Gadgets.rotate64(x, 12, 'left')
  );
  equivalent({ from: [uint(length)], to: field })(
    (x) => Fp.leftShift(x, 12),
    (x) => Gadgets.leftShift64(x, 12)
  );
  equivalent({ from: [uint(length)], to: field })(
    (x) => Fp.rightShift(x, 12),
    (x) => Gadgets.rightShift64(x, 12)
  );
});

[2, 4, 8, 16, 32].forEach((length) => {
  equivalent({ from: [uint(length)], to: field })(
    (x) => Fp.rot(x, 12n, 'left', 32n),
    (x) => Gadgets.rotate32(x, 12, 'left')
  );
  equivalent({ from: [uint(length)], to: field })(
    (x) => Fp.leftShift(x, 12, 32),
    (x) => Gadgets.leftShift32(x, 12)
  );
});

const runs = 2;

await equivalentAsync({ from: [uint(64), uint(64)], to: field }, { runs })(
  (x, y) => {
    return x ^ y;
  },
  async (x, y) => {
    let proof = await Bitwise.xor(x, y);
    return proof.publicOutput;
  }
);

await equivalentAsync({ from: [maybeField], to: field }, { runs })(
  (x) => {
    return Fp.not(x, 254);
  },
  async (x) => {
    let proof = await Bitwise.notUnchecked(x);
    return proof.publicOutput;
  }
);
await equivalentAsync({ from: [maybeField], to: field }, { runs })(
  (x) => {
    if (x > 2n ** 254n) throw Error('Does not fit into 254 bit');
    return Fp.not(x, 254);
  },
  async (x) => {
    let proof = await Bitwise.notChecked(x);
    return proof.publicOutput;
  }
);

await equivalentAsync({ from: [maybeField, maybeField], to: field }, { runs })(
  (x, y) => {
    if (x >= 2n ** 64n || y >= 2n ** 64n)
      throw Error('Does not fit into 64 bits');
    return x & y;
  },
  async (x, y) => {
    let proof = await Bitwise.and(x, y);
    return proof.publicOutput;
  }
);

await equivalentAsync({ from: [field], to: field }, { runs })(
  (x) => {
    if (x >= 2n ** 64n) throw Error('Does not fit into 64 bits');
    return Fp.rot(x, 12n, 'left');
  },
  async (x) => {
    let proof = await Bitwise.rot64(x);
    return proof.publicOutput;
  }
);

await equivalentAsync({ from: [uint(32)], to: uint(32) }, { runs })(
  (x) => {
    return Fp.rot(x, 12n, 'left', 32n);
  },
  async (x) => {
    let proof = await Bitwise.rot32(x);
    return proof.publicOutput;
  }
);

await equivalentAsync({ from: [field], to: field }, { runs })(
  (x) => {
    if (x >= 2n ** 64n) throw Error('Does not fit into 64 bits');
    return Fp.leftShift(x, 12);
  },
  async (x) => {
    let proof = await Bitwise.leftShift64(x);
    return proof.publicOutput;
  }
);

await equivalentAsync({ from: [field], to: field }, { runs })(
  (x) => {
    if (x >= 1n << 32n) throw Error('Does not fit into 32 bits');
    return Fp.leftShift(x, 12, 32);
  },
  async (x) => {
    let proof = await Bitwise.leftShift32(x);
    return proof.publicOutput;
  }
);

await equivalentAsync({ from: [field], to: field }, { runs })(
  (x) => {
    if (x >= 2n ** 64n) throw Error('Does not fit into 64 bits');
    return Fp.rightShift(x, 12);
  },
  async (x) => {
    let proof = await Bitwise.rightShift64(x);
    return proof.publicOutput;
  }
);

// check that gate chains stay intact

function xorChain(bits: number) {
  return repeat(Math.ceil(bits / 16), 'Xor16').concat('Zero');
}

constraintSystem.fromZkProgram(
  Bitwise,
  'xor',
  ifNotAllConstant(contains(xorChain(254)))
);

constraintSystem.fromZkProgram(
  Bitwise,
  'notChecked',
  ifNotAllConstant(contains(xorChain(254)))
);

constraintSystem.fromZkProgram(
  Bitwise,
  'notUnchecked',
  ifNotAllConstant(contains('Generic'))
);

constraintSystem.fromZkProgram(
  Bitwise,
  'and',
  ifNotAllConstant(contains(xorChain(64)))
);

let rotChain: GateType[] = ['Rot64', 'RangeCheck0'];
let isJustRotate = ifNotAllConstant(
  and(contains(rotChain), withoutGenerics(equals(rotChain)))
);

constraintSystem.fromZkProgram(Bitwise, 'rot64', isJustRotate);
constraintSystem.fromZkProgram(Bitwise, 'leftShift64', isJustRotate);
constraintSystem.fromZkProgram(Bitwise, 'rightShift64', isJustRotate);