Improvements for H-02

src/base/BaseCustomAccounting.sol

@@ -48,7 +48,7 @@ abstract contract BaseCustomAccounting is BaseHook {
     error PoolNotInitialized();
 
     /**
-     * @dev Liquidity modification delta resulted in too much slippage.
+     * @dev Principal delta of liquidity modification resulted in too much slippage.
      */
     error TooMuchSlippage();
 
@@ -123,10 +123,11 @@ abstract contract BaseCustomAccounting is BaseHook {
      *
      * NOTE: This function doesn't revert if currency0 is not native and msg.value is non-zero, i.e.
      * the hook accepts native currency even if currency0 is not native to allow hooks for out-of-pool
-     * use cases.
+     * use cases. Additionally, the `amount0Min` and `amount1Min` parameters are relative to the
+     * principal delta, which excludes fees accrued from the liquidity modification delta.
      *
      * @param params The parameters for the liquidity addition.
-     * @return delta The balance delta of the liquidity addition from the `PoolManager`.
+     * @return delta The principal delta of the liquidity addition.
      */
     function addLiquidity(AddLiquidityParams calldata params)
         external
@@ -151,10 +152,13 @@ abstract contract BaseCustomAccounting is BaseHook {
         (bytes memory modifyParams, uint256 shares) = _getAddLiquidity(sqrtPriceX96, params);
 
         // Apply the liquidity modification
-        delta = _modifyLiquidity(modifyParams);
+        (BalanceDelta callerDelta, BalanceDelta feesAccrued) = _modifyLiquidity(modifyParams);
 
         // Mint the liquidity units to the `params.to` address
-        _mint(params, delta, shares);
+        _mint(params, callerDelta, shares);
+
+        // Get the principal delta by subtracting the fee delta from the caller delta (-= is not supported)
+        delta = callerDelta - feesAccrued;
 
         // Check for slippage
         uint128 amount0 = uint128(-delta.amount0());
@@ -173,8 +177,11 @@ abstract contract BaseCustomAccounting is BaseHook {
      *
      * @dev `msg.sender` should have already given the hook allowance of at least liquidity on the pool.
      *
+     * NOTE: The `amount0Min` and `amount1Min` parameters are relative to the principal delta, which
+     * excludes fees accrued from the liquidity modification delta.
+     *
      * @param params The parameters for the liquidity removal.
-     * @return delta The balance delta of the liquidity removal from the `PoolManager`.
+     * @return delta The principal delta of the liquidity removal.
      */
     function removeLiquidity(RemoveLiquidityParams calldata params)
         external
@@ -190,15 +197,16 @@ abstract contract BaseCustomAccounting is BaseHook {
         (bytes memory modifyParams, uint256 shares) = _getRemoveLiquidity(params);
 
         // Apply the liquidity modification
-        delta = _modifyLiquidity(modifyParams);
+        (BalanceDelta callerDelta, BalanceDelta feesAccrued) = _modifyLiquidity(modifyParams);
 
         // Burn the liquidity shares from the sender
-        _burn(params, delta, shares);
+        _burn(params, callerDelta, shares);
+
+        // Get the principal delta by subtracting the fee delta from the caller delta (-= is not supported)
+        delta = callerDelta - feesAccrued;
 
         // Check for slippage
-        uint128 amount0 = delta.amount0() < 0 ? uint128(-delta.amount0()) : uint128(delta.amount0());
-        uint128 amount1 = delta.amount1() < 0 ? uint128(-delta.amount1()) : uint128(delta.amount1());
-        if (amount0 < params.amount0Min || amount1 < params.amount1Min) {
+        if (uint128(delta.amount0()) < params.amount0Min || uint128(delta.amount1()) < params.amount1Min) {
             revert TooMuchSlippage();
         }
     }
@@ -207,15 +215,20 @@ abstract contract BaseCustomAccounting is BaseHook {
      * @dev Calls the `PoolManager` to unlock and call back the hook's `_unlockCallback` function.
      *
      * @param params The encoded parameters for the liquidity modification based on the `ModifyLiquidityParams` struct.
-     * @return delta The balance delta of the liquidity modification from the `PoolManager`.
+     * @return callerDelta The balance delta from the liquidity modification. This is the total of both principal and fee deltas.
+     * @return feesAccrued The balance delta of the fees generated in the liquidity range.
      */
     // slither-disable-next-line dead-code
-    function _modifyLiquidity(bytes memory params) internal virtual returns (BalanceDelta delta) {
-        delta = abi.decode(
+    function _modifyLiquidity(bytes memory params)
+        internal
+        virtual
+        returns (BalanceDelta callerDelta, BalanceDelta feesAccrued)
+    {
+        (callerDelta, feesAccrued) = abi.decode(
             poolManager.unlock(
                 abi.encode(CallbackData(msg.sender, abi.decode(params, (IPoolManager.ModifyLiquidityParams))))
             ),
-            (BalanceDelta)
+            (BalanceDelta, BalanceDelta)
         );
     }
 
@@ -228,7 +241,7 @@ abstract contract BaseCustomAccounting is BaseHook {
      * accordingly.
      *
      * @param rawData The encoded `CallbackData` struct.
-     * @return returnData The encoded balance delta of the liquidity modification from the `PoolManager`.
+     * @return returnData The encoded caller and fees accrued deltas.
      */
     function unlockCallback(bytes calldata rawData)
         external
@@ -240,29 +253,27 @@ abstract contract BaseCustomAccounting is BaseHook {
         PoolKey memory key = poolKey;
 
         // Get liquidity modification deltas
-        (BalanceDelta delta, BalanceDelta feeDelta) = poolManager.modifyLiquidity(key, data.params, "");
-
-        // Get the releveant delta by substracting the fee delta from the principal delta (-= is not supported)
-        delta = delta - feeDelta;
+        (BalanceDelta callerDelta, BalanceDelta feesAccrued) = poolManager.modifyLiquidity(key, data.params, "");
 
         // Handle each currency amount based on its sign
-        if (delta.amount0() < 0) {
+        if (callerDelta.amount0() < 0) {
             // If amount0 is negative, send tokens from the sender to the pool
-            key.currency0.settle(poolManager, data.sender, uint256(int256(-delta.amount0())), false);
+            key.currency0.settle(poolManager, data.sender, uint256(int256(-callerDelta.amount0())), false);
         } else {
             // If amount0 is positive, send tokens from the pool to the sender
-            key.currency0.take(poolManager, data.sender, uint256(int256(delta.amount0())), false);
+            key.currency0.take(poolManager, data.sender, uint256(int256(callerDelta.amount0())), false);
         }
 
-        if (delta.amount1() < 0) {
+        if (callerDelta.amount1() < 0) {
             // If amount1 is negative, send tokens from the sender to the pool
-            key.currency1.settle(poolManager, data.sender, uint256(int256(-delta.amount1())), false);
+            key.currency1.settle(poolManager, data.sender, uint256(int256(-callerDelta.amount1())), false);
         } else {
             // If amount1 is positive, send tokens from the pool to the sender
-            key.currency1.take(poolManager, data.sender, uint256(int256(delta.amount1())), false);
+            key.currency1.take(poolManager, data.sender, uint256(int256(callerDelta.amount1())), false);
         }
 
-        return abi.encode(delta);
+        // Return both deltas so that slippage checks can be done on the principal delta
+        return abi.encode(callerDelta, feesAccrued);
     }
 
     /**
@@ -347,7 +358,7 @@ abstract contract BaseCustomAccounting is BaseHook {
      * @dev Mint liquidity shares to the sender.
      *
      * @param params The parameters for the liquidity addition.
-     * @param delta The balance delta of the liquidity addition from the `PoolManager`.
+     * @param delta The balance delta from the liquidity addition. This is the total of both principal and fee deltas.
      * @param shares The liquidity shares to mint.
      */
     function _mint(AddLiquidityParams memory params, BalanceDelta delta, uint256 shares) internal virtual;
@@ -356,7 +367,7 @@ abstract contract BaseCustomAccounting is BaseHook {
      * @dev Burn liquidity shares from the sender.
      *
      * @param params The parameters for the liquidity removal.
-     * @param delta The balance delta of the liquidity removal from the `PoolManager`.
+     * @param delta The balance delta from the liquidity removal. This is the total of both principal and fee deltas.
      * @param shares The liquidity shares to burn.
      */
     function _burn(RemoveLiquidityParams memory params, BalanceDelta delta, uint256 shares) internal virtual;

src/base/BaseCustomCurve.sol

@@ -11,7 +11,7 @@ import {Currency} from "v4-core/src/types/Currency.sol";
 import {SafeCast} from "v4-core/src/libraries/SafeCast.sol";
 import {CurrencySettler} from "src/utils/CurrencySettler.sol";
 import {BeforeSwapDelta, toBeforeSwapDelta} from "v4-core/src/types/BeforeSwapDelta.sol";
-import {BalanceDelta, toBalanceDelta} from "v4-core/src/types/BalanceDelta.sol";
+import {BalanceDelta, toBalanceDelta, BalanceDeltaLibrary} from "v4-core/src/types/BalanceDelta.sol";
 
 /**
  * @dev Base implementation for custom curves, inheriting from {BaseCustomAccounting}.
@@ -127,20 +127,28 @@ abstract contract BaseCustomCurve is BaseCustomAccounting {
      *
      * @param params The parameters for the liquidity modification, encoded in the
      * {_getAddLiquidity} or {_getRemoveLiquidity} function.
-     * @return delta The balance delta of the liquidity modifications.
+     * @return callerDelta The balance delta from the liquidity modification. This is the total of both principal and fee deltas.
+     * @return feesAccrued The balance delta of the fees generated in the liquidity range.
      */
-    function _modifyLiquidity(bytes memory params) internal virtual override returns (BalanceDelta delta) {
+    function _modifyLiquidity(bytes memory params)
+        internal
+        virtual
+        override
+        returns (BalanceDelta callerDelta, BalanceDelta feesAccrued)
+    {
         (int128 amount0, int128 amount1) = abi.decode(params, (int128, int128));
-        delta =
-            abi.decode(poolManager.unlock(abi.encode(CallbackDataCustom(msg.sender, amount0, amount1))), (BalanceDelta));
+        (callerDelta, feesAccrued) = abi.decode(
+            poolManager.unlock(abi.encode(CallbackDataCustom(msg.sender, amount0, amount1))),
+            (BalanceDelta, BalanceDelta)
+        );
     }
 
     /**
      * @dev Decodes the callback data and applies the liquidity modifications, overriding the custom
      * accounting logic to mint and burn ERC-6909 claim tokens which are used in swaps.
      *
      * @param rawData The callback data encoded in the {_modifyLiquidity} function.
-     * @return returnData The encoded balance delta of the liquidity modification from the `PoolManager`.
+     * @return returnData The encoded caller and fees accrued deltas.
      */
     function unlockCallback(bytes calldata rawData)
         external
@@ -199,7 +207,8 @@ abstract contract BaseCustomCurve is BaseCustomAccounting {
             amount1 = -data.amount1;
         }
 
-        return abi.encode(toBalanceDelta(amount0, amount1));
+        // Return the encoded caller and fees accrued (zero by default) deltas
+        return abi.encode(toBalanceDelta(amount0, amount1), BalanceDeltaLibrary.ZERO_DELTA);
     }
 
     /**

test/base/BaseCustomAccounting.t.sol

@@ -233,6 +233,66 @@ contract BaseCustomAccountingTest is Test, Deployers {
         assertEq(liquidityTokenBal, 14545454545454545454);
     }
 
+    function test_addLiquidity_swapFeeThenAdd_succeeds(uint24 lpFee) public {
+        lpFee = uint24(bound(lpFee, 0, 1e6));
+        vm.prank(address(hook));
+        manager.updateDynamicLPFee(key, lpFee);
+
+        uint256 prevBalance0 = key.currency0.balanceOf(address(this));
+        uint256 prevBalance1 = key.currency1.balanceOf(address(this));
+
+        hook.addLiquidity(
+            BaseCustomAccounting.AddLiquidityParams(
+                10 ether, 10 ether, 9 ether, 9 ether, address(this), MAX_DEADLINE, MIN_TICK, MAX_TICK, bytes32(0)
+            )
+        );
+
+        uint256 liquidityTokenBal = hook.balanceOf(address(this));
+
+        assertEq(manager.getLiquidity(id), liquidityTokenBal);
+        assertEq(liquidityTokenBal, 10 ether);
+        assertEq(key.currency0.balanceOf(address(this)), prevBalance0 - 10 ether);
+        assertEq(key.currency1.balanceOf(address(this)), prevBalance1 - 10 ether);
+
+        if (lpFee == 0) {
+            vm.expectEmit(true, true, true, true, address(manager));
+            emit Swap(
+                id,
+                address(swapRouter),
+                -4142135623730950489,
+                2928932188134524755,
+                56022770974786139918731938227,
+                10 ether,
+                -6932,
+                0
+            );
+        }
+
+        IPoolManager.SwapParams memory params =
+            IPoolManager.SwapParams({zeroForOne: true, amountSpecified: -10 ether, sqrtPriceLimitX96: SQRT_PRICE_1_2});
+        PoolSwapTest.TestSettings memory settings =
+            PoolSwapTest.TestSettings({takeClaims: false, settleUsingBurn: false});
+
+        swapRouter.swap(key, params, settings, ZERO_BYTES);
+
+        if (lpFee == 0) {
+            assertEq(key.currency0.balanceOf(address(this)), prevBalance0 - (10 ether + 4142135623730950489));
+            assertEq(key.currency1.balanceOf(address(this)), prevBalance1 - 7071067811865475245);
+        }
+
+        hook.addLiquidity(
+            BaseCustomAccounting.AddLiquidityParams(
+                5 ether, 5 ether, 0, 0, address(this), MAX_DEADLINE, MIN_TICK, MAX_TICK, bytes32(0)
+            )
+        );
+
+        if (lpFee == 0) {
+            liquidityTokenBal = hook.balanceOf(address(this));
+            assertEq(manager.getLiquidity(id), liquidityTokenBal);
+            assertEq(liquidityTokenBal, 13535533905932737622);
+        }
+    }
+
     function test_addLiquidity_expired_revert() public {
         vm.expectRevert(BaseCustomAccounting.ExpiredPastDeadline.selector);
         hook.addLiquidity(