Add Experimental Merkle Tree Implementation

src/index.ts

@@ -62,6 +62,7 @@ import {
 } from './lib/circuit_value';
 import { jsLayout, asFieldsAndAux } from './snarky/types';
 import { packToFields } from './lib/hash';
+import { MerkleTree, MerkleWitness } from './lib/merkle_tree';
 export { Experimental };
 
 /**
@@ -76,6 +77,8 @@ const Experimental = {
   jsLayout,
   asFieldsAndAux,
   packToFields,
+  MerkleTree,
+  MerkleWitness,
 };
 namespace Experimental {
   export type AsFieldsAndAux<T, TJson> = AsFieldsAndAux_<T, TJson>;

src/lib/merkle_tree.test.ts

@@ -0,0 +1,61 @@
+import {
+  isReady,
+  shutdown,
+  Poseidon,
+  Field,
+  Experimental,
+} from '../../dist/server';
+
+describe('Merkle Tree', () => {
+  beforeAll(async () => {
+    await isReady;
+  });
+  afterAll(async () => {
+    setTimeout(shutdown, 0);
+  });
+
+  it('root of empty tree of size 1', () => {
+    const tree = new Experimental.MerkleTree(1);
+    expect(tree.getRoot().toString()).toEqual(Field(0).toString());
+  });
+
+  it('root is correct', () => {
+    const tree = new Experimental.MerkleTree(2);
+    tree.setLeaf(0n, Field(1));
+    tree.setLeaf(1n, Field(2));
+    expect(tree.getRoot().toString()).toEqual(
+      Poseidon.hash([Field(1), Field(2)]).toString()
+    );
+  });
+
+  it('builds correct tree', () => {
+    const tree = new Experimental.MerkleTree(4);
+    tree.setLeaf(0n, Field(1));
+    tree.setLeaf(1n, Field(2));
+    tree.setLeaf(2n, Field(3));
+    expect(tree.validate(0n)).toBe(true);
+    expect(tree.validate(1n)).toBe(true);
+    expect(tree.validate(2n)).toBe(true);
+    expect(tree.validate(3n)).toBe(true);
+  });
+
+  it('tree of height 128', () => {
+    const tree = new Experimental.MerkleTree(128);
+
+    const index = 2n ** 64n;
+    expect(tree.validate(index)).toBe(true);
+
+    tree.setLeaf(index, Field(1));
+    expect(tree.validate(index)).toBe(true);
+  });
+
+  it('tree of height 256', () => {
+    const tree = new Experimental.MerkleTree(256);
+
+    const index = 2n ** 128n;
+    expect(tree.validate(index)).toBe(true);
+
+    tree.setLeaf(index, Field(1));
+    expect(tree.validate(index)).toBe(true);
+  });
+});

src/lib/merkle_tree.ts

@@ -0,0 +1,115 @@
+import { Circuit, CircuitValue, arrayProp } from './circuit_value';
+import { Poseidon } from './hash';
+import { Bool, Field } from './core';
+
+export type Witness = {
+  isLeft: boolean;
+  sibling: Field;
+}[];
+
+/**
+ * Levels are indexed from leafs (level 0) to root (level N - 1).
+ */
+export class MerkleTree {
+  private nodes: Record<number, Record<string, Field>> = {};
+  private zeroes: Field[];
+
+  constructor(public readonly height: number) {
+    this.zeroes = [Field(0)];
+    for (let i = 1; i < height; i++) {
+      this.zeroes.push(Poseidon.hash([this.zeroes[i - 1], this.zeroes[i - 1]]));
+    }
+  }
+
+  getNode(level: number, index: bigint): Field {
+    return this.nodes[level]?.[index.toString()] ?? this.zeroes[level];
+  }
+
+  getRoot(): Field {
+    return this.getNode(this.height - 1, 0n);
+  }
+
+  private setNode(level: number, index: bigint, value: Field) {
+    (this.nodes[level] ??= {})[index.toString()] = value;
+  }
+
+  setLeaf(index: bigint, leaf: Field) {
+    this.setNode(0, index, leaf);
+    let currIndex = index;
+    for (let level = 1; level < this.height; level++) {
+      currIndex = currIndex / 2n;
+
+      const left = this.getNode(level - 1, currIndex * 2n);
+      const right = this.getNode(level - 1, currIndex * 2n + 1n);
+
+      this.setNode(level, currIndex, Poseidon.hash([left, right]));
+    }
+  }
+
+  getWitness(index: bigint): Witness {
+    const witness = [];
+    for (let level = 0; level < this.height - 1; level++) {
+      const isLeft = index % 2n === 0n;
+      witness.push({
+        isLeft,
+        sibling: this.getNode(level, isLeft ? index + 1n : index - 1n),
+      });
+      index = index / 2n;
+    }
+    return witness;
+  }
+
+  validate(index: bigint): boolean {
+    const path = this.getWitness(index);
+    let hash = this.getNode(0, index);
+    for (const node of path) {
+      hash = Poseidon.hash(
+        node.isLeft ? [hash, node.sibling] : [node.sibling, hash]
+      );
+    }
+
+    return hash.toString() === this.getRoot().toString();
+  }
+
+  fill(leaves: Field[]) {
+    leaves.forEach((value, index) => {
+      this.setLeaf(BigInt(index), value);
+    });
+  }
+}
+
+export class MerkleWitness extends CircuitValue {
+  @arrayProp(Field, 255) path: Field[];
+  @arrayProp(Bool, 255) isLeft: Bool[];
+
+  constructor(witness: Witness) {
+    super(witness);
+
+    this.path = witness.map((item) => item.sibling);
+    this.isLeft = witness.map((item) => Bool(item.isLeft));
+  }
+
+  calculateRoot(leaf: Field): Field {
+    let hash = leaf;
+
+    for (let i = 1; i < 256; ++i) {
+      const left = Circuit.if(this.isLeft[i - 1], hash, this.path[i - 1]);
+      const right = Circuit.if(this.isLeft[i - 1], this.path[i - 1], hash);
+      hash = Poseidon.hash([left, right]);
+    }
+
+    return hash;
+  }
+
+  calculateIndex(): Field {
+    let powerOfTwo = Field(1);
+    let index = Field(0);
+
+    for (let i = 1; i < 256; ++i) {
+      index = Circuit.if(this.isLeft[i - 1], index, index.add(powerOfTwo));
+      powerOfTwo = powerOfTwo.mul(2);
+    }
+
+    return index;
+  }
+}