reduce solidity memory usage and run the script in chunks

GenerationSoftware · Mar 21, 2024 · af6c111 · af6c111
1 parent ba155fa
commit af6c111
Show file tree

Hide file tree

Showing 5 changed files with 125 additions and 103 deletions.
diff --git a/index.js b/index.js
@@ -12,7 +12,7 @@ const rpcUrl = process.env.FWC_RPC_URL
 const run = async () => {
   const { chainId, multicallBatchSize, prizePoolAddress, vaultAddress, userAddresses } = JSON.parse(fs.readFileSync(inputFile, 'utf8'))
 
-  const allParams = Object.entries(await getAllParams(chainId, rpcUrl, prizePoolAddress, vaultAddress, userAddresses, { multicallBatchSize }))
+  const paramBatches = await getAllParams(chainId, rpcUrl, prizePoolAddress, vaultAddress, userAddresses, { multicallBatchSize })
 
   const scriptPath = join(rootDir, './sol/script/WinnerCalc.s.sol')
   const configPath = join(rootDir, './sol/foundry.toml')
@@ -22,16 +22,20 @@ const run = async () => {
 
   const winnerMap = new Map()
 
-  for(const [tier, params] of allParams) {
+  for(const { tier, params, chunkSize, prizeCount } of paramBatches) {
+    console.log(`Calculating prizes for tier ${tier}, checking ${chunkSize} addresses (${chunkSize * prizeCount} total picks)...`)
+
     try { fs.mkdirSync(dirname(paramsPath)) } catch {}
     fs.writeFileSync(paramsPath, params)
 
     await new Promise((resolve, reject) => {
-      exec(`forge script ${scriptPath}:WinnerCalcScript --config-path ${configPath} --cache-path ${cachePath}`, (error, stdout, stderr) => {
+      exec(`forge script ${scriptPath}:WinnerCalcScript --config-path ${configPath} --cache-path ${cachePath} --memory-limit 268435456 --skip-simulation`, (error, stdout, stderr) => {
         if(error ?? stderr) {
+          console.log(stderr)
           reject(error ?? stderr)
         } else {
           const winners = decodeWinners(fs.readFileSync(resultsPath, 'utf8'))
+          console.log(`Found ${winners.length} winner${winners.length == 1 ? "" : "s"}${winners.length > 0 ? "!" : "."}`)
 
           for(const winner of winners) {
             if(!winnerMap.has(winner.user)) {

diff --git a/sol/script/WinnerCalc.s.sol b/sol/script/WinnerCalc.s.sol
@@ -2,19 +2,35 @@
 pragma solidity ^0.8.21;
 
 import "forge-std/Script.sol";
-import { console2 } from "forge-std/console2.sol";
-import { Params, WinnerCalcLib, Winner } from "../src/WinnerCalcLib.sol";
+import { Params, WinnerCalc, Winner } from "../src/WinnerCalc.sol";
 
 contract WinnerCalcScript is Script {
 
-    Winner[] public winners;
+    WinnerCalc public winnerCalc;
 
     function run() public {
         string memory paramFilename = "files/params.txt";
         string memory outputFilename = "files/results.txt";
         Params memory params = abi.decode(vm.parseBytes(vm.readFile(paramFilename)), (Params));
-        WinnerCalcLib.getWinningPicks(params, winners);
-        vm.writeFile(outputFilename, vm.toString(abi.encode(winners)));
+
+        winnerCalc = new WinnerCalc();
+        for (uint256 userIndex = 0; userIndex < params.user.length; userIndex++) {
+            for (uint32 prizeIndex = 0; prizeIndex < params.tierPrizeCount; prizeIndex++) {
+                winnerCalc.checkWin(
+                    params.winningRandomNumber,
+                    params.lastAwardedDrawId,
+                    params.vault,
+                    params.vaultTotalSupplyTwab,
+                    params.vaultPortion,
+                    params.user[userIndex],
+                    params.userTwab[userIndex],
+                    params.tier,
+                    params.tierOdds,
+                    prizeIndex
+                );
+            }
+        }
+        winnerCalc.writeResults(vm, outputFilename);
     }
 
 }
diff --git a/sol/src/WinnerCalc.sol b/sol/src/WinnerCalc.sol
@@ -0,0 +1,70 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.21;
+
+import { Vm } from "forge-std/Vm.sol";
+import { TierCalculationLib, SD59x18 } from "pt-v5-prize-pool/PrizePool.sol";
+
+struct Params {
+    uint256 winningRandomNumber;
+    uint24 lastAwardedDrawId;
+    address vault;
+    uint8 tier;
+    uint32 tierPrizeCount;
+    SD59x18 tierOdds;
+    SD59x18 vaultPortion;
+    uint256 vaultTotalSupplyTwab;
+    address[] user;
+    uint256[] userTwab;
+}
+
+struct Winner {
+    address user;
+    uint32 prizeIndex;
+}
+
+contract WinnerCalc {
+
+    Winner[] public winners;
+
+    /// @dev pushes winners to the `winners` array based off the prize param info
+    function checkWin(
+        uint256 winningRandomNumber,
+        uint24 lastAwardedDrawId,
+        address vault,
+        uint256 vaultTotalSupplyTwab,
+        SD59x18 vaultPortion,
+        address user,
+        uint256 userTwab,
+        uint8 tier,
+        SD59x18 tierOdds,
+        uint32 prizeIndex
+    ) external {
+        uint256 _userSpecificRandomNumber = TierCalculationLib.calculatePseudoRandomNumber(
+            lastAwardedDrawId,
+            vault,
+            user,
+            tier,
+            prizeIndex,
+            winningRandomNumber
+        );
+        if (
+            TierCalculationLib.isWinner(
+                _userSpecificRandomNumber,
+                userTwab,
+                vaultTotalSupplyTwab,
+                vaultPortion,
+                tierOdds
+            )
+        ) {
+            winners.push(Winner({
+                user: user,
+                prizeIndex: prizeIndex
+            }));
+        }
+    }
+
+    function writeResults(Vm vm, string memory outputFilename) external {
+        vm.writeFile(outputFilename, vm.toString(abi.encode(winners)));
+    }
+
+}
diff --git a/sol/src/WinnerCalcLib.sol b/sol/src/WinnerCalcLib.sol
diff --git a/utils/params.ts b/utils/params.ts
@@ -5,7 +5,7 @@ import { getTwabs } from "./twab.js"
 import type { Address } from "viem"
 
 export const getAllParams = async (chainId: number, rpcUrl: string, prizePoolAddress: Address, vaultAddress: Address, userAddresses: Address[], options?: { multicallBatchSize?: number }) => {
-  const params: { [tier: number]: `0x${string}` } = {}
+  const params: { tier: number, params: `0x${string}`, chunkSize: number, prizeCount: number }[] = []
   const cachedTwabs: { [startTimestamp: number]: ReturnType<typeof getTwabs> } = {}
 
   const client = getClient(chainId, rpcUrl, options)
@@ -23,20 +23,32 @@ export const getAllParams = async (chainId: number, rpcUrl: string, prizePoolAdd
     }
     const { vaultTwab, userTwabs } = await cachedTwabs[startTwabTimestamp]
 
-    const encodedParams = encodeParams({
-      winningRandomNumber: prizePoolInfo.randomNumber,
-      lastAwardedDrawId: prizePoolInfo.lastAwardedDrawId,
-      vaultAddress,
-      tier,
-      tierPrizeCount: tierInfo[tier].prizeCount,
-      tierOdds: tierInfo[tier].odds,
-      vaultPortion,
-      vaultTwab,
-      userAddresses: userTwabs.map(user => user.address),
-      userTwabs: userTwabs.map(user => user.twab)
-    })
-
-    params[tier] = encodedParams
+    const chunkSize = Math.min(10_000, Math.ceil(1_000_000 / tierInfo[tier].prizeCount)) // varies between 1-1000 based on prize count
+    const userChunks: (typeof userTwabs)[] = []
+    for (let chunk = 0; userChunks.length < Math.ceil(userTwabs.length / chunkSize); chunk++) {
+      userChunks.push(
+        userTwabs.slice(
+          chunk * chunkSize,
+          Math.min((chunk + 1) * chunkSize, userTwabs.length)
+        )
+      )
+    }
+
+    for (const userChunk of userChunks) {
+      const encodedParams = encodeParams({
+        winningRandomNumber: prizePoolInfo.randomNumber,
+        lastAwardedDrawId: prizePoolInfo.lastAwardedDrawId,
+        vaultAddress,
+        tier,
+        tierPrizeCount: tierInfo[tier].prizeCount,
+        tierOdds: tierInfo[tier].odds,
+        vaultPortion,
+        vaultTwab,
+        userAddresses: userChunk.map(user => user.address),
+        userTwabs: userChunk.map(user => user.twab)
+      })
+      params.push({ tier, params: encodedParams, chunkSize: userChunk.length, prizeCount: tierInfo[tier].prizeCount })
+    }
   }))
 
   return params