improve: Import MerkleTree from @uma/common instead of redefining (#210)

* improve: Import MerkleTree from @uma/common instead of redefining

* update uma/common

* Update MerkleTree.ts

* Fix tests

Author: nicholaspai

package.json

@@ -35,7 +35,7 @@
     "@defi-wonderland/smock": "^2.0.7",
     "@eth-optimism/contracts": "^0.5.11",
     "@openzeppelin/contracts": "^4.7.3",
-    "@uma/common": "^2.28.4",
+    "@uma/common": "^2.29.0",
     "@uma/contracts-node": "^0.3.18",
     "@uma/core": "^2.41.0",
     "arb-bridge-eth": "^0.7.4",

test/MerkleLib.utils.ts

@@ -125,5 +125,5 @@ export async function buildSlowRelayTree(relays: RelayData[]) {
   const hashFn = (input: RelayData) => {
     return keccak256(defaultAbiCoder.encode([paramType!], [input]));
   };
-  return new MerkleTree(relays, hashFn);
+  return new MerkleTree<RelayData>(relays, hashFn);
 }

test/merkle-distributor/MerkleDistributor.ts

@@ -1,6 +1,16 @@
 /* eslint-disable no-unused-expressions */
 
-import { ethers, getContractFactory, SignerWithAddress, Contract, toWei, toBN, expect } from "../utils";
+import {
+  ethers,
+  getContractFactory,
+  SignerWithAddress,
+  Contract,
+  toWei,
+  toBN,
+  expect,
+  keccak256,
+  defaultAbiCoder,
+} from "../utils";
 import { deployErc20 } from "../gas-analytics/utils";
 import { MAX_UINT_VAL, MerkleTree } from "@uma/common";
 
@@ -19,7 +29,9 @@ let contractCreator: SignerWithAddress;
 let otherAddress: SignerWithAddress;
 
 // Test variables
-let merkleTree: MerkleTree;
+let merkleTree: MerkleTree<Buffer>;
+const hashFn = (input: Buffer) => input.toString("hex");
+
 let windowIndex: number;
 const sampleIpfsHash = "";
 
@@ -58,7 +70,8 @@ describe("AcrossMerkleDistributor", () => {
         amount: totalRewardAmount,
         accountIndex: 0,
       });
-      merkleTree = new MerkleTree([leaf]);
+
+      merkleTree = new MerkleTree<Buffer>([leaf], hashFn);
       // Expect this merkle root to be at the first index.
       windowIndex = 0;
       // Seed the merkleDistributor with the root of the tree and additional information.
@@ -118,7 +131,7 @@ describe("AcrossMerkleDistributor", () => {
         amount: totalRewardAmount,
         accountIndex: 1,
       });
-      merkleTree = new MerkleTree([leaf1, leaf2]);
+      merkleTree = new MerkleTree<Buffer>([leaf1, leaf2], hashFn);
       // Expect this merkle root to be at the first index.
       windowIndex = 0;
       // Seed the merkleDistributor with the root of the tree and additional information.
@@ -171,7 +184,7 @@ describe("AcrossMerkleDistributor", () => {
         amount: totalRewardAmount,
         accountIndex: 0,
       });
-      merkleTree = new MerkleTree([leaf]);
+      merkleTree = new MerkleTree<Buffer>([leaf], hashFn);
       // Expect this merkle root to be at the first index.
       windowIndex = 0;
       // Seed the merkleDistributor with the root of the tree and additional information.

utils/MerkleTree.ts

@@ -1,145 +1 @@
-// This script provides some useful methods for building MerkleTrees. It is essentially the uniswap implementation
-// https://github.com/Uniswap/merkle-distributor/blob/master/src/merkle-tree.ts with some added convenience methods
-// to take the leaves and conversion functions, so the user never has to work with buffers.
-import { bufferToHex, keccak256 } from "ethereumjs-util";
-
-export const EMPTY_MERKLE_ROOT = "0x0000000000000000000000000000000000000000000000000000000000000000";
-export class MerkleTree<T> {
-  private readonly elements: Buffer[];
-  private readonly bufferElementPositionIndex: { [hexElement: string]: number };
-  private readonly layers: Buffer[][];
-
-  constructor(leaves: T[], public readonly hashFn: (element: T) => string) {
-    this.elements = leaves.map((leaf) => this.leafToBuf(leaf));
-    // Sort elements
-    this.elements.sort(Buffer.compare);
-    // Deduplicate elements
-    this.elements = MerkleTree.bufDedup(this.elements);
-
-    this.bufferElementPositionIndex = this.elements.reduce<{ [hexElement: string]: number }>((memo, el, index) => {
-      memo[bufferToHex(el)] = index;
-      return memo;
-    }, {});
-
-    // Create layers
-    this.layers = this.getLayers(this.elements);
-  }
-
-  isEmpty(): boolean {
-    return this.layers.length === 0;
-  }
-
-  getLayers(elements: Buffer[]): Buffer[][] {
-    const layers: Buffer[][] = [];
-    if (elements.length === 0) return layers;
-
-    layers.push(elements);
-
-    // Get next layer until we reach the root
-    while (layers[layers.length - 1].length > 1) {
-      layers.push(this.getNextLayer(layers[layers.length - 1]));
-    }
-
-    return layers;
-  }
-
-  getNextLayer(elements: Buffer[]): Buffer[] {
-    return elements.reduce<Buffer[]>((layer, el, idx, arr) => {
-      if (idx % 2 === 0) {
-        // Hash the current element with its pair element
-        layer.push(MerkleTree.combinedHash(el, arr[idx + 1]));
-      }
-
-      return layer;
-    }, []);
-  }
-
-  static combinedHash(first: Buffer, second: Buffer): Buffer {
-    if (!first) {
-      return second;
-    }
-    if (!second) {
-      return first;
-    }
-
-    return keccak256(MerkleTree.sortAndConcat(first, second));
-  }
-
-  getRoot(): Buffer {
-    return this.layers[this.layers.length - 1][0];
-  }
-
-  getHexRoot(): string {
-    if (this.isEmpty()) return EMPTY_MERKLE_ROOT;
-    return bufferToHex(this.getRoot());
-  }
-
-  getProof(leaf: T) {
-    return this.getProofRawBuf(this.leafToBuf(leaf));
-  }
-
-  getHexProof(leaf: T) {
-    return this.getHexProofRawBuf(this.leafToBuf(leaf));
-  }
-
-  leafToBuf(element: T): Buffer {
-    const hash = this.hashFn(element);
-    const hexString = hash.startsWith("0x") ? hash.substring(2) : hash;
-    return Buffer.from(hexString.toLowerCase(), "hex");
-  }
-
-  // Methods that take the raw buffers (hashes).
-  getProofRawBuf(element: Buffer) {
-    let idx = this.bufferElementPositionIndex[bufferToHex(element)];
-
-    if (typeof idx !== "number") {
-      throw new Error("Element does not exist in Merkle tree");
-    }
-
-    return this.layers.reduce((proof, layer) => {
-      const pairElement = MerkleTree.getPairElement(idx, layer);
-
-      if (pairElement) {
-        proof.push(pairElement);
-      }
-
-      idx = Math.floor(idx / 2);
-
-      return proof;
-    }, []);
-  }
-
-  getHexProofRawBuf(el: Buffer): string[] {
-    const proof = this.getProofRawBuf(el);
-
-    return MerkleTree.bufArrToHexArr(proof);
-  }
-
-  private static getPairElement(idx: number, layer: Buffer[]): Buffer | null {
-    const pairIdx = idx % 2 === 0 ? idx + 1 : idx - 1;
-
-    if (pairIdx < layer.length) {
-      return layer[pairIdx];
-    } else {
-      return null;
-    }
-  }
-
-  private static bufDedup(elements: Buffer[]): Buffer[] {
-    return elements.filter((el, idx) => {
-      return idx === 0 || !elements[idx - 1].equals(el);
-    });
-  }
-
-  private static bufArrToHexArr(arr: Buffer[]): string[] {
-    if (arr.some((el) => !Buffer.isBuffer(el))) {
-      throw new Error("Array is not an array of buffers");
-    }
-
-    return arr.map((el) => "0x" + el.toString("hex"));
-  }
-
-  private static sortAndConcat(...args: Buffer[]): Buffer {
-    return Buffer.concat([...args].sort(Buffer.compare));
-  }
-}
+export { MerkleTree, EMPTY_MERKLE_ROOT } from "@uma/common";