Merge pull request #1694 from o1-labs/feat/fix-merkle-map

Fix potential collisions with merkle map indicies

CHANGELOG.md

@@ -24,6 +24,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
   - Feature flags are required to tell Pickles what proof structure it should expect when side loading dynamic proofs and verification keys.
   - `FeatureFlags` is now exported and provides a set of helper functions to compute feature flags correctly.
 
+### Deprecated
+
+- `MerkleMap.computeRootAndKey()` deprecated in favor of `MerkleMap.computeRootAndKeyV2()` due to a potential issue of computing hash collisions in key indicies https://github.com/o1-labs/o1js/pull/1694
+
 ## [1.3.1](https://github.com/o1-labs/o1js/compare/1ad7333e9e...40c597775) - 2024-06-11
 
 ### Breaking Changes

src/lib/provable/merkle-map.ts

@@ -13,6 +13,10 @@ class MerkleMap {
   /**
    * Creates a new, empty Merkle Map.
    * @returns A new MerkleMap
+   * @example
+   * ```ts
+   * const merkleMap = new MerkleMap();
+   * ```
    */
   constructor() {
     this.tree = new MerkleTree(256);
@@ -22,6 +26,13 @@ class MerkleMap {
     // the bit map is reversed to make reconstructing the key during proving more convenient
     let bits = BinableFp.toBits(key.toBigInt()).reverse();
 
+    // Make sure that the key fits in 254 bits, in order to avoid collisions since the Pasta field modulus is smaller than 2^255
+    if (bits[0]) {
+      throw Error(
+        'Key must be less than 2^254, to avoid collisions in the field modulus. Please use a smaller key.'
+      );
+    }
+
     let n = 0n;
     for (let i = bits.length - 1; i >= 0; i--) {
       n = (n << 1n) | BigInt(bits[i]);
@@ -34,6 +45,12 @@ class MerkleMap {
    * Sets a key of the merkle map to a given value.
    * @param key The key to set in the map.
    * @param value The value to set.
+   * @example
+   * ```ts
+   * const key = Field(5);
+   * const value = Field(10);
+   * merkleMap.set(key, value);
+   * ```
    */
   set(key: Field, value: Field) {
     const index = this._keyToIndex(key);
@@ -44,6 +61,12 @@ class MerkleMap {
    * Returns a value given a key. Values are by default Field(0).
    * @param key The key to get the value from.
    * @returns The value stored at the key.
+   * @example
+   * ```ts
+   * const key = Field(5);
+   * const value = merkleMap.get(key);
+   * console.log(value); // Output: the value at key 5 or Field(0) if key does not exist
+   * ```
    */
   get(key: Field) {
     const index = this._keyToIndex(key);
@@ -53,6 +76,10 @@ class MerkleMap {
   /**
    * Returns the root of the Merkle Map.
    * @returns The root of the Merkle Map.
+   * @example
+   * ```ts
+   * const root = merkleMap.getRoot();
+   * ```
    */
   getRoot() {
     return this.tree.getRoot();
@@ -62,6 +89,11 @@ class MerkleMap {
    * Returns a circuit-compatible witness (also known as [Merkle Proof or Merkle Witness](https://computersciencewiki.org/index.php/Merkle_proof)) for the given key.
    * @param key The key to make a witness for.
    * @returns A MerkleMapWitness, which can be used to assert changes to the MerkleMap, and the witness's key.
+   * @example
+   * ```ts
+   * const key = Field(5);
+   * const witness = merkleMap.getWitness(key);
+   * ```
    */
   getWitness(key: Field) {
     const index = this._keyToIndex(key);
@@ -82,11 +114,38 @@ class MerkleMapWitness extends CircuitValue {
   }
 
   /**
-   * computes the merkle tree root for a given value and the key for this witness
+   * @deprecated This method is deprecated and will be removed in the next release. Please use {@link computeRootAndKeyV2} instead.
+   */
+  computeRootAndKey(value: Field) {
+    let hash = value;
+
+    const isLeft = this.isLefts;
+    const siblings = this.siblings;
+
+    let key = Field(0);
+
+    for (let i = 0; i < 255; i++) {
+      const left = Provable.if(isLeft[i], hash, siblings[i]);
+      const right = Provable.if(isLeft[i], siblings[i], hash);
+      hash = Poseidon.hash([left, right]);
+
+      const bit = Provable.if(isLeft[i], Field(0), Field(1));
+
+      key = key.mul(2).add(bit);
+    }
+
+    return [hash, key];
+  }
+
+  /**
+   * Computes the merkle tree root for a given value and the key for this witness
    * @param value The value to compute the root for.
    * @returns A tuple of the computed merkle root, and the key that is connected to the path updated by this witness.
    */
-  computeRootAndKey(value: Field) {
+  computeRootAndKeyV2(value: Field) {
+    // Check that the computed key is less than 2^254, in order to avoid collisions since the Pasta field modulus is smaller than 2^255
+    this.isLefts[0].assertTrue();
+
     let hash = value;
 
     const isLeft = this.isLefts;

src/lib/provable/test/merkle-tree.unit-test.ts

@@ -37,7 +37,7 @@ console.log(
       let mapWitness = Provable.witness(MerkleMapWitness, () =>
         throwError('unused')
       );
-      let [actualRoot, actualKey] = mapWitness.computeRootAndKey(value);
+      let [actualRoot, actualKey] = mapWitness.computeRootAndKeyV2(value);
       key.assertEquals(actualKey);
       root.assertEquals(actualRoot);
     }
@@ -71,11 +71,11 @@ console.log(
       let mapWitness = Provable.witness(MerkleMapWitness, () =>
         throwError('unused')
       );
-      let [actualRoot, actualKey] = mapWitness.computeRootAndKey(oldValue);
+      let [actualRoot, actualKey] = mapWitness.computeRootAndKeyV2(oldValue);
       key.assertEquals(actualKey);
       root.assertEquals(actualRoot);
 
-      let [_newRoot] = mapWitness.computeRootAndKey(value);
+      let [_newRoot] = mapWitness.computeRootAndKeyV2(value);
     }
   )
 );
@@ -296,7 +296,15 @@ test(Random.bool, Random.field, Random.field, (b, x, y) => {
 
 test(Random.field, (key) => {
   let map = new MerkleMap();
-  let witness = map.getWitness(Field(key));
-  let [, calculatedKey] = witness.computeRootAndKey(Field(0));
-  expect(calculatedKey.toBigInt()).toEqual(key);
+  // Check that the key fits in 254 bits, if it doesn't, we should throw an error (since the Pasta field modulus is smaller than 2^255)
+  if (key > 2n ** 254n) {
+    expect(() => {
+      let witness = map.getWitness(Field(key));
+      witness.computeRootAndKeyV2(Field(0));
+    }).toThrowError();
+  } else {
+    let witness = map.getWitness(Field(key));
+    let [, calculatedKey] = witness.computeRootAndKeyV2(Field(0));
+    expect(calculatedKey.toBigInt()).toEqual(key);
+  }
 });