Fix Merkle map keys

CHANGELOG.md

@@ -23,6 +23,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
   - Fixes soundness of ECDSA; slightly increases its constraints from ~28k to 29k
   - Breaks circuits that used EC addition, like ECDSA
 
+### Fixed
+
+- Fix error when computing Merkle map witnesses, introduced in the last version due to the `toBits()` change https://github.com/o1-labs/o1js/pull/1559
+
 ## [0.18.0](https://github.com/o1-labs/o1js/compare/74948acac...1b6fd8b8e) - 2024-04-09
 
 ### Breaking changes

src/lib/provable/merkle-map.ts

@@ -3,53 +3,33 @@ import { Field, Bool } from './wrapped.js';
 import { Poseidon } from './crypto/poseidon.js';
 import { MerkleTree, MerkleWitness } from './merkle-tree.js';
 import { Provable } from './provable.js';
+import { BinableFp } from '../../mina-signer/src/field-bigint.js';
 
-const bits = 255;
-const printDebugs = false;
+export { MerkleMap, MerkleMapWitness };
 
-export class MerkleMap {
-  tree: InstanceType<typeof MerkleTree>;
-
-  // ------------------------------------------------
+class MerkleMap {
+  tree: MerkleTree;
 
   /**
    * Creates a new, empty Merkle Map.
    * @returns A new MerkleMap
    */
   constructor() {
-    if (bits > 255) {
-      throw Error('bits must be <= 255');
-    }
-    if (bits !== 255) {
-      console.warn(
-        'bits set to',
-        bits + '. Should be set to 255 in production to avoid collisions'
-      );
-    }
-    this.tree = new MerkleTree(bits + 1);
+    this.tree = new MerkleTree(256);
   }
 
-  // ------------------------------------------------
-
   _keyToIndex(key: Field) {
     // the bit map is reversed to make reconstructing the key during proving more convenient
-    let keyBits = key
-      .toBits()
-      .slice(0, bits)
-      .reverse()
-      .map((b) => b.toBoolean());
+    let bits = BinableFp.toBits(key.toBigInt()).reverse();
 
     let n = 0n;
-    for (let i = 0; i < keyBits.length; i++) {
-      const b = keyBits[i] ? 1 : 0;
-      n += 2n ** BigInt(i) * BigInt(b);
+    for (let i = bits.length - 1; i >= 0; i--) {
+      n = (n << 1n) | BigInt(bits[i]);
     }
 
     return n;
   }
 
-  // ------------------------------------------------
-
   /**
    * Sets a key of the merkle map to a given value.
    * @param key The key to set in the map.
@@ -60,8 +40,6 @@ export class MerkleMap {
     this.tree.setLeaf(index, value);
   }
 
-  // ------------------------------------------------
-
   /**
    * Returns a value given a key. Values are by default Field(0).
    * @param key The key to get the value from.
@@ -72,8 +50,6 @@ export class MerkleMap {
     return this.tree.getNode(0, index);
   }
 
-  // ------------------------------------------------
-
   /**
    * Returns the root of the Merkle Map.
    * @returns The root of the Merkle Map.
@@ -89,34 +65,15 @@ export class MerkleMap {
    */
   getWitness(key: Field) {
     const index = this._keyToIndex(key);
-    class MyMerkleWitness extends MerkleWitness(bits + 1) {}
+    class MyMerkleWitness extends MerkleWitness(256) {}
     const witness = new MyMerkleWitness(this.tree.getWitness(index));
-
-    if (printDebugs) {
-      // witness bits and key bits should be the reverse of each other, so
-      // we can calculate the key during recursively traversing the path
-      console.log(
-        'witness bits',
-        witness.isLeft.map((l) => (l.toBoolean() ? '0' : '1')).join(', ')
-      );
-      console.log(
-        'key bits',
-        key
-          .toBits()
-          .slice(0, bits)
-          .map((l) => (l.toBoolean() ? '1' : '0'))
-          .join(', ')
-      );
-    }
     return new MerkleMapWitness(witness.isLeft, witness.path);
   }
 }
 
-// =======================================================
-
-export class MerkleMapWitness extends CircuitValue {
-  @arrayProp(Bool, bits) isLefts: Bool[];
-  @arrayProp(Field, bits) siblings: Field[];
+class MerkleMapWitness extends CircuitValue {
+  @arrayProp(Bool, 255) isLefts: Bool[];
+  @arrayProp(Field, 255) siblings: Field[];
 
   constructor(isLefts: Bool[], siblings: Field[]) {
     super();
@@ -137,7 +94,7 @@ export class MerkleMapWitness extends CircuitValue {
 
     let key = Field(0);
 
-    for (let i = 0; i < bits; i++) {
+    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]);

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

@@ -2,6 +2,7 @@ import { Bool, Field } from '../wrapped.js';
 import { maybeSwap } from '../merkle-tree.js';
 import { Random, test } from '../../testing/property.js';
 import { expect } from 'expect';
+import { MerkleMap } from '../merkle-map.js';
 
 test(Random.bool, Random.field, Random.field, (b, x, y) => {
   let [x0, y0] = maybeSwap(Bool(!!b), Field(x), Field(y));
@@ -15,3 +16,10 @@ test(Random.bool, Random.field, Random.field, (b, x, y) => {
     expect(y0.toBigInt()).toEqual(x);
   }
 });
+
+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);
+});