Merge pull request #1671 from o1-labs/feature/indexed-merkle-map-2

IndexedMerkleMap: Support 0 and -1 keys

src/lib/provable/field.ts

@@ -654,7 +654,7 @@ class Field {
    * Assert that this {@link Field} is less than another "field-like" value.
    *
    * Note: This uses fewer constraints than `x.lessThan(y).assertTrue()`.
-   * See {@link Field.lessThan} for more details.
+   * See {@link lessThan} for more details.
    *
    * **Important**: If an assertion fails, the code throws an error.
    *

src/lib/provable/merkle-tree-indexed.ts

@@ -49,6 +49,10 @@ export { Leaf };
  *   }
  * })
  * ```
+ *
+ * Initially, every `IndexedMerkleMap` is populated by a single key-value pair: `(0, 0)`. The value for key `0` can be updated like any other.
+ * When keys and values are hash outputs, `(0, 0)` can serve as a convenient way to represent a dummy update to the tree, since 0 is not
+ * effciently computable as a hash image, and this update doesn't affect the Merkle root.
  */
 function IndexedMerkleMap(height: number): typeof IndexedMerkleMapBase {
   assert(height > 0, 'height must be positive');
@@ -98,19 +102,6 @@ class IndexedMerkleMapBase {
     readonly sortedLeaves: StoredLeaf[];
   }>;
 
-  clone() {
-    let cloned = new (this.constructor as typeof IndexedMerkleMapBase)();
-    cloned.root = this.root;
-    cloned.length = this.length;
-    cloned.data.updateAsProver(({ nodes, sortedLeaves }) => {
-      return {
-        nodes: nodes.map((row) => [...row]),
-        sortedLeaves: [...sortedLeaves],
-      };
-    });
-    return cloned;
-  }
-
   // we'd like to do `abstract static provable` here but that's not supported
   static provable: Provable<
     IndexedMerkleMapBase,
@@ -140,11 +131,31 @@ class IndexedMerkleMapBase {
   static _firstLeaf = {
     key: 0n,
     value: 0n,
-    // maximum, which is always greater than any key that is a hash
-    nextKey: Field.ORDER - 1n,
+    // the 0 key encodes the minimum and maximum at the same time
+    // so, if a second node is inserted, it will get `nextKey = 0`, and thus point to the first node
+    nextKey: 0n,
     index: 0,
   };
 
+  /**
+   * Clone the entire Merkle map.
+   *
+   * This method is provable.
+   */
+  clone() {
+    let cloned = new (this.constructor as typeof IndexedMerkleMapBase)();
+    cloned.root = this.root;
+    cloned.length = this.length;
+    cloned.data.updateAsProver(() => {
+      let { nodes, sortedLeaves } = this.data.get();
+      return {
+        nodes: nodes.map((row) => [...row]),
+        sortedLeaves: [...sortedLeaves],
+      };
+    });
+    return cloned;
+  }
+
   /**
    * Insert a new leaf `(key, value)`.
    *
@@ -161,10 +172,13 @@ class IndexedMerkleMapBase {
     // prove that the key doesn't exist yet by presenting a valid low node
     let low = Provable.witness(Leaf, () => this._findLeaf(key).low);
     let lowPath = this._proveInclusion(low, 'Invalid low node (root)');
-    low.key.assertLessThan(key, 'Invalid low node (key)');
-
     // if the key does exist, we have lowNode.nextKey == key, and this line fails
-    key.assertLessThan(low.nextKey, 'Key already exists in the tree');
+    assertStrictlyBetween(
+      low.key,
+      key,
+      low.nextKey,
+      'Key already exists in the tree'
+    );
 
     // at this point, we know that we have a valid insertion; so we can mutate internal data
 
@@ -230,8 +244,7 @@ class IndexedMerkleMapBase {
     // prove whether the key exists or not, by showing a valid low node
     let { low, self } = Provable.witness(LeafPair, () => this._findLeaf(key));
     let lowPath = this._proveInclusion(low, 'Invalid low node (root)');
-    low.key.assertLessThan(key, 'Invalid low node (key)');
-    key.assertLessThanOrEqual(low.nextKey, 'Invalid low node (next key)');
+    assertBetween(low.key, key, low.nextKey, 'Invalid low node (key)');
 
     // the key exists iff lowNode.nextKey == key
     let keyExists = low.nextKey.equals(key);
@@ -301,8 +314,7 @@ class IndexedMerkleMapBase {
     // prove whether the key exists or not, by showing a valid low node
     let { low, self } = Provable.witness(LeafPair, () => this._findLeaf(key));
     this._proveInclusion(low, 'Invalid low node (root)');
-    low.key.assertLessThan(key, 'Invalid low node (key)');
-    key.assertLessThanOrEqual(low.nextKey, 'Invalid low node (next key)');
+    assertBetween(low.key, key, low.nextKey, 'Invalid low node (key)');
 
     // the key exists iff lowNode.nextKey == key
     let keyExists = low.nextKey.equals(key);
@@ -337,8 +349,9 @@ class IndexedMerkleMapBase {
     // prove that the key does not exist yet, by showing a valid low node
     let low = Provable.witness(Leaf, () => this._findLeaf(key).low);
     this._proveInclusion(low, 'Invalid low node (root)');
-    low.key.assertLessThan(key, 'Invalid low node (key)');
-    key.assertLessThan(
+    assertStrictlyBetween(
+      low.key,
+      key,
       low.nextKey,
       message ?? 'Key already exists in the tree'
     );
@@ -353,8 +366,7 @@ class IndexedMerkleMapBase {
     // prove that the key does not exist yet, by showing a valid low node
     let low = Provable.witness(Leaf, () => this._findLeaf(key).low);
     this._proveInclusion(low, 'Invalid low node (root)');
-    low.key.assertLessThan(key, 'Invalid low node (key)');
-    key.assertLessThanOrEqual(low.nextKey, 'Invalid low node (next key)');
+    assertBetween(low.key, key, low.nextKey, 'Invalid low node (key)');
 
     return low.nextKey.equals(key);
   }
@@ -494,7 +506,7 @@ class IndexedMerkleMapBase {
     // and reject it using comparison constraints
     if (key === 0n)
       return {
-        low: Leaf.toValue(Leaf.empty()),
+        low: Leaf.fromStored(leaves[leaves.length - 1], leaves.length - 1),
         self: Leaf.fromStored(leaves[0], 0),
       };
 
@@ -528,7 +540,9 @@ class IndexedMerkleMapBase {
       let iSorted = leaf.sortedIndex.get();
 
       if (Bool(leafExists).toBoolean()) {
-        sortedLeaves[iSorted] = leafValue;
+        // for key=0, the sorted index overflows the length because we compute it as low.sortedIndex + 1
+        // in that case, it should wrap back to 0
+        sortedLeaves[iSorted % sortedLeaves.length] = leafValue;
       } else {
         sortedLeaves.splice(iSorted, 0, leafValue);
       }
@@ -701,3 +715,57 @@ function bisectUnique(
 
   return { lowIndex: iLow, foundValue: getValue(iLow) === target };
 }
+
+// custom comparison methods where 0 can act as the min and max value simultaneously
+
+/**
+ * Assert that `x in (low, high)`, i.e. low < x < high, with the following exceptions:
+ *
+ * - high=0 is treated as the maximum value, so `x in (low, 0)` always succeeds if only low < x; except for x = 0.
+ * - x=0 is also treated as the maximum value, so `0 in (low, high)` always fails, because x >= high.
+ */
+function assertStrictlyBetween(
+  low: Field,
+  x: Field,
+  high: Field,
+  message?: string
+) {
+  // exclude x=0
+  x.assertNotEquals(0n, message ?? '0 is not in any strict range');
+
+  // normal assertion for low < x
+  low.assertLessThan(x, message);
+
+  // for x < high, use a safe comparison that also works if high=0
+  let highIsZero = high.equals(0n);
+  let xSafe = Provable.witness(Field, () => (highIsZero.toBoolean() ? 0n : x));
+  let highSafe = Provable.witness(Field, () =>
+    highIsZero.toBoolean() ? 1n : high
+  );
+  xSafe.assertLessThan(highSafe, message);
+  assert(xSafe.equals(x).or(highIsZero), message);
+  assert(highSafe.equals(high).or(highIsZero), message);
+}
+
+/**
+ * Assert that `x in (low, high]`, i.e. low < x <= high, with the following exceptions:
+ *
+ * - high=0 is treated as the maximum value, so `x in (low, 0]` always succeeds if only low < x.
+ * - x=0 is also treated as the maximum value, so `0 in (low, high]` fails except if high=0.
+ * - note: `0 in (n, 0]` succeeds for any n!
+ */
+function assertBetween(low: Field, x: Field, high: Field, message?: string) {
+  // for low < x, we need to handle the x=0 case separately
+  let xIsZero = x.equals(0n);
+  let lowSafe = Provable.witness(Field, () => (xIsZero.toBoolean() ? 0n : low));
+  let xSafe1 = Provable.witness(Field, () => (xIsZero.toBoolean() ? 1n : x));
+  lowSafe.assertLessThan(xSafe1, message);
+  assert(lowSafe.equals(low).or(xIsZero), message);
+  assert(xSafe1.equals(x).or(xIsZero), message);
+
+  // for x <= high, we need to handle the high=0 case separately
+  let highIsZero = high.equals(0n);
+  let xSafe0 = Provable.witness(Field, () => (highIsZero.toBoolean() ? 0n : x));
+  xSafe0.assertLessThanOrEqual(high, message);
+  assert(xSafe0.equals(x).or(highIsZero), message);
+}

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

@@ -102,6 +102,7 @@ console.log(
 // some manual tests for IndexedMerkleMap
 {
   let map = new (IndexedMerkleMap(3))();
+  const minus1 = Field.ORDER - 1n;
 
   // there's 1 element in the map at the beginning
   // check initial root against `MerkleTree` implementation
@@ -110,24 +111,22 @@ console.log(
   initialTree.setLeaf(0n, Leaf.hashNode(IndexedMerkleMap(3)._firstLeaf));
   expect(map.root).toEqual(initialTree.getRoot());
 
-  map.insert(2n, 14n);
-  map.insert(1n, 13n);
+  // the initial value at key 0 is 0
+  expect(map.getOption(0n).assertSome().toBigInt()).toEqual(0n);
 
-  // -1 (the max value) can't be inserted, because there's always a value pointing to it,
-  // and yet it's not included as a leaf
-  expect(() => map.insert(-1n, 11n)).toThrow('Key already exists');
-  expect(() => map.set(-1n, 12n)).toThrow('Invalid leaf');
+  map.insert(-1n, 14n); // -1 is the largest possible value
+  map.insert(1n, 13n);
 
   expect(map.getOption(1n).assertSome().toBigInt()).toEqual(13n);
-  expect(map.getOption(2n).assertSome().toBigInt()).toEqual(14n);
+  expect(map.getOption(-1n).assertSome().toBigInt()).toEqual(14n);
   expect(map.getOption(3n).isSome.toBoolean()).toEqual(false);
 
-  map.update(2n, 15n);
+  map.update(-1n, 15n);
   map.update(0n, 12n);
-  expect(map.getOption(2n).assertSome().toBigInt()).toEqual(15n);
+  expect(map.getOption(-1n).assertSome().toBigInt()).toEqual(15n);
 
-  // TODO get() doesn't work on 0n because the low node checks fail
-  // expect(map.get(0n).assertSome().toBigInt()).toEqual(12n);
+  expect(map.get(0n).toBigInt()).toEqual(12n);
+  expect(map.getOption(0n).assertSome().toBigInt()).toEqual(12n);
 
   // can't insert the same key twice
   expect(() => map.insert(1n, 17n)).toThrow('Key already exists');
@@ -137,6 +136,7 @@ console.log(
 
   map.set(4n, 16n);
   map.set(1n, 17n);
+  map.set(0n, 12n);
   expect(map.get(4n).toBigInt()).toEqual(16n);
   expect(map.getOption(1n).assertSome().toBigInt()).toEqual(17n);
   expect(map.getOption(5n).isSome.toBoolean()).toEqual(false);
@@ -149,7 +149,7 @@ console.log(
   expect(map.length.toBigInt()).toEqual(4n);
 
   // check that internal nodes exactly match `MerkleTree` implementation
-  let keys = [0n, 2n, 1n, 4n]; // insertion order
+  let keys = [0n, minus1, 1n, 4n]; // insertion order
   let leafNodes = keys.map((key) => Leaf.hashNode(map._findLeaf(key).self));
   let tree = new MerkleTree(3);
   tree.fill(leafNodes);
@@ -167,16 +167,16 @@ console.log(
   let sorted = [
     { key: 0n, value: 12n, index: 0 },
     { key: 1n, value: 17n, index: 2 },
-    { key: 2n, value: 15n, index: 1 },
     { key: 4n, value: 16n, index: 3 },
+    { key: minus1, value: 15n, index: 1 },
   ];
   let sortedLeaves = map.data.get().sortedLeaves;
 
   for (let i = 0; i < 4; i++) {
     expect(sortedLeaves[i]).toEqual({
       key: sorted[i].key,
       value: sorted[i].value,
-      nextKey: sorted[i + 1]?.key ?? Field.ORDER - 1n,
+      nextKey: sorted[i + 1]?.key ?? 0n,
       index: sorted[i].index,
     });
   }