From f5cab4a6208049599b9179508c9b398d5552434b Mon Sep 17 00:00:00 2001
From: Channing <channing@magiceden.io>
Date: Wed, 13 Mar 2024 15:31:36 -0700
Subject: [PATCH] Update scripts

---
 scripts/setStages.ts | 78 ++++++++++++++++++++++++++++++--------------
 1 file changed, 53 insertions(+), 25 deletions(-)

diff --git a/scripts/setStages.ts b/scripts/setStages.ts
index cd92b9b..e3a99a4 100644
--- a/scripts/setStages.ts
+++ b/scripts/setStages.ts
@@ -20,6 +20,7 @@ interface StageConfig {
   walletLimit?: number;
   maxSupply?: number;
   whitelistPath?: string;
+  variableWalletLimitPath?: string;
 }
 
 export const setStages = async (
@@ -41,35 +42,62 @@ export const setStages = async (
   if (args.gaslimit) {
     overrides.gasLimit = ethers.BigNumber.from(args.gaslimit);
   }
+
   const merkleRoots = await Promise.all(
     stagesConfig.map((stage) => {
-      if (!stage.whitelistPath) {
-        return ethers.utils.hexZeroPad('0x', 32);
-      }
-      const whitelist = JSON.parse(
-        fs.readFileSync(stage.whitelistPath, 'utf-8'),
-      );
-
-      // Clean up whitelist
-      const filteredWhitelist=  whitelist.filter((address: string) => ethers.utils.isAddress(address));
-      console.log(`Filtered whitelist: ${filteredWhitelist.length} addresses. ${whitelist.length - filteredWhitelist.length} invalid addresses removed.`);
-      const invalidWhitelist=  whitelist.filter((address: string) => !ethers.utils.isAddress(address));
-      console.log(`❌ Invalid whitelist: ${invalidWhitelist.length} addresses.\r\n${invalidWhitelist.join(', \r\n')}`);
-
-      if (invalidWhitelist.length > 0) {
-        console.log(`🔄 🚨 updating whitelist file: ${stage.whitelistPath}`);
-        fs.writeFileSync(stage.whitelistPath, JSON.stringify(filteredWhitelist, null, 2))
-      }
+      if (stage.whitelistPath) {
+        const whitelist = JSON.parse(
+          fs.readFileSync(stage.whitelistPath, 'utf-8'),
+        );
+
+        // Clean up whitelist
+        const filteredWhitelist=  whitelist.filter((address: string) => ethers.utils.isAddress(address));
+        console.log(`Filtered whitelist: ${filteredWhitelist.length} addresses. ${whitelist.length - filteredWhitelist.length} invalid addresses removed.`);
+        const invalidWhitelist=  whitelist.filter((address: string) => !ethers.utils.isAddress(address));
+        console.log(`❌ Invalid whitelist: ${invalidWhitelist.length} addresses.\r\n${invalidWhitelist.join(', \r\n')}`);
+
+        if (invalidWhitelist.length > 0) {
+          console.log(`🔄 🚨 updating whitelist file: ${stage.whitelistPath}`);
+          fs.writeFileSync(stage.whitelistPath, JSON.stringify(filteredWhitelist, null, 2))
+        }
 
-      const mt = new MerkleTree(
-        filteredWhitelist.map(ethers.utils.getAddress),
-        ethers.utils.keccak256,
-        {
+        const mt = new MerkleTree(
+          filteredWhitelist.map(ethers.utils.getAddress),
+          ethers.utils.keccak256,
+          {
+            sortPairs: true,
+            hashLeaves: true,
+          },
+        );
+        return mt.getHexRoot();
+      } else if (stage.variableWalletLimitPath) {
+        const leaves: any[] = [];
+        const file = fs.readFileSync(stage.variableWalletLimitPath, 'utf-8');
+        file
+          .split('\n')
+          .filter((line) => line)
+          .forEach((line) => {
+            const [addressStr, limitStr] = line.split(',');
+            const limit = parseInt(limitStr, 10);
+
+            if (!ethers.utils.isAddress(addressStr.trim().toLowerCase()) && limit < 0) {
+              console.log(`Filtered address: ${addressStr}`);
+              return;
+            }
+
+            const address = ethers.utils.getAddress(addressStr.trim().toLowerCase());
+            const digest = ethers.utils.solidityKeccak256(['address', 'uint32'], [address, limit]);
+            leaves.push(digest);
+          });
+
+        const mt = new MerkleTree(leaves, ethers.utils.keccak256, {
           sortPairs: true,
-          hashLeaves: true,
-        },
-      );
-      return mt.getHexRoot();
+          hashLeaves: false,
+        });
+        return mt.getHexRoot();
+      }
+
+      return ethers.utils.hexZeroPad('0x', 32);
     }),
   );