Merge pull request #546 from o1-labs/merkle-map

implement merkle map

CHANGELOG.md

@@ -15,7 +15,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
     _Security_ in case of vulnerabilities.
  -->
 
-## [Unreleased](https://github.com/o1-labs/snarkyjs/compare/9a94231c...HEAD)
+## [Unreleased](https://github.com/o1-labs/snarkyjs/compare/705f58d3...HEAD)
+
+### Added
+
+- `MerkleMap` and `MerkleMapWitness` https://github.com/o1-labs/snarkyjs/pull/546
+
+## [0.7.1](https://github.com/o1-labs/snarkyjs/compare/f0837188...705f58d3)
+
+### Fixed
+
+- Testnet-incompatible signatures in v0.7.0 https://github.com/o1-labs/snarkyjs/pull/565
 
 ## [0.7.0](https://github.com/o1-labs/snarkyjs/compare/f0837188...9a94231c)
 

src/index.ts

@@ -62,6 +62,7 @@ export * as Encryption from './lib/encryption.js';
 export * as Encoding from './lib/encoding.js';
 export { Character, CircuitString } from './lib/string.js';
 export { MerkleTree, MerkleWitness } from './lib/merkle_tree.js';
+export { MerkleMap, MerkleMapWitness } from './lib/merkle_map.js';
 
 // experimental APIs
 import { ZkProgram } from './lib/proof_system.js';

src/lib/merkle_map.test.ts

@@ -0,0 +1,45 @@
+import { isReady, shutdown, Field, MerkleMap } from 'snarkyjs';
+
+describe('Merkle Map', () => {
+  beforeAll(async () => {
+    await isReady;
+  });
+  afterAll(async () => {
+    setTimeout(shutdown, 0);
+  });
+
+  it('set and get a value from a key', () => {
+    const map = new MerkleMap();
+
+    const key = Field.random();
+    const value = Field.random();
+
+    map.set(key, value);
+
+    expect(map.get(key).equals(value).toBoolean());
+  });
+
+  it('check merkle map witness computes the correct root and key', () => {
+    const map = new MerkleMap();
+
+    const key = Field.random();
+    const value = Field.random();
+
+    map.set(key, value);
+
+    const witness = map.getWitness(key);
+
+    const emptyMap = new MerkleMap();
+
+    const [emptyLeafWitnessRoot, witnessKey] = witness.computeRootAndKey(
+      Field(0)
+    );
+    const [witnessRoot, _] = witness.computeRootAndKey(value);
+
+    expect(
+      emptyLeafWitnessRoot.equals(emptyMap.getRoot()).toBoolean() &&
+        witnessKey.equals(key).toBoolean() &&
+        witnessRoot.equals(map.getRoot()).toBoolean()
+    );
+  });
+});

src/lib/merkle_map.ts

@@ -0,0 +1,151 @@
+import { arrayProp, CircuitValue, Circuit } from './circuit_value.js';
+import { Field, Bool } from './core.js';
+import { Poseidon } from './hash.js';
+import { MerkleTree, MerkleWitness } from './merkle_tree.js';
+
+const bits = 255;
+const printDebugs = false;
+
+export class MerkleMap {
+  tree: InstanceType<typeof 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);
+  }
+
+  // ------------------------------------------------
+
+  _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 n = 0n;
+    for (let i = 0; i < keyBits.length; i++) {
+      const b = keyBits[i] ? 1 : 0;
+      n += 2n ** BigInt(i) * BigInt(b);
+    }
+
+    return n;
+  }
+
+  // ------------------------------------------------
+
+  /**
+   * Sets a key of the merkle map to a given value.
+   * @param key The key to set in the map.
+   * @param key The value to set.
+   */
+  set(key: Field, value: Field) {
+    const index = this._keyToIndex(key);
+    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.
+   * @returns The value stored at the key.
+   */
+  get(key: Field) {
+    const index = this._keyToIndex(key);
+    return this.tree.getNode(0, index);
+  }
+
+  // ------------------------------------------------
+
+  /**
+   * Returns the root of the Merkle Map.
+   * @returns The root of the Merkle Map.
+   */
+  getRoot() {
+    return this.tree.getRoot();
+  }
+
+  /**
+   * 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.
+   */
+  getWitness(key: Field) {
+    const index = this._keyToIndex(key);
+    class MyMerkleWitness extends MerkleWitness(bits + 1) {}
+    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[];
+
+  constructor(isLefts: Bool[], siblings: Field[]) {
+    super();
+    this.isLefts = isLefts;
+    this.siblings = siblings;
+  }
+
+  /**
+   * 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) {
+    let hash = value;
+
+    const isLeft = this.isLefts;
+    const siblings = this.siblings;
+
+    let key = Field(0);
+
+    for (let i = 0; i < bits; i++) {
+      const left = Circuit.if(isLeft[i], hash, siblings[i]);
+      const right = Circuit.if(isLeft[i], siblings[i], hash);
+      hash = Poseidon.hash([left, right]);
+
+      const bit = Circuit.if(isLeft[i], Field(0), Field(1));
+
+      key = key.mul(2).add(bit);
+    }
+
+    return [hash, key];
+  }
+}