feat: add merkle trees

using pedersen hash

__tests__/utils/merkle.test.ts

@@ -0,0 +1,145 @@
+import { MerkleTree, proofMerklePath } from '../../src/utils/merkle';
+
+describe('MerkleTree class', () => {
+  describe('generate roots', () => {
+    test('should generate valid root for 1 elemements', async () => {
+      const leaves = ['0x1'];
+      const tree = new MerkleTree(leaves);
+
+      const manualMerkle = leaves[0];
+
+      expect(tree.root).toBe(manualMerkle);
+    });
+    test('should generate valid root for 2 elemements', async () => {
+      const leaves = ['0x1', '0x2'];
+      const tree = new MerkleTree(leaves);
+
+      const manualMerkle = MerkleTree.hash(leaves[0], leaves[1]);
+
+      expect(tree.root).toBe(manualMerkle);
+    });
+    test('should generate valid root for 4 elemements', async () => {
+      const leaves = ['0x1', '0x2', '0x3', '0x4'];
+      const tree = new MerkleTree(leaves);
+
+      const manualMerkle = MerkleTree.hash(
+        MerkleTree.hash(leaves[0], leaves[1]),
+        MerkleTree.hash(leaves[2], leaves[3])
+      );
+
+      expect(tree.root).toBe(manualMerkle);
+    });
+    test('should generate valid root for 6 elemements', async () => {
+      const leaves = ['0x1', '0x2', '0x3', '0x4', '0x5', '0x6'];
+      const tree = new MerkleTree(leaves);
+
+      const manualMerkle = MerkleTree.hash(
+        MerkleTree.hash(
+          MerkleTree.hash(leaves[0], leaves[1]),
+          MerkleTree.hash(leaves[2], leaves[3])
+        ),
+        MerkleTree.hash(leaves[4], leaves[5])
+      );
+
+      expect(tree.root).toBe(manualMerkle);
+    });
+    test('should generate valid root for 7 elemements', async () => {
+      const leaves = ['0x1', '0x2', '0x3', '0x4', '0x5', '0x6', '0x7'];
+      const tree = new MerkleTree(leaves);
+
+      const manualMerkle = MerkleTree.hash(
+        MerkleTree.hash(
+          MerkleTree.hash(leaves[0], leaves[1]),
+          MerkleTree.hash(leaves[2], leaves[3])
+        ),
+        MerkleTree.hash(MerkleTree.hash(leaves[4], leaves[5]), leaves[6])
+      );
+
+      expect(tree.root).toBe(manualMerkle);
+    });
+  });
+  describe('generate proofs', () => {
+    let tree: MerkleTree;
+    beforeAll(() => {
+      const leaves = ['0x1', '0x2', '0x3', '0x4', '0x5', '0x6', '0x7'];
+      tree = new MerkleTree(leaves);
+    });
+    test('should return proof path for valid child', async () => {
+      const proof = tree.getProof('0x3');
+
+      const manualProof = [
+        '0x4',
+        MerkleTree.hash('0x1', '0x2'),
+        MerkleTree.hash(MerkleTree.hash('0x5', '0x6'), '0x7'),
+      ];
+
+      expect(proof).toEqual(manualProof);
+    });
+    test('should return proof path for valid child', async () => {
+      const proof = tree.getProof('0x7');
+
+      const manualProof = [
+        MerkleTree.hash('0x5', '0x6'),
+        MerkleTree.hash(MerkleTree.hash('0x1', '0x2'), MerkleTree.hash('0x3', '0x4')),
+      ];
+
+      expect(proof).toEqual(manualProof);
+    });
+    test('should throw for invalid child', () => {
+      expect(() => tree.getProof('0x8')).toThrow('leaf not found');
+    });
+  });
+  describe('verify proofs', () => {
+    let tree: MerkleTree;
+    beforeAll(() => {
+      const leaves = ['0x1', '0x2', '0x3', '0x4', '0x5', '0x6', '0x7'];
+      tree = new MerkleTree(leaves);
+    });
+
+    test('should return true for valid manual proof', async () => {
+      const manualProof = [
+        MerkleTree.hash('0x5', '0x6'),
+        MerkleTree.hash(MerkleTree.hash('0x1', '0x2'), MerkleTree.hash('0x3', '0x4')),
+      ];
+      const leaf = '0x7';
+      const { root } = tree;
+
+      expect(proofMerklePath(root, leaf, manualProof)).toBe(true);
+    });
+    test('should return true for valid proof', async () => {
+      const proof = tree.getProof('0x3');
+      const leaf = '0x3';
+      const { root } = tree;
+
+      expect(proofMerklePath(root, leaf, proof)).toBe(true);
+    });
+    test('should return false for invalid proof (root)', async () => {
+      const proof = tree.getProof('0x3');
+      const leaf = '0x3';
+      const root = '0x4';
+
+      expect(proofMerklePath(root, leaf, proof)).toBe(false);
+    });
+    test('should return false for invalid proof (proof[0])', async () => {
+      const proof = tree.getProof('0x3');
+      const leaf = '0x3';
+      const { root } = tree;
+      proof[0] = '0x7';
+      expect(proofMerklePath(root, leaf, proof)).toBe(false);
+    });
+    test('should return false for invalid proof (proof[1])', async () => {
+      const proof = tree.getProof('0x3');
+      const leaf = '0x3';
+      const { root } = tree;
+      proof[1] = '0x4';
+      expect(proofMerklePath(root, leaf, proof)).toBe(false);
+    });
+    test('should return false for invalid proof (proof[2])', async () => {
+      const proof = tree.getProof('0x3');
+      const leaf = '0x3';
+      const { root } = tree;
+      proof[2] = '0x4';
+      expect(proofMerklePath(root, leaf, proof)).toBe(false);
+    });
+  });
+});

