Merge pull request #1767 from o1-labs/feature/blake2b

BLAKE2b gadget

CHANGELOG-v2.md

@@ -35,3 +35,7 @@ Showing all changes since the last release (v.1.5.0)
 ### Added
 
 - New method `toCanonical()` in the `Provable<T>` interface to protect against incompleteness of certain operations on malicious witness inputs https://github.com/o1-labs/o1js/pull/1759
+- `divMod64()` division modulo 2^64 that returns the remainder and quotient of the operation
+- `addMod64()` addition modulo 2^64
+- Bitwise OR via `{UInt32, UInt64}.or()`
+- **BLAKE2B hash function** gadget [#1285](https://github.com/o1-labs/o1js/pull/1285)

src/examples/crypto/blake2b/blake2b.ts

@@ -0,0 +1,18 @@
+import { Bytes, Gadgets, ZkProgram } from 'o1js';
+
+export { BLAKE2BProgram, Bytes12 };
+
+class Bytes12 extends Bytes(12) {}
+
+let BLAKE2BProgram = ZkProgram({
+  name: 'blake2b',
+  publicOutput: Bytes(32),
+  methods: {
+    blake2b: {
+      privateInputs: [Bytes12],
+      async method(xs: Bytes12) {
+        return { publicOutput: Gadgets.BLAKE2B.hash(xs, 32) };
+      },
+    },
+  },
+});

src/examples/crypto/blake2b/run.ts

@@ -0,0 +1,23 @@
+import { Bytes12, BLAKE2BProgram } from './blake2b.js';
+
+console.time('compile');
+await BLAKE2BProgram.compile();
+console.timeEnd('compile');
+
+let preimage = Bytes12.fromString('hello world!');
+
+console.log('blake2b rows:', (await BLAKE2BProgram.analyzeMethods()).blake2b.rows);
+
+console.time('prove');
+let { proof } = await BLAKE2BProgram.blake2b(preimage);
+console.timeEnd('prove');
+let isValid = await BLAKE2BProgram.verify(proof);
+
+console.log('digest:', proof.publicOutput.toHex());
+
+if (
+  proof.publicOutput.toHex() !==
+  '4fccfb4d98d069558aa93e9565f997d81c33b080364efd586e77a433ddffc5e2'
+)
+  throw new Error('Invalid blake2b digest!');
+if (!isValid) throw new Error('Invalid proof');

src/lib/provable/crypto/hash.ts

@@ -138,4 +138,10 @@ const Hash = {
       return Keccak.preNist(512, bytes);
     },
   },
+
+  BLAKE2B: {
+    hash(bytes: Bytes) {
+      return Gadgets.BLAKE2B.hash(bytes);
+    },
+  },
 };

src/lib/provable/gadgets/arithmetic.ts

@@ -2,9 +2,9 @@ import { provableTuple } from '../types/struct.js';
 import { Field } from '../wrapped.js';
 import { assert } from '../../util/errors.js';
 import { Provable } from '../provable.js';
-import { rangeCheck32, rangeCheckN } from './range-check.js';
+import { rangeCheck32, rangeCheck64, rangeCheckN } from './range-check.js';
 
