update fullRange and tests

Uniswap · Jun 26, 2024 · fe5c0e0 · fe5c0e0
1 parent 2d3586c
commit fe5c0e0
Show file tree

Hide file tree

Showing 7 changed files with 21 additions and 17 deletions.
diff --git a/.forge-snapshots/FullRangeAddInitialLiquidity.snap b/.forge-snapshots/FullRangeAddInitialLiquidity.snap
@@ -1 +1 @@
-311181
+311621
diff --git a/.forge-snapshots/FullRangeAddLiquidity.snap b/.forge-snapshots/FullRangeAddLiquidity.snap
@@ -1 +1 @@
-122990
+125693
diff --git a/.forge-snapshots/FullRangeInitialize.snap b/.forge-snapshots/FullRangeInitialize.snap
@@ -1 +1 @@
-1015181
+1015209
diff --git a/.forge-snapshots/FullRangeRemoveLiquidity.snap b/.forge-snapshots/FullRangeRemoveLiquidity.snap
@@ -1 +1 @@
-110566
+110668
diff --git a/.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap b/.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap
@@ -1 +1 @@
-240044
+240370
diff --git 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
@@ -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 {