From fe5c0e0b2af35e10bfe50a7982e47167cc051726 Mon Sep 17 00:00:00 2001
From: Junion <69495294+Jun1on@users.noreply.github.com>
Date: Wed, 26 Jun 2024 00:58:31 -0400
Subject: [PATCH] update fullRange and tests

---
 .../FullRangeAddInitialLiquidity.snap         |  2 +-
 .forge-snapshots/FullRangeAddLiquidity.snap   |  2 +-
 .forge-snapshots/FullRangeInitialize.snap     |  2 +-
 .../FullRangeRemoveLiquidity.snap             |  2 +-
 .../FullRangeRemoveLiquidityAndRebalance.snap |  2 +-
 contracts/hooks/examples/FullRange.sol        | 23 +++++++++++--------
 test/FullRange.t.sol                          |  5 ++--
 7 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/.forge-snapshots/FullRangeAddInitialLiquidity.snap b/.forge-snapshots/FullRangeAddInitialLiquidity.snap
index 404cf12a..3c0a4b77 100644
--- a/.forge-snapshots/FullRangeAddInitialLiquidity.snap
+++ b/.forge-snapshots/FullRangeAddInitialLiquidity.snap
@@ -1 +1 @@
-311181
\ No newline at end of file
+311621
\ No newline at end of file
diff --git a/.forge-snapshots/FullRangeAddLiquidity.snap b/.forge-snapshots/FullRangeAddLiquidity.snap
index a4a14676..c8909575 100644
--- a/.forge-snapshots/FullRangeAddLiquidity.snap
+++ b/.forge-snapshots/FullRangeAddLiquidity.snap
@@ -1 +1 @@
-122990
\ No newline at end of file
+125693
\ No newline at end of file
diff --git a/.forge-snapshots/FullRangeInitialize.snap b/.forge-snapshots/FullRangeInitialize.snap
index 7a0170eb..6fa606d0 100644
--- a/.forge-snapshots/FullRangeInitialize.snap
+++ b/.forge-snapshots/FullRangeInitialize.snap
@@ -1 +1 @@
-1015181
\ No newline at end of file
+1015209
\ No newline at end of file
diff --git a/.forge-snapshots/FullRangeRemoveLiquidity.snap b/.forge-snapshots/FullRangeRemoveLiquidity.snap
index feea4936..998dcafd 100644
--- a/.forge-snapshots/FullRangeRemoveLiquidity.snap
+++ b/.forge-snapshots/FullRangeRemoveLiquidity.snap
@@ -1 +1 @@
-110566
\ No newline at end of file
+110668
\ No newline at end of file
diff --git a/.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap b/.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap
index e0df7eb7..0acd884d 100644
--- a/.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap
+++ b/.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap
@@ -1 +1 @@
-240044
\ No newline at end of file
+240370
\ No newline at end of file
diff --git a/contracts/hooks/examples/FullRange.sol b/contracts/hooks/examples/FullRange.sol
index acc41b13..722ed28f 100644
--- a/contracts/hooks/examples/FullRange.sol
+++ b/contracts/hooks/examples/FullRange.sol
@@ -145,6 +145,7 @@ contract FullRange is BaseHook {
         if (poolLiquidity == 0 && liquidity <= MINIMUM_LIQUIDITY) {
             revert LiquidityDoesntMeetMinimum();
         }
+
         BalanceDelta addedDelta = modifyLiquidity(
             key,
             IPoolManager.ModifyLiquidityParams({
@@ -157,12 +158,14 @@ contract FullRange is BaseHook {
 
         if (poolLiquidity == 0) {
             // permanently lock the first MINIMUM_LIQUIDITY tokens
-            liquidity -= MINIMUM_LIQUIDITY;
             UniswapV4ERC20(pool.liquidityToken).mint(address(0), MINIMUM_LIQUIDITY);
+            UniswapV4ERC20(pool.liquidityToken).mint(params.to, liquidity - MINIMUM_LIQUIDITY);
+        } else {
+            uint256 liquidityMinted = uint256(liquidity) * UniswapV4ERC20(pool.liquidityToken).totalSupply()
+                / uint256(manager.getLiquidity(poolId) - liquidity);
+            UniswapV4ERC20(pool.liquidityToken).mint(params.to, liquidityMinted);
         }
 
-        UniswapV4ERC20(pool.liquidityToken).mint(params.to, liquidity);
-
         if (uint128(-addedDelta.amount0()) < params.amount0Min || uint128(-addedDelta.amount1()) < params.amount1Min) {
             revert TooMuchSlippage();
         }
@@ -281,10 +284,6 @@ contract FullRange is BaseHook {
         PoolId poolId = key.toId();
         PoolInfo storage pool = poolInfo[poolId];
 
-        if (pool.hasAccruedFees) {
-            _rebalance(key);
-        }
-
         uint256 liquidityToRemove = FullMath.mulDiv(
             uint256(-params.liquidityDelta),
             manager.getLiquidity(poolId),
@@ -293,13 +292,19 @@ contract FullRange is BaseHook {
 
         params.liquidityDelta = -(liquidityToRemove.toInt256());
         (delta,) = manager.modifyLiquidity(key, params, ZERO_BYTES);
-        pool.hasAccruedFees = false;
     }
 
     function _unlockCallback(bytes calldata rawData) internal override returns (bytes memory) {
         CallbackData memory data = abi.decode(rawData, (CallbackData));
         BalanceDelta delta;
 
+        PoolId poolId = data.key.toId();
+        PoolInfo storage pool = poolInfo[data.key.toId()];
+        if (pool.hasAccruedFees) {
+            pool.hasAccruedFees = false;
+            rebalance(data.key);
+        }
+
         if (data.params.liquidityDelta < 0) {
             delta = _removeLiquidity(data.key, data.params);
             _takeDeltas(data.sender, data.key, delta);
@@ -310,7 +315,7 @@ contract FullRange is BaseHook {
         return abi.encode(delta);
     }
 
-    function _rebalance(PoolKey memory key) public {
+    function rebalance(PoolKey memory key) public {
         PoolId poolId = key.toId();
         (BalanceDelta balanceDelta,) = manager.modifyLiquidity(
             key,
diff --git a/test/FullRange.t.sol b/test/FullRange.t.sol
index 49698434..10edd943 100644
--- a/test/FullRange.t.sol
+++ b/test/FullRange.t.sol
@@ -295,9 +295,8 @@ contract TestFullRange is Test, Deployers, GasSnapshot {
         (hasAccruedFees,) = fullRange.poolInfo(id);
         liquidityTokenBal = UniswapV4ERC20(liquidityToken).balanceOf(address(this));
 
-        assertEq(manager.getLiquidity(id), liquidityTokenBal + LOCKED_LIQUIDITY);
-        assertEq(liquidityTokenBal, 14546694553059925434 - LOCKED_LIQUIDITY);
-        assertEq(hasAccruedFees, true);
+        assertTrue(manager.getLiquidity(id) > liquidityTokenBal + LOCKED_LIQUIDITY);
+        assertEq(hasAccruedFees, false);
     }
 
     function testFullRange_addLiquidity_FailsIfTooMuchSlippage() public {