-export { divMod32, addMod32 };
+export { divMod32, addMod32, divMod64, addMod64 };
 
 function divMod32(n: Field, nBits = 64) {
   assert(
@@ -56,3 +56,57 @@ function divMod32(n: Field, nBits = 64) {
 function addMod32(x: Field, y: Field) {
   return divMod32(x.add(y), 33).remainder;
 }
+
+function divMod64(n: Field, nBits = 128) {
+  assert(
+    nBits >= 0 && nBits < 255,
+    `nBits must be in the range [0, 255), got ${nBits}`
+  );
+
+  // calculate the number of bits allowed for the quotient to avoid overflow
+  const quotientBits = Math.max(0, nBits - 64);
+
+  if (n.isConstant()) {
+    assert(
+      n.toBigInt() < 1n << BigInt(nBits),
+      `n needs to fit into ${nBits} bits, but got ${n.toBigInt()}`
+    );
+    let nBigInt = n.toBigInt();
+    let q = nBigInt >> 64n;
+    let r = nBigInt - (q << 64n);
+    return {
+      remainder: new Field(r),
+      quotient: new Field(q),
+    };
+  }
+
+  let [quotient, remainder] = Provable.witness(
+    provableTuple([Field, Field]),
+    () => {
+      let nBigInt = n.toBigInt();
+      let q = nBigInt >> 64n;
+      let r = nBigInt - (q << 64n);
+      return [q, r] satisfies [bigint, bigint];
+    }
+  );
+
+  if (quotientBits === 1) {
+    quotient.assertBool();
+  } else if (quotientBits === 64) {
+    rangeCheck64(quotient);
+  } else {
+    rangeCheckN(quotientBits, quotient);
+  }
+  rangeCheck64(remainder);
+
+  n.assertEquals(quotient.mul(1n << 64n).add(remainder));
+
+  return {
+    remainder,
+    quotient,
+  };
+}
+
+function addMod64(x: Field, y: Field) {
+  return divMod64(x.add(y), 65).remainder;
+}

src/lib/provable/gadgets/bitwise.ts

@@ -13,6 +13,7 @@ export {
   rotate64,
   rotate32,
   and,
+  or,
   rightShift64,
   leftShift64,
   leftShift32,
@@ -173,6 +174,10 @@ function and(a: Field, b: Field, length: number) {
   return outputAnd;
 }
 
+function or(a: Field, b: Field, length: number) {
+  return not(and(not(a, length), not(b, length), length), length);
+}
+
 function rotate64(
   field: Field,
   bits: number,

src/lib/provable/gadgets/blake2b.ts

@@ -0,0 +1,236 @@
+// https://datatracker.ietf.org/doc/html/rfc7693.html
+import { UInt64, UInt8 } from '../int.js';
+import { FlexibleBytes } from '../bytes.js';
+import { Bytes } from '../wrapped-classes.js';
+import { Gadgets } from './gadgets.js';
+import { assert } from '../../util/errors.js';
+import { Provable } from '../provable.js';
+import { wordToBytes } from './bit-slices.js';
+
+export { BLAKE2B };
+
+/**
+ * IV[0..7]  Initialization Vector (constant).
+ * SIGMA[0..9]  Message word permutations (constant).
+ * p[0..7]  Parameter block (defines hash and key sizes).
+ * m[0..15]  Sixteen words of a single message block.
+ * h[0..7]  Internal state of the hash.
+ * d[0..dd-1]  Padded input blocks.  Each has "bb" bytes.
+ * t  Message byte offset at the end of the current block.
+ * f  Flag indicating the last block.
+ *
+ * All mathematical operations are on 64-bit words in BLAKE2b.
+ *
+ * Byte (octet) streams are interpreted as words in little-endian order,
+ * with the least-significant byte first.
+ */
+
+type State = {
+  h: UInt64[];
+  t: [bigint, bigint];
+  buf: UInt8[];
+  buflen: number;
+  outlen: number;
+};
+
+const BLAKE2BConstants = {
+  IV: [
+    UInt64.from(0x6a09e667f3bcc908n),
+    UInt64.from(0xbb67ae8584caa73bn),
+    UInt64.from(0x3c6ef372fe94f82bn),
+    UInt64.from(0xa54ff53a5f1d36f1n),
+    UInt64.from(0x510e527fade682d1n),
+    UInt64.from(0x9b05688c2b3e6c1fn),
+    UInt64.from(0x1f83d9abfb41bd6bn),
+    UInt64.from(0x5be0cd19137e2179n),
+  ],
+
+  SIGMA: [
+    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
+    [14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3],
+    [11, 8, 12, 0, 5, 2, 15, 13, 10, 14, 3, 6, 7, 1, 9, 4],
+    [7, 9, 3, 1, 13, 12, 11, 14, 2, 6, 5, 10, 4, 0, 15, 8],
+    [9, 0, 5, 7, 2, 4, 10, 15, 14, 1, 11, 12, 6, 8, 3, 13],
+    [2, 12, 6, 10, 0, 11, 8, 3, 4, 13, 7, 5, 15, 14, 1, 9],
+    [12, 5, 1, 15, 14, 13, 4, 10, 0, 7, 6, 3, 9, 2, 8, 11],
+    [13, 11, 7, 14, 12, 1, 3, 9, 5, 0, 15, 4, 8, 6, 2, 10],
+    [6, 15, 14, 9, 11, 3, 0, 8, 12, 2, 13, 7, 1, 4, 10, 5],
+    [10, 2, 8, 4, 7, 6, 1, 5, 15, 11, 9, 14, 3, 12, 13, 0],
+  ],
+};
+
+const BLAKE2B = {
+  hash(data: FlexibleBytes, digestLength = 64) {
+    assert(
+      digestLength >= 1 && digestLength <= 64,
+      `digestLength must be in the range [1, 64], got ${digestLength}`
+    );
+    assert(
+      data.length >= 0 && data.length < 2 ** 128,
+      `data byte length must be in the range [0, 2**128), got ${data.length}`
+    );
+    const state = initialize(digestLength);
+    const updatedState = update(state, Bytes.from(data).bytes);
+    const out = final(updatedState);
+    return Bytes.from(out);
+  },
+  get IV() {
+    return BLAKE2BConstants.IV;
+  },
+};
+
+function G(
+  v: UInt64[],
+  a: number,
+  b: number,
+  c: number,
+  d: number,
+  x: UInt64,
+  y: UInt64
+) {
+  v[a] = UInt64.Unsafe.fromField(
+    Gadgets.divMod64(v[a].value.add(v[b].value.add(x.value)), 128).remainder
+  );
+  v[d] = v[d].xor(v[a]).rotate(32, 'right');
+
+  v[c] = UInt64.Unsafe.fromField(
+    Gadgets.divMod64(v[c].value.add(v[d].value), 128).remainder
+  );
+  v[b] = v[b].xor(v[c]).rotate(24, 'right');
+
+  v[a] = UInt64.Unsafe.fromField(
+    Gadgets.divMod64(v[a].value.add(v[b].value.add(y.value)), 128).remainder
+  );
+  v[d] = v[d].xor(v[a]).rotate(16, 'right');
+
+  v[c] = UInt64.Unsafe.fromField(
+    Gadgets.divMod64(v[c].value.add(v[d].value), 128).remainder
+  );
+  v[b] = v[b].xor(v[c]).rotate(63, 'right');
+}
+
+/**
+ * Compression function. "last" flag indicates last block.
+ * @param {State} state
+ * @param {boolean} last
+ */
+function compress(state: State, last: boolean): State {
+  const { h, t, buf } = state;
+  const v = h.concat(BLAKE2B.IV); // initalize local work vector. First half from state and second half from IV.
+
+  v[12] = v[12].xor(UInt64.from(t[0])); // low word of the offset
+  v[13] = v[13].xor(UInt64.from(t[1])); // high word of the offset
+
+  if (last) {
+    // last block flag set ?
+    v[14] = v[14].not();
+  }
+
+  const m: UInt64[] = [];
+  for (let i = 0; i < 16; i++) {
+    // get little-endian words
+    m.push(
+      UInt64.Unsafe.fromField(
+        buf[i * 8].value
+          .add(buf[i * 8 + 1].value.mul(1n << 8n))
+          .add(buf[i * 8 + 2].value.mul(1n << 16n))
+          .add(buf[i * 8 + 3].value.mul(1n << 24n))
+          .add(buf[i * 8 + 4].value.mul(1n << 32n))
+          .add(buf[i * 8 + 5].value.mul(1n << 40n))
+          .add(buf[i * 8 + 6].value.mul(1n << 48n))
+          .add(buf[i * 8 + 7].value.mul(1n << 56n))
+          .seal()
+      )
+    );
+  }
+
+  for (let i = 0; i < 12; i++) {
+    // twelve rounds
+    const s = BLAKE2BConstants.SIGMA[i % 10];
+    G(v, 0, 4, 8, 12, m[s[0]], m[s[1]]);
+    G(v, 1, 5, 9, 13, m[s[2]], m[s[3]]);
+    G(v, 2, 6, 10, 14, m[s[4]], m[s[5]]);
+    G(v, 3, 7, 11, 15, m[s[6]], m[s[7]]);
+
+    G(v, 0, 5, 10, 15, m[s[8]], m[s[9]]);
+    G(v, 1, 6, 11, 12, m[s[10]], m[s[11]]);
+    G(v, 2, 7, 8, 13, m[s[12]], m[s[13]]);
+    G(v, 3, 4, 9, 14, m[s[14]], m[s[15]]);
+  }
+
+  for (let i = 0; i < 8; i++) {
+    // XOR the two halves
+    h[i] = v[i].xor(v[i + 8]).xor(h[i]);
+  }
+  return state;
+}
+
+/**
+ * Initializes the state with the given digest length.
+ *
+ * @param {number} outlen - Digest length in bits
+ * @returns {State}
+ */
+function initialize(outlen: number): State {
+  const h = BLAKE2B.IV.slice(); // shallow copy IV to h
+  h[0] = UInt64.from(0x01010000).xor(UInt64.from(outlen)).xor(h[0]); // state "param block"
+
+  return {
+    h,
+    t: [0n, 0n],
+    buf: [],
+    buflen: 0,
+    outlen,
+  };
+}
+
+/**
+ * Updates hash state
+ * @param {State} state
+ * @param {UInt8[]} input
+ * @returns {State} updated state
+ */
+function update(state: State, input: UInt8[]): State {
+  for (let i = 0; i < input.length; i++) {
+    if (state.buflen === 128) {
+      // buffer full ?
+      state.t[0] = UInt64.from(state.t[0])
+        .addMod64(UInt64.from(state.buflen))
+        .toBigInt(); // add counters
+      if (state.t[0] < state.buflen) {
+        // carry overflow ?
+        state.t[1] = UInt64.from(state.t[1]).addMod64(UInt64.one).toBigInt(); // high word
+      }
+      state = compress(state, false); // compress (not last)
+      state.buflen = 0; // counter to zero
+    }
+    state.buf[state.buflen++] = input[i];
+  }
+  return state;
+}
+
+/**
+ * Finalizes the hash state and returns digest
+ * @param {State} state
+ * @returns {UInt8[]} digest
+ */
+function final(state: State): UInt8[] {
+  state.t[0] = UInt64.from(state.t[0])
+    .addMod64(UInt64.from(state.buflen))
+    .toBigInt(); // add counters
+  if (state.t[0] < state.buflen) {
+    // carry overflow ?
+    state.t[1] = UInt64.from(state.t[1]).addMod64(UInt64.one).toBigInt(); // high word
+  }
+
+  while (state.buflen < 128) {
+    state.buf[state.buflen++] = UInt8.from(0); // fill up with zeroes
+  }
+  compress(state, true);
+
+  // little endian convert and store
+  const out: UInt8[] = state.h
+    .slice(0, state.outlen / 8)
+    .flatMap((x) => wordToBytes(x.value));
+  return out;
+}