src/index.ts

@@ -17,6 +17,7 @@ export * as json from './utils/json';
 export * as number from './utils/number';
 export * as transaction from './utils/transaction';
 export * as stark from './utils/stark';
+export * as merkle from './utils/merkle';
 export * as ec from './utils/ellipticCurve';
 export * as uint256 from './utils/uint256';
 export * as shortString from './utils/shortString';

src/utils/merkle.ts

@@ -0,0 +1,70 @@
+import { pedersen } from './hash';
+
+export class MerkleTree {
+  public leaves: string[];
+
+  public branches: string[][] = [];
+
+  public root: string;
+
+  constructor(leafHashes: string[]) {
+    this.leaves = leafHashes;
+    this.root = this.build(leafHashes);
+  }
+
+  build(leaves: string[]): string {
+    if (leaves.length === 1) {
+      return leaves[0];
+    }
+    if (leaves.length !== this.leaves.length) {
+      this.branches.push(leaves);
+    }
+    const newLeaves = [];
+    for (let i = 0; i < leaves.length; i += 2) {
+      if (i + 1 === leaves.length) {
+        newLeaves.push(leaves[i]);
+      } else {
+        newLeaves.push(MerkleTree.hash(leaves[i], leaves[i + 1]));
+      }
+    }
+    return this.build(newLeaves);
+  }
+
+  static hash(a: string, b: string) {
+    const [aSorted, bSorted] = [a, b].sort();
+    return pedersen([aSorted, bSorted]);
+  }
+
+  getProof(leaf: string, branch = this.leaves, hashPath: string[] = []): string[] {
+    if (branch.length === 1) {
+      return hashPath;
+    }
+    const index = branch.indexOf(leaf);
+    if (index === -1) {
+      throw new Error('leaf not found');
+    }
+    const isLeft = index % 2 === 0;
+    const neededBranch = (isLeft ? branch[index + 1] : branch[index - 1]) ?? branch[index];
+    const newHashPath = [...hashPath, neededBranch];
+    const currentBranchLevelIndex =
+      this.leaves.length === branch.length
+        ? -1
+        : this.branches.findIndex((b) => b.length === branch.length);
+    const nextBranch = this.branches[currentBranchLevelIndex + 1] ?? [this.root];
+    return this.getProof(
+      neededBranch === leaf
+        ? leaf
+        : MerkleTree.hash(isLeft ? leaf : neededBranch, isLeft ? neededBranch : leaf),
+      nextBranch,
+      neededBranch === leaf ? hashPath : newHashPath
+    );
+  }
+}
+
+export function proofMerklePath(root: string, leaf: string, path: string[]): boolean {
+  if (path.length === 0) {
+    return root === leaf;
+  }
+  const [next, ...rest] = path;
+  return proofMerklePath(root, MerkleTree.hash(leaf, next), rest);
+}