Add SimpleMerkleTree (#36)

Co-authored-by: Hadrien Croubois <hadrien.croubois@gmail.com>

Author: ernestognw

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 1.0.6
 
 - Added an option to disable leaf sorting.
+- Added `SimpleMerkleTree` class that supports `bytes32` leaves with no extra hashing.
 
 ## 1.0.5
 

package-lock.json

@@ -24,7 +24,9 @@
   "license": "MIT",
   "dependencies": {
     "@ethersproject/abi": "^5.7.0",
-    "ethereum-cryptography": "^1.1.2"
+    "@ethersproject/bytes": "^5.7.0",
+    "@ethersproject/constants": "^5.7.0",
+    "@ethersproject/keccak256": "^5.7.0"
   },
   "devDependencies": {
     "@types/mocha": "^10.0.0",

package.json

@@ -1,19 +1,12 @@
-import { bytesToHex } from 'ethereum-cryptography/utils';
+import type { BytesLike } from '@ethersproject/bytes';
+type HexString = string;
 
-export type Bytes = Uint8Array;
+import { arrayify as toBytes, hexlify as toHex, concat } from '@ethersproject/bytes';
 
-export function compareBytes(a: Bytes, b: Bytes): number {
-  const n = Math.min(a.length, b.length);
-
-  for (let i = 0; i < n; i++) {
-    if (a[i] !== b[i]) {
-      return a[i]! - b[i]!;
-    }
-  }
-
-  return a.length - b.length;
+function compare(a: BytesLike, b: BytesLike): number {
+  const diff = BigInt(toHex(a)) - BigInt(toHex(b));
+  return diff > 0 ? 1 : diff < 0 ? -1 : 0;
 }
 
-export function hex(b: Bytes): string {
-  return '0x' + bytesToHex(b);
-}
+export type { HexString, BytesLike };
+export { toBytes, toHex, concat, compare };

src/bytes.ts

@@ -1,6 +1,7 @@
 import fc from 'fast-check';
 import assert from 'assert/strict';
-import { equalsBytes } from 'ethereum-cryptography/utils';
+import { HashZero as zero } from '@ethersproject/constants';
+import { keccak256 } from '@ethersproject/keccak256';
 import {
   makeMerkleTree,
   getProof,
@@ -10,12 +11,9 @@ import {
   isValidMerkleTree,
   renderMerkleTree,
 } from './core';
-import { compareBytes, hex } from './bytes';
-import { keccak256 } from 'ethereum-cryptography/keccak';
+import { toHex, compare } from './bytes';
 
-const zero = new Uint8Array(32);
-
-const leaf = fc.uint8Array({ minLength: 32, maxLength: 32 }).map(x => PrettyBytes.from(x));
+const leaf = fc.uint8Array({ minLength: 32, maxLength: 32 }).map(toHex);
 const leaves = fc.array(leaf, { minLength: 1 });
 const leavesAndIndex = leaves.chain(xs => fc.tuple(fc.constant(xs), fc.nat({ max: xs.length - 1 })));
 const leavesAndIndices = leaves.chain(xs => fc.tuple(fc.constant(xs), fc.uniqueArray(fc.nat({ max: xs.length - 1 }))));
@@ -33,7 +31,7 @@ describe('core properties', () => {
         const proof = getProof(tree, treeIndex);
         const leaf = leaves[leafIndex]!;
         const impliedRoot = processProof(leaf, proof);
-        return equalsBytes(root, impliedRoot);
+        return root === impliedRoot;
       }),
     );
   });
@@ -49,46 +47,40 @@ describe('core properties', () => {
         if (leafIndices.length !== proof.leaves.length) return false;
         if (leafIndices.some(i => !proof.leaves.includes(leaves[i]!))) return false;
         const impliedRoot = processMultiProof(proof);
-        return equalsBytes(root, impliedRoot);
+        return root === impliedRoot;
       }),
     );
   });
 });
 
 describe('core error conditions', () => {
   it('zero leaves', () => {
-    assert.throws(() => makeMerkleTree([]), /^Error: Expected non-zero number of leaves$/);
+    assert.throws(() => makeMerkleTree([]), /^InvalidArgumentError: Expected non-zero number of leaves$/);
   });
 
   it('multiproof duplicate index', () => {
     const tree = makeMerkleTree(new Array(2).fill(zero));
-    assert.throws(() => getMultiProof(tree, [1, 1]), /^Error: Cannot prove duplicated index$/);
+    assert.throws(() => getMultiProof(tree, [1, 1]), /^InvalidArgumentError: Cannot prove duplicated index$/);
   });
 
   it('tree validity', () => {
     assert(!isValidMerkleTree([]), 'empty tree');
     assert(!isValidMerkleTree([zero, zero]), 'even number of nodes');
     assert(!isValidMerkleTree([zero, zero, zero]), 'inner node not hash of children');
 
-    assert.throws(() => renderMerkleTree([]), /^Error: Expected non-zero number of nodes$/);
+    assert.throws(() => renderMerkleTree([]), /^InvalidArgumentError: Expected non-zero number of nodes$/);
   });
 
   it('multiproof invariants', () => {
     const leaf = keccak256(Uint8Array.of(42));
     const tree = makeMerkleTree([leaf, zero]);
 
     const badMultiProof = {
-      leaves: [128, 129].map(n => keccak256(Uint8Array.of(n))).sort(compareBytes),
+      leaves: [128, 129].map(n => keccak256(Uint8Array.of(n))).sort(compare),
       proof: [leaf, leaf],
       proofFlags: [true, true, false],
     };
 
-    assert.throws(() => processMultiProof(badMultiProof), /^Error: Broken invariant$/);
+    assert.throws(() => processMultiProof(badMultiProof), /^InvariantError$/);
   });
 });
-
-class PrettyBytes extends Uint8Array {
-  [fc.toStringMethod]() {
-    return hex(this);
-  }
-}

src/core.test.ts

@@ -1,9 +1,8 @@
-import { keccak256 } from 'ethereum-cryptography/keccak';
-import { concatBytes, bytesToHex, equalsBytes } from 'ethereum-cryptography/utils';
-import { Bytes, compareBytes } from './bytes';
-import { throwError } from './utils/throw-error';
+import { keccak256 } from '@ethersproject/keccak256';
+import { BytesLike, HexString, toHex, toBytes, concat, compare } from './bytes';
+import { invariant, throwError, validateArgument } from './utils/errors';
 
-const hashPair = (a: Bytes, b: Bytes) => keccak256(concatBytes(...[a, b].sort(compareBytes)));
+const hashPair = (a: BytesLike, b: BytesLike): HexString => keccak256(concat([a, b].sort(compare)));
 
 const leftChildIndex = (i: number) => 2 * i + 1;
 const rightChildIndex = (i: number) => 2 * i + 2;
@@ -13,26 +12,21 @@ const siblingIndex = (i: number) => (i > 0 ? i - (-1) ** (i % 2) : throwError('R
 const isTreeNode = (tree: unknown[], i: number) => i >= 0 && i < tree.length;
 const isInternalNode = (tree: unknown[], i: number) => isTreeNode(tree, leftChildIndex(i));
 const isLeafNode = (tree: unknown[], i: number) => isTreeNode(tree, i) && !isInternalNode(tree, i);
-const isValidMerkleNode = (node: Bytes) => node instanceof Uint8Array && node.length === 32;
+const isValidMerkleNode = (node: BytesLike) => toBytes(node).length === 32;
 
-const checkTreeNode = (tree: unknown[], i: number) => void (isTreeNode(tree, i) || throwError('Index is not in tree'));
-const checkInternalNode = (tree: unknown[], i: number) =>
-  void (isInternalNode(tree, i) || throwError('Index is not an internal tree node'));
 const checkLeafNode = (tree: unknown[], i: number) => void (isLeafNode(tree, i) || throwError('Index is not a leaf'));
-const checkValidMerkleNode = (node: Bytes) =>
+const checkValidMerkleNode = (node: BytesLike) =>
   void (isValidMerkleNode(node) || throwError('Merkle tree nodes must be Uint8Array of length 32'));
 
-export function makeMerkleTree(leaves: Bytes[]): Bytes[] {
+export function makeMerkleTree(leaves: BytesLike[]): HexString[] {
   leaves.forEach(checkValidMerkleNode);
 
-  if (leaves.length === 0) {
-    throw new Error('Expected non-zero number of leaves');
-  }
+  validateArgument(leaves.length !== 0, 'Expected non-zero number of leaves');
 
-  const tree = new Array<Bytes>(2 * leaves.length - 1);
+  const tree = new Array<HexString>(2 * leaves.length - 1);
 
   for (const [i, leaf] of leaves.entries()) {
-    tree[tree.length - 1 - i] = leaf;
+    tree[tree.length - 1 - i] = toHex(leaf);
   }
   for (let i = tree.length - 1 - leaves.length; i >= 0; i--) {
     tree[i] = hashPair(tree[leftChildIndex(i)]!, tree[rightChildIndex(i)]!);
@@ -41,22 +35,22 @@ export function makeMerkleTree(leaves: Bytes[]): Bytes[] {
   return tree;
 }
 
-export function getProof(tree: Bytes[], index: number): Bytes[] {
+export function getProof(tree: BytesLike[], index: number): HexString[] {
   checkLeafNode(tree, index);
 
   const proof = [];
   while (index > 0) {
-    proof.push(tree[siblingIndex(index)]!);
+    proof.push(toHex(tree[siblingIndex(index)]!));
     index = parentIndex(index);
   }
   return proof;
 }
 
-export function processProof(leaf: Bytes, proof: Bytes[]): Bytes {
+export function processProof(leaf: BytesLike, proof: BytesLike[]): HexString {
   checkValidMerkleNode(leaf);
   proof.forEach(checkValidMerkleNode);
 
-  return proof.reduce(hashPair, leaf);
+  return toHex(proof.reduce(hashPair, leaf));
 }
 
 export interface MultiProof<T, L = T> {
@@ -65,13 +59,14 @@ export interface MultiProof<T, L = T> {
   proofFlags: boolean[];
 }
 
-export function getMultiProof(tree: Bytes[], indices: number[]): MultiProof<Bytes> {
+export function getMultiProof(tree: BytesLike[], indices: number[]): MultiProof<HexString> {
   indices.forEach(i => checkLeafNode(tree, i));
   indices.sort((a, b) => b - a);
 
-  if (indices.slice(1).some((i, p) => i === indices[p])) {
-    throw new Error('Cannot prove duplicated index');
-  }
+  validateArgument(
+    indices.slice(1).every((i, p) => i !== indices[p]),
+    'Cannot prove duplicated index',
+  );
 
   const stack = indices.concat(); // copy
   const proof = [];
@@ -87,54 +82,51 @@ export function getMultiProof(tree: Bytes[], indices: number[]): MultiProof<Byte
       stack.shift(); // consume from the stack
     } else {
       proofFlags.push(false);
-      proof.push(tree[s]!);
+      proof.push(toHex(tree[s]!));
     }
     stack.push(p);
   }
 
   if (indices.length === 0) {
-    proof.push(tree[0]!);
+    proof.push(toHex(tree[0]!));
   }
 
   return {
-    leaves: indices.map(i => tree[i]!),
+    leaves: indices.map(i => toHex(tree[i]!)),
     proof,
     proofFlags,
   };
 }
 
-export function processMultiProof(multiproof: MultiProof<Bytes>): Bytes {
+export function processMultiProof(multiproof: MultiProof<BytesLike>): HexString {
   multiproof.leaves.forEach(checkValidMerkleNode);
   multiproof.proof.forEach(checkValidMerkleNode);
 
-  if (multiproof.proof.length < multiproof.proofFlags.filter(b => !b).length) {
-    throw new Error('Invalid multiproof format');
-  }
-
-  if (multiproof.leaves.length + multiproof.proof.length !== multiproof.proofFlags.length + 1) {
-    throw new Error('Provided leaves and multiproof are not compatible');
-  }
+  validateArgument(
+    multiproof.proof.length >= multiproof.proofFlags.filter(b => !b).length,
+    'Invalid multiproof format',
+  );
+  validateArgument(
+    multiproof.leaves.length + multiproof.proof.length === multiproof.proofFlags.length + 1,
+    'Provided leaves and multiproof are not compatible',
+  );
 
   const stack = multiproof.leaves.concat(); // copy
   const proof = multiproof.proof.concat(); // copy
 
   for (const flag of multiproof.proofFlags) {
     const a = stack.shift();
     const b = flag ? stack.shift() : proof.shift();
-    if (a === undefined || b === undefined) {
-      throw new Error('Broken invariant');
-    }
+    invariant(a !== undefined && b !== undefined);
     stack.push(hashPair(a, b));
   }
 
-  if (stack.length + proof.length !== 1) {
-    throw new Error('Broken invariant');
-  }
+  invariant(stack.length + proof.length === 1);
 
-  return stack.pop() ?? proof.shift()!;
+  return toHex(stack.pop() ?? proof.shift()!);
 }
 
-export function isValidMerkleTree(tree: Bytes[]): boolean {
+export function isValidMerkleTree(tree: BytesLike[]): boolean {
   for (const [i, node] of tree.entries()) {
     if (!isValidMerkleNode(node)) {
       return false;
@@ -147,18 +139,16 @@ export function isValidMerkleTree(tree: Bytes[]): boolean {
       if (l < tree.length) {
         return false;
       }
-    } else if (!equalsBytes(node, hashPair(tree[l]!, tree[r]!))) {
+    } else if (compare(node, hashPair(tree[l]!, tree[r]!))) {
       return false;
     }
   }
 
   return tree.length > 0;
 }
 
-export function renderMerkleTree(tree: Bytes[]): string {
-  if (tree.length === 0) {
-    throw new Error('Expected non-zero number of nodes');
-  }
+export function renderMerkleTree(tree: BytesLike[]): HexString {
+  validateArgument(tree.length !== 0, 'Expected non-zero number of nodes');
 
   const stack: [number, number[]][] = [[0, []]];
 
@@ -178,7 +168,7 @@ export function renderMerkleTree(tree: Bytes[]): string {
           .join('') +
         i +
         ') ' +
-        bytesToHex(tree[i]!),
+        toHex(tree[i]!),
     );
 
     if (rightChildIndex(i) < tree.length) {

src/core.ts

@@ -0,0 +1,9 @@
+import assert from 'assert/strict';
+import { SimpleMerkleTree, StandardMerkleTree } from '.';
+
+describe('index properties', () => {
+  it('classes are exported', () => {
+    assert.notEqual(SimpleMerkleTree, undefined);
+    assert.notEqual(StandardMerkleTree, undefined);
+  });
+});

src/index.test.ts

@@ -1 +1,2 @@
+export { SimpleMerkleTree } from './simple';
 export { StandardMerkleTree } from './standard';