fix tests and use BalanceDeltas (#130)

saucepoint · web-flow · commit bc17c2270974 · 2024-06-26T13:59:05.000-04:00
* fix some assertions

* use BalanceDeltas for arithmetic

* cleanest code in the game???

* additional cleaning

* typo lol

* autocompound gas benchmarks

* autocompound excess credit gas benchmark

* save 600 gas, cleaner code when moving caller delta to tokensOwed
diff --git a/.forge-snapshots/autocompound_exactUnclaimedFees.snap b/.forge-snapshots/autocompound_exactUnclaimedFees.snap
@@ -0,0 +1 @@
+258477
diff --git a/.forge-snapshots/autocompound_exactUnclaimedFees_exactCustodiedFees.snap b/.forge-snapshots/autocompound_exactUnclaimedFees_exactCustodiedFees.snap
@@ -0,0 +1 @@
+190850
diff --git a/.forge-snapshots/autocompound_excessFeesCredit.snap b/.forge-snapshots/autocompound_excessFeesCredit.snap
@@ -0,0 +1 @@
+279016
diff --git a/.forge-snapshots/increaseLiquidity_erc20.snap b/.forge-snapshots/increaseLiquidity_erc20.snap
@@ -1 +1 @@
-176386
+171241
diff --git a/.forge-snapshots/increaseLiquidity_erc6909.snap b/.forge-snapshots/increaseLiquidity_erc6909.snap
@@ -1 +1 @@
-151968
+146823
diff --git a/.forge-snapshots/mintWithLiquidity.snap b/.forge-snapshots/mintWithLiquidity.snap
@@ -1 +1 @@
-471675
+466530
diff --git a/contracts/base/BaseLiquidityManagement.sol b/contracts/base/BaseLiquidityManagement.sol
@@ -23,6 +23,7 @@ import {FeeMath} from "../libraries/FeeMath.sol";
 import {LiquiditySaltLibrary} from "../libraries/LiquiditySaltLibrary.sol";
 import {IBaseLiquidityManagement} from "../interfaces/IBaseLiquidityManagement.sol";
 import {PositionLibrary} from "../libraries/Position.sol";
+import {BalanceDeltaExtensionLibrary} from "../libraries/BalanceDeltaExtensionLibrary.sol";
 
 import "forge-std/console2.sol";
 
@@ -38,6 +39,7 @@ contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallback {
     using SafeCast for uint256;
     using LiquiditySaltLibrary for IHooks;
     using PositionLibrary for IBaseLiquidityManagement.Position;
+    using BalanceDeltaExtensionLibrary for BalanceDelta;
 
     mapping(address owner => mapping(LiquidityRangeId rangeId => Position)) public positions;
 
@@ -106,7 +108,7 @@ contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallback {
         LiquidityRange memory range,
         uint256 liquidityToAdd,
         bytes memory hookData
-    ) internal returns (BalanceDelta, BalanceDelta) {
+    ) internal returns (BalanceDelta callerDelta, BalanceDelta thisDelta) {
         // Note that the liquidityDelta includes totalFeesAccrued. The totalFeesAccrued is returned separately for accounting purposes.
         (BalanceDelta liquidityDelta, BalanceDelta totalFeesAccrued) =
             _modifyLiquidity(owner, range, liquidityToAdd.toInt256(), hookData);
@@ -126,66 +128,37 @@ contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallback {
             position.liquidity
         );
 
-        console2.log(callerFeesAccrued.amount0());
-        console2.log(callerFeesAccrued.amount1());
-        console2.log("totalFees");
-        console2.log(totalFeesAccrued.amount0());
-        console2.log(totalFeesAccrued.amount1());
+        if (totalFeesAccrued == callerFeesAccrued) {
+            // when totalFeesAccrued == callerFeesAccrued, the caller is not sharing the range
+            // therefore, the caller is responsible for the entire liquidityDelta
+            callerDelta = liquidityDelta;
+        } else {
+            // the delta for increasing liquidity assuming that totalFeesAccrued was not applied
+            BalanceDelta principalDelta = liquidityDelta - totalFeesAccrued;
+
+            // outstanding deltas the caller is responsible for, after their fees are credited to the principal delta
+            callerDelta = principalDelta + callerFeesAccrued;
 
-        // Calculate the accurate tokens owed to the caller.
-        // If the totalFeesAccrued equals the callerFeesAccrued then the total owed to the caller is just the liquidityDelta.
-        // If the totalFeesAccrued is greater than the callerFeesAccrued, we must account for the difference.
-        // TODO: If totalFeesAccrued == callerFeesAccrued, I think we can just apply the entire delta onto the caller, even if this implicitly collects on behalf of another user in the same range.
-        (int128 callerDelta0, int128 callerDelta1) = totalFeesAccrued != callerFeesAccrued
-            ? _calculateCallerDeltas(liquidityDelta, totalFeesAccrued, callerFeesAccrued)
-            : (liquidityDelta.amount0(), liquidityDelta.amount1());
+            // outstanding deltas this contract is responsible for, intuitively the contract is responsible for taking fees external to the caller's accrued fees
+            thisDelta = totalFeesAccrued - callerFeesAccrued;
+        }
 
         // Update position storage, flushing the callerDelta value to tokensOwed first if necessary.
         // If callerDelta > 0, then even after investing callerFeesAccrued, the caller still has some amount to collect that were not added into the position so they are accounted to tokensOwed and removed from the final callerDelta returned.
-        uint128 tokensOwed0 = 0;
-        uint128 tokensOwed1 = 0;
-        (tokensOwed0, callerDelta0) = callerDelta0 > 0 ? (uint128(callerDelta0), int128(0)) : (uint128(0), callerDelta0);
-        (tokensOwed1, callerDelta1) = callerDelta1 > 0 ? (uint128(callerDelta1), int128(0)) : (uint128(0), callerDelta1);
+        BalanceDelta tokensOwed;
+        if (callerDelta.amount0() > 0) {
+            (tokensOwed, callerDelta, thisDelta) =
+                _moveCallerDeltaToTokensOwed(true, tokensOwed, callerDelta, thisDelta);
+        }
+
+        if (callerDelta.amount1() > 0) {
+            (tokensOwed, callerDelta, thisDelta) =
+                _moveCallerDeltaToTokensOwed(false, tokensOwed, callerDelta, thisDelta);
+        }
 
-        position.addTokensOwed(tokensOwed0, tokensOwed1);
+        position.addTokensOwed(tokensOwed);
         position.addLiquidity(liquidityToAdd);
         position.updateFeeGrowthInside(feeGrowthInside0X128, feeGrowthInside1X128);
-
-        // The delta owed or credited by this contract.
-        // TODO @sauce check that if callerDelta == 0 (zerod out from above), then this line just credits the posm to takes on behalf of the caller
-        int128 thisDelta0 = liquidityDelta.amount0() - callerDelta0;
-        int128 thisDelta1 = liquidityDelta.amount1() - callerDelta1;
-
-        return (toBalanceDelta(callerDelta0, callerDelta1), toBalanceDelta(thisDelta0, thisDelta1));
-    }
-
-    // Returns the delta paid/credited by/to the caller.
-    function _calculateCallerDeltas(
-        BalanceDelta liquidityDelta,
-        BalanceDelta totalFeesAccrued,
-        BalanceDelta callerFeesAccrued
-    ) private pure returns (int128 callerDelta0, int128 callerDelta1) {
-        (int128 liquidityDelta0, int128 liquidityDelta1) = (liquidityDelta.amount0(), liquidityDelta.amount1());
-        (int128 totalFeesAccrued0, int128 totalFeesAccrued1) = (totalFeesAccrued.amount0(), totalFeesAccrued.amount1());
-        (int128 callerFeesAccrued0, int128 callerFeesAccrued1) =
-            (callerFeesAccrued.amount0(), callerFeesAccrued.amount1());
-
-        callerDelta0 = _calculateCallerDelta(liquidityDelta0, totalFeesAccrued0, callerFeesAccrued0);
-        callerDelta1 = _calculateCallerDelta(liquidityDelta1, totalFeesAccrued1, callerFeesAccrued1);
-    }
-
-    function _calculateCallerDelta(int128 liquidityDelta, int128 totalFeesAccrued, int128 callerFeesAccrued)
-        private
-        pure
-        returns (int128 callerDelta)
-    {
-        unchecked {
-            // The principle delta owed/debited to the caller before any LP fees are deducted.
-            int128 principleDelta = liquidityDelta - totalFeesAccrued;
-            // The new caller delta is this principle delta plus the callerFeesAccrued which consists of
-            // the custodied fees by posm and unclaimed fees from the modifyLiq call.
-            callerDelta = principleDelta + callerFeesAccrued;
-        }
     }
 
     function _increaseLiquidityAndZeroOut(
@@ -230,6 +203,26 @@ contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallback {
         if (delta1 < 0) currency1.settle(manager, address(this), uint256(int256(-delta1)), true);
     }
 
+    function _moveCallerDeltaToTokensOwed(
+        bool useAmount0,
+        BalanceDelta tokensOwed,
+        BalanceDelta callerDelta,
+        BalanceDelta thisDelta
+    ) private returns (BalanceDelta, BalanceDelta, BalanceDelta) {
+        // credit the excess tokens to the position's tokensOwed
+        tokensOwed =
+            useAmount0 ? tokensOwed.setAmount0(callerDelta.amount0()) : tokensOwed.setAmount1(callerDelta.amount1());
+
+        // this contract is responsible for custodying the excess tokens
+        thisDelta =
+            useAmount0 ? thisDelta.addAmount0(callerDelta.amount0()) : thisDelta.addAmount1(callerDelta.amount1());
+
+        // the caller is not expected to collect the excess tokens
+        callerDelta = useAmount0 ? callerDelta.setAmount0(0) : callerDelta.setAmount1(0);
+
+        return (tokensOwed, callerDelta, thisDelta);
+    }
+
     function _lockAndIncreaseLiquidity(
         address owner,
         LiquidityRange memory range,
diff --git a/contracts/libraries/BalanceDeltaExtensionLibrary.sol b/contracts/libraries/BalanceDeltaExtensionLibrary.sol
@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+
+library BalanceDeltaExtensionLibrary {
+    function setAmount0(BalanceDelta a, int128 amount0) internal pure returns (BalanceDelta) {
+        assembly {
+            // set the upper 128 bits of a to amount0
+            a := or(shl(128, amount0), and(sub(shl(128, 1), 1), a))
+        }
+        return a;
+    }
+
+    function setAmount1(BalanceDelta a, int128 amount1) internal pure returns (BalanceDelta) {
+        assembly {
+            // set the lower 128 bits of a to amount1
+            a := or(and(shl(128, sub(shl(128, 1), 1)), a), amount1)
+        }
+        return a;
+    }
+
+    function addAmount0(BalanceDelta a, int128 amount0) internal pure returns (BalanceDelta) {
+        assembly {
+            let a0 := sar(128, a)
+            let res0 := add(a0, amount0)
+            a := or(shl(128, res0), and(sub(shl(128, 1), 1), a))
+        }
+        return a;
+    }
+
+    function addAmount1(BalanceDelta a, int128 amount1) internal pure returns (BalanceDelta) {
+        assembly {
+            let a1 := signextend(15, a)
+            let res1 := add(a1, amount1)
+            a := or(and(shl(128, sub(shl(128, 1), 1)), a), res1)
+        }
+        return a;
+    }
+
+    function addAndAssign(BalanceDelta a, BalanceDelta b) internal pure returns (BalanceDelta) {
+        assembly {
+            let a0 := sar(128, a)
+            let a1 := signextend(15, a)
+            let b0 := sar(128, b)
+            let b1 := signextend(15, b)
+            let res0 := add(a0, b0)
+            let res1 := add(a1, b1)
+            a := or(shl(128, res0), and(sub(shl(128, 1), 1), res1))
+        }
+        return a;
+    }
+}
diff --git a/contracts/libraries/Position.sol b/contracts/libraries/Position.sol
@@ -2,15 +2,14 @@
 pragma solidity >=0.8.20;
 
 import {IBaseLiquidityManagement} from "../interfaces/IBaseLiquidityManagement.sol";
+import {BalanceDelta} from "v4-core/types/BalanceDelta.sol";
 
 // Updates Position storage
 library PositionLibrary {
     // TODO ensure this is one sstore.
-    function addTokensOwed(IBaseLiquidityManagement.Position storage position, uint128 tokensOwed0, uint128 tokensOwed1)
-        internal
-    {
-        position.tokensOwed0 += tokensOwed0;
-        position.tokensOwed1 += tokensOwed1;
+    function addTokensOwed(IBaseLiquidityManagement.Position storage position, BalanceDelta tokensOwed) internal {
+        position.tokensOwed0 += uint128(tokensOwed.amount0());
+        position.tokensOwed1 += uint128(tokensOwed.amount1());
     }
 
     function addLiquidity(IBaseLiquidityManagement.Position storage position, uint256 liquidity) internal {
diff --git a/test/position-managers/Gas.t.sol b/test/position-managers/Gas.t.sol
@@ -56,6 +56,20 @@ contract GasTest is Test, Deployers, GasSnapshot {
         IERC20(Currency.unwrap(currency0)).approve(address(lpm), type(uint256).max);
         IERC20(Currency.unwrap(currency1)).approve(address(lpm), type(uint256).max);
 
+        // Give tokens to Alice and Bob, with approvals
+        IERC20(Currency.unwrap(currency0)).transfer(alice, STARTING_USER_BALANCE);
+        IERC20(Currency.unwrap(currency1)).transfer(alice, STARTING_USER_BALANCE);
+        IERC20(Currency.unwrap(currency0)).transfer(bob, STARTING_USER_BALANCE);
+        IERC20(Currency.unwrap(currency1)).transfer(bob, STARTING_USER_BALANCE);
+        vm.startPrank(alice);
+        IERC20(Currency.unwrap(currency0)).approve(address(lpm), type(uint256).max);
+        IERC20(Currency.unwrap(currency1)).approve(address(lpm), type(uint256).max);
+        vm.stopPrank();
+        vm.startPrank(bob);
+        IERC20(Currency.unwrap(currency0)).approve(address(lpm), type(uint256).max);
+        IERC20(Currency.unwrap(currency1)).approve(address(lpm), type(uint256).max);
+        vm.stopPrank();
+
         // mint some ERC6909 tokens
         claimsRouter.deposit(currency0, address(this), 100_000_000 ether);
         claimsRouter.deposit(currency1, address(this), 100_000_000 ether);
@@ -102,6 +116,119 @@ contract GasTest is Test, Deployers, GasSnapshot {
         snapLastCall("increaseLiquidity_erc6909");
     }
 
+    function test_gas_autocompound_exactUnclaimedFees() public {
+        // Alice and Bob provide liquidity on the range
+        // Alice uses her exact fees to increase liquidity (compounding)
+
+        uint256 liquidityAlice = 3_000e18;
+        uint256 liquidityBob = 1_000e18;
+
+        // alice provides liquidity
+        vm.prank(alice);
+        (uint256 tokenIdAlice,) = lpm.mint(range, liquidityAlice, block.timestamp + 1, alice, ZERO_BYTES);
+
+        // bob provides liquidity
+        vm.prank(bob);
+        lpm.mint(range, liquidityBob, block.timestamp + 1, bob, ZERO_BYTES);
+
+        // donate to create fees
+        donateRouter.donate(key, 0.2e18, 0.2e18, ZERO_BYTES);
+
+        // alice uses her exact fees to increase liquidity
+        (uint256 token0Owed, uint256 token1Owed) = lpm.feesOwed(tokenIdAlice);
+
+        (uint160 sqrtPriceX96,,,) = StateLibrary.getSlot0(manager, range.poolKey.toId());
+        uint256 liquidityDelta = LiquidityAmounts.getLiquidityForAmounts(
+            sqrtPriceX96,
+            TickMath.getSqrtPriceAtTick(range.tickLower),
+            TickMath.getSqrtPriceAtTick(range.tickUpper),
+            token0Owed,
+            token1Owed
+        );
+
+        vm.prank(alice);
+        lpm.increaseLiquidity(tokenIdAlice, liquidityDelta, ZERO_BYTES, false);
+        snapLastCall("autocompound_exactUnclaimedFees");
+    }
+
+    function test_gas_autocompound_exactUnclaimedFees_exactCustodiedFees() public {
+        // Alice and Bob provide liquidity on the range
+        // Alice uses her fees to increase liquidity. Both unclaimed fees and cached fees are used to exactly increase the liquidity
+        uint256 liquidityAlice = 3_000e18;
+        uint256 liquidityBob = 1_000e18;
+
+        // alice provides liquidity
+        vm.prank(alice);
+        (uint256 tokenIdAlice,) = lpm.mint(range, liquidityAlice, block.timestamp + 1, alice, ZERO_BYTES);
+
+        // bob provides liquidity
+        vm.prank(bob);
+        (uint256 tokenIdBob,) = lpm.mint(range, liquidityBob, block.timestamp + 1, bob, ZERO_BYTES);
+
+        // donate to create fees
+        donateRouter.donate(key, 20e18, 20e18, ZERO_BYTES);
+
+        // bob collects fees so some of alice's fees are now cached
+        vm.prank(bob);
+        lpm.collect(tokenIdBob, bob, ZERO_BYTES, false);
+
+        // donate to create more fees
+        donateRouter.donate(key, 20e18, 20e18, ZERO_BYTES);
+
+        (uint256 newToken0Owed, uint256 newToken1Owed) = lpm.feesOwed(tokenIdAlice);
+
+        // alice will use ALL of her fees to increase liquidity
+        {
+            (uint160 sqrtPriceX96,,,) = StateLibrary.getSlot0(manager, range.poolKey.toId());
+            uint256 liquidityDelta = LiquidityAmounts.getLiquidityForAmounts(
+                sqrtPriceX96,
+                TickMath.getSqrtPriceAtTick(range.tickLower),
+                TickMath.getSqrtPriceAtTick(range.tickUpper),
+                newToken0Owed,
+                newToken1Owed
+            );
+
+            vm.prank(alice);
+            lpm.increaseLiquidity(tokenIdAlice, liquidityDelta, ZERO_BYTES, false);
+            snapLastCall("autocompound_exactUnclaimedFees_exactCustodiedFees");
+        }
+    }
+
+    // autocompounding but the excess fees are credited to tokensOwed
+    function test_gas_autocompound_excessFeesCredit() public {
+        // Alice and Bob provide liquidity on the range
+        // Alice uses her fees to increase liquidity. Excess fees are accounted to alice
+        uint256 liquidityAlice = 3_000e18;
+        uint256 liquidityBob = 1_000e18;
+
+        // alice provides liquidity
+        vm.prank(alice);
+        (uint256 tokenIdAlice,) = lpm.mint(range, liquidityAlice, block.timestamp + 1, alice, ZERO_BYTES);
+
+        // bob provides liquidity
+        vm.prank(bob);
+        (uint256 tokenIdBob,) = lpm.mint(range, liquidityBob, block.timestamp + 1, bob, ZERO_BYTES);
+
+        // donate to create fees
+        donateRouter.donate(key, 20e18, 20e18, ZERO_BYTES);
+
+        // alice will use half of her fees to increase liquidity
+        (uint256 token0Owed, uint256 token1Owed) = lpm.feesOwed(tokenIdAlice);
+
+        (uint160 sqrtPriceX96,,,) = StateLibrary.getSlot0(manager, range.poolKey.toId());
+        uint256 liquidityDelta = LiquidityAmounts.getLiquidityForAmounts(
+            sqrtPriceX96,
+            TickMath.getSqrtPriceAtTick(range.tickLower),
+            TickMath.getSqrtPriceAtTick(range.tickUpper),
+            token0Owed / 2,
+            token1Owed / 2
+        );
+
+        vm.prank(alice);
+        lpm.increaseLiquidity(tokenIdAlice, liquidityDelta, ZERO_BYTES, false);
+        snapLastCall("autocompound_excessFeesCredit");
+    }
+
     function test_gas_decreaseLiquidity_erc20() public {
         (uint256 tokenId,) = lpm.mint(range, 10_000 ether, block.timestamp + 1, address(this), ZERO_BYTES);
 
diff --git a/test/position-managers/IncreaseLiquidity.t.sol b/test/position-managers/IncreaseLiquidity.t.sol
