From 12957aaa51e4af39c6c5a45c52c49801accfefda Mon Sep 17 00:00:00 2001
From: tim <timolsoncrypto@gmail.com>
Date: Tue, 7 Oct 2025 15:57:04 -0400
Subject: [PATCH] flashLoan gas test

---
 src/PartyPool.sol    |  3 +--
 test/GasTest.sol     |  6 +-----
 test/PartyPool.t.sol | 16 ++++++----------
 3 files changed, 8 insertions(+), 17 deletions(-)

diff --git a/src/PartyPool.sol b/src/PartyPool.sol
index 8ccbc02..2cc0664 100644
--- a/src/PartyPool.sol
+++ b/src/PartyPool.sol
@@ -411,13 +411,13 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
     {
         IERC20 token = IERC20(tokenAddr);
         require(amount <= token.balanceOf(address(this)));
+        uint256 tokenIndex = tokenAddressToIndexPlusOne[token] - 1;
         (uint256 fee, ) = _computeFee(amount, FLASH_FEE_PPM);
 
         // Compute protocol share of flash fee
         if (PROTOCOL_FEE_PPM > 0 && fee > 0) {
             uint256 protoShare = (fee * PROTOCOL_FEE_PPM) / 1_000_000; // floor
             if (protoShare > 0) {
-                uint256 tokenIndex = tokenAddressToIndexPlusOne[token] - 1;
                 protocolFeesOwed[tokenIndex] += protoShare;
             }
         }
@@ -427,7 +427,6 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         require(token.transferFrom(address(receiver), address(this), amount + fee));
 
         // Update cached balance for the borrowed token
-        uint256 tokenIndex = tokenAddressToIndexPlusOne[token] - 1;
         uint256 balAfter = token.balanceOf(address(this));
         _recordCachedBalance(tokenIndex, balAfter);
 
diff --git a/test/GasTest.sol b/test/GasTest.sol
index e851831..55168c3 100644
--- a/test/GasTest.sol
+++ b/test/GasTest.sol
@@ -34,10 +34,6 @@ contract FlashBorrower is IERC3156FlashBorrower {
         payer = _payer;
     }
 
-    function flash(address token, uint256 amount) external {
-        PartyPool(pool).flashLoan(IERC3156FlashBorrower(address(this)), token, amount, "");
-    }
-
     function onFlashLoan(
         address /*initiator*/,
         address token,
@@ -416,7 +412,7 @@ contract GasTest is Test {
 
         // Execute flash loan 10 times to measure gas
         for (uint256 i = 0; i < 10; i++) {
-            borrower.flash(token, amount);
+            pool2.flashLoan(borrower, token, amount, "");
         }
     }
 }
diff --git a/test/PartyPool.t.sol b/test/PartyPool.t.sol
index 499ac79..73eab90 100644
--- a/test/PartyPool.t.sol
+++ b/test/PartyPool.t.sol
@@ -37,10 +37,6 @@ contract FlashBorrower is IERC3156FlashBorrower {
         payer = _payer;
     }
 
-    function flash(address token, uint256 amount) external {
-        PartyPool(pool).flashLoan(IERC3156FlashBorrower(address(this)), token, amount, "");
-    }
-
     function onFlashLoan(
         address /*initiator*/,
         address token,
@@ -850,7 +846,7 @@ contract PartyPoolTest is Test {
         uint256 poolToken0Before = token0.balanceOf(address(pool));
 
         // Execute flash loan
-        borrower.flash(address(token0), amount);
+        pool.flashLoan(borrower, address(token0), amount, "");
 
         // Net change for alice should equal the flash fee (principal is returned during repayment)
         uint256 fee = (amount * pool.flashFeePpm() + 1_000_000 - 1) / 1_000_000; // ceil fee calculation
@@ -882,7 +878,7 @@ contract PartyPoolTest is Test {
 
         // Execute flash loan - should revert due to insufficient allowance when pool tries to pull repayment
         vm.expectRevert();
-        borrower.flash(address(token0), amount);
+        pool.flashLoan(borrower, address(token0), amount, "");
     }
 
     /// @notice Test flash loan with partial repayment (should revert)
@@ -897,7 +893,7 @@ contract PartyPoolTest is Test {
 
         // Execute flash loan - should revert due to insufficient allowance when pool tries to pull full repayment
         vm.expectRevert();
-        borrower.flash(address(token0), amount);
+        pool.flashLoan(borrower, address(token0), amount, "");
     }
 
     /// @notice Test flash loan with principal repayment but no fee (should revert)
@@ -913,10 +909,10 @@ contract PartyPoolTest is Test {
         // Execute flash loan - should revert due to insufficient allowance if fee > 0
         if (pool.flashFeePpm() > 0) {
             vm.expectRevert();
-            borrower.flash(address(token0), amount);
+            pool.flashLoan(borrower, address(token0), amount, "");
         } else {
             // If fee is zero, this should succeed
-            borrower.flash(address(token0), amount);
+            pool.flashLoan(borrower, address(token0), amount, "");
         }
     }
 
@@ -935,7 +931,7 @@ contract PartyPoolTest is Test {
         uint256 poolToken0Before = token0.balanceOf(address(pool));
 
         // Execute flash loan
-        borrower.flash(address(token0), amount);
+        pool.flashLoan(borrower, address(token0), amount, "");
 
         // Check balances: net change for alice should equal the fee
         uint256 fee = (amount * pool.flashFeePpm() + 1_000_000 - 1) / 1_000_000; // ceil fee calculation