From c57ee9fc1ce6a43999f96f7b821bc2b0025a42ce Mon Sep 17 00:00:00 2001
From: Michael Standen <mstan@horizon.io>
Date: Wed, 15 May 2024 11:07:15 +1200
Subject: [PATCH] Fix generate tree script

---
 scripts/generateMerkleTree.ts | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/scripts/generateMerkleTree.ts b/scripts/generateMerkleTree.ts
index 6b33cf0..3a46cf3 100644
--- a/scripts/generateMerkleTree.ts
+++ b/scripts/generateMerkleTree.ts
@@ -1,9 +1,14 @@
 import { MerkleTree } from 'merkletreejs'
-import { utils } from 'ethers'
+import { BigNumberish, utils } from 'ethers'
 import keccak256 from 'keccak256'
 
-const generateTree = (elements: string[]) => {
-  const hashed = elements.map(e => utils.solidityKeccak256(['uint256'], [e]))
+export type TreeElement = {
+  address: string
+  tokenId: BigNumberish
+}
+
+const generateTree = (elements: TreeElement[]) => {
+  const hashed = elements.map(e => getLeaf(e))
 
   const merkleTree = new MerkleTree(hashed, keccak256, {
     sort: true,
@@ -17,7 +22,13 @@ const generateTree = (elements: string[]) => {
   }
 }
 
-const generateProof = (tree: MerkleTree, element: string) =>
-  tree.getHexProof(utils.solidityKeccak256(['uint256'], [element]))
+const getLeaf = (element: TreeElement) =>
+  utils.solidityKeccak256(
+    ['address', 'uint256'],
+    [element.address.toLowerCase(), element.tokenId],
+  )
+
+const generateProof = (tree: MerkleTree, element: TreeElement) =>
+  tree.getHexProof(getLeaf(element))
 
-export { generateTree, generateProof }
+export { generateTree, generateProof, getLeaf }