refactor handleLiquidity

bnb-party · Sep 25, 2024 · 1cad278 · 1cad278
1 parent df52a30
commit 1cad278
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 43 deletions.
diff --git a/contracts/BNBPartyLiquidity.sol b/contracts/BNBPartyLiquidity.sol
@@ -8,63 +8,69 @@ import "./BNBPartyLiquidityHelper.sol";
 abstract contract BNBPartyLiquidity is BNBPartyLiquidityHelper {
     using SafeERC20 for IERC20;
 
-    /// @notice Handles liquidity by decreasing liquidity, collecting tokens, and creating a new liquidity pool.
+    /// @notice Handles liquidity by decreasing the liquidity, collecting tokens, and creating a new liquidity pool.
     /// @param recipient Address receiving the bonus BNB
     /// @dev Decreases liquidity, collects tokens, creates a new pool, and sends bonuses
-    function _handleLiquidity(address recipient) internal returns (address liquidityPool, uint256 tokenId) {
+    function _handleLiquidity(address recipient) internal returns (address liquidityPool, uint256 tokenId){
         IUniswapV3Pool pool = IUniswapV3Pool(msg.sender);
         (uint160 sqrtPriceX96, , , , , , ) = pool.slot0();
+        (address token0, address token1) = (pool.token0(), pool.token1());
         uint128 liquidity = pool.liquidity();
         uint256 unwrapAmount = party.bonusTargetReach + party.bonusPartyCreator + party.targetReachFee;
-
-        (address token0, address token1, address memeToken) = _getTokens(pool);
-
+        uint160 newSqrtPriceX96;
+        // Decrease liquidity and collect tokens
         (uint256 amount0, uint256 amount1) = _decreaseAndCollect(lpToTokenId[msg.sender], liquidity);
-        uint160 newSqrtPriceX96 = _getNewSqrtPrice(liquidity, sqrtPriceX96, unwrapAmount, token0);
-
-        _isToken0WBNB(token0) ? amount0 -= unwrapAmount : amount1 -= unwrapAmount;
-        isTokenTargetReached[memeToken] = true;
-
+        // Calculate new amounts and price
+        (newSqrtPriceX96, amount0, amount1) = _calculateAmountsAndPrice(token0, token1, amount0, amount1, unwrapAmount, sqrtPriceX96, liquidity);
         // Approve tokens for the new liquidity pool creation
         _approveTokensForLP(token0, token1, amount0, amount1);
-
         // Create new Liquidity Pool
         (liquidityPool, tokenId) = _createLP(positionManager, token0, token1, amount0, amount1, newSqrtPriceX96, party.lpFee, _getTicks(token0, party.lpTicks));
 
-        // Send bonuses and burn meme tokens
+        // Send bonuses
         _unwrapAndSendBNB(recipient, unwrapAmount);
-        _burnMemeToken(memeToken);
+        // burn meme tokens
+        _burnMemeToken(token0 == address(WBNB) ? token1 : token0);
     }
 
-    function _getTokens(IUniswapV3Pool pool) internal view returns (address token0, address token1, address memeToken) {
-        token0 = pool.token0();
-        token1 = pool.token1();
-        memeToken = _isToken0WBNB(token0) ? token1 : token0;
+    function _approveTokensForLP(
+        address token0,
+        address token1,
+        uint256 amount0,
+        uint256 amount1
+    ) internal {
+        IERC20(token0).safeIncreaseAllowance(address(positionManager), amount0);
+        IERC20(token1).safeIncreaseAllowance(address(positionManager), amount1);
     }
 
-    function _getNewSqrtPrice(
-        uint128 liquidity,
-        uint160 sqrtPriceX96,
+    function _calculateAmountsAndPrice(
+        address token0,
+        address token1,
+        uint256 amount0,
+        uint256 amount1,
         uint256 unwrapAmount,
-        address token0
-    )
-        internal
-        view
-        returns (uint160 newSqrtPriceX96)
-    {
-        if (_isToken0WBNB(token0)) {
+        uint160 sqrtPriceX96,
+        uint128 liquidity
+    ) internal returns (uint160 newSqrtPriceX96, uint256, uint256) {
+        if (token0 == address(WBNB)) {
+            amount0 -= unwrapAmount;
+            isTokenTargetReached[token1] = true;
             newSqrtPriceX96 = sqrtPriceCalculator.getNextSqrtPriceFromAmount0RoundingUp(
-                sqrtPriceX96, liquidity, unwrapAmount, false
-            );
+                    sqrtPriceX96,
+                    liquidity,
+                    unwrapAmount,
+                    false
+                );
         } else {
+            amount1 -= unwrapAmount;
+            isTokenTargetReached[token0] = true;
             newSqrtPriceX96 = sqrtPriceCalculator.getNextSqrtPriceFromAmount1RoundingDown(
-                sqrtPriceX96, liquidity, unwrapAmount, false
-            );
+                    sqrtPriceX96,
+                    liquidity,
+                    unwrapAmount,
+                    false
+                );
         }
-    }
-
-    function _approveTokensForLP(address token0, address token1, uint256 amount0, uint256 amount1) internal {
-        IERC20(token0).safeIncreaseAllowance(address(positionManager), amount0);
-        IERC20(token1).safeIncreaseAllowance(address(positionManager), amount1);
+        return (newSqrtPriceX96, amount0, amount1);
     }
 }
diff --git a/contracts/BNBPartyView.sol b/contracts/BNBPartyView.sol
@@ -16,13 +16,7 @@ abstract contract BNBPartyView is BNBPartyManageable {
         if (liquidityPool == IUniswapV3Pool(address(0))) {
             revert ZeroAddress();
         }
-        return _isToken0WBNB(liquidityPool.token0()); // Checks if WBNB is token0
-    }
-
-    function _isToken0WBNB(
-        address token0
-    ) internal view returns (bool) {
-        return token0 == address(WBNB); // Checks if WBNB is token0
+        return liquidityPool.token0() == address(WBNB); // Checks if WBNB is token0
     }
 
     /// @notice Calculates the fees earned based on liquidity and global fee growth