Refactor and optimize SwapMath (#280)

* Refactor and optimize `SwapMath` * Update `SwapMath` calculations and gas snapshots, change `test_entireInputAmountTakenAsFee` * No assembly in `computeSwapStep` Replace inline assembly with native Solidity for `amountRemainingAbs` calculation in `computeSwapStep` function. The previous implementation was not intuitive and easy to read. This change simplifies the code and improves readability while maintaining the exact functionality. * Avoid casting `feePips` multiple times In `SwapMath.sol`, we reduced redundancy by storing the value of `feePips` into `_feePips` and using it throughout the function instead of casting `feePips` multiple times. This refactoring improved efficiency and reduced bytecode size, leading to reduced gas cost for function calls. * Fix natspec * Rename `sqrtRatio` to `sqrtPrice` in SwapMath and Pool libraries * Refactor variable declaration in `SwapMath` The declaration and assignment of `amountRemainingAbs` has been moved inside the relevant if-else blocks in the SwapMath.sol file. This change makes the code more concise and clear by limiting the scope and use of `amountRemainingAbs` to where it is needed. * Update test function to pure in SwapMath.t.sol * Optimize `getSqrtPriceTarget` * Switch if else blocks for ease of review * Refactor and reduce diffs * Recompute `amountIn` when it's capped by `amountRemaining` instead of the target price for `exactIn` to achieve equivalence with V3 * Remove `amountRemainingAbs` * Remove .gas-snapshot * Address review comment * fmt
Uniswap · May 26, 2024 · d50bfc5 · d50bfc5
1 parent 0d3a9a7
commit d50bfc5
Show file tree

Hide file tree

Showing 25 changed files with 50 additions and 58 deletions.
diff --git a/.forge-snapshots/SwapMath_oneForZero_exactInCapped.snap b/.forge-snapshots/SwapMath_oneForZero_exactInCapped.snap
@@ -1 +1 @@
-1947
+1662
diff --git a/.forge-snapshots/SwapMath_oneForZero_exactInPartial.snap b/.forge-snapshots/SwapMath_oneForZero_exactInPartial.snap
@@ -1 +1 @@
-2390
+2152
diff --git a/.forge-snapshots/SwapMath_oneForZero_exactOutCapped.snap b/.forge-snapshots/SwapMath_oneForZero_exactOutCapped.snap
@@ -1 +1 @@
-1686
+1413
diff --git a/.forge-snapshots/SwapMath_oneForZero_exactOutPartial.snap b/.forge-snapshots/SwapMath_oneForZero_exactOutPartial.snap
@@ -1 +1 @@
-2831
+2120
diff --git a/.forge-snapshots/SwapMath_zeroForOne_exactInCapped.snap b/.forge-snapshots/SwapMath_zeroForOne_exactInCapped.snap
@@ -1 +1 @@
-2059
+1753
diff --git a/.forge-snapshots/SwapMath_zeroForOne_exactInPartial.snap b/.forge-snapshots/SwapMath_zeroForOne_exactInPartial.snap
@@ -1 +1 @@
-2941
+2671
diff --git a/.forge-snapshots/SwapMath_zeroForOne_exactOutCapped.snap b/.forge-snapshots/SwapMath_zeroForOne_exactOutCapped.snap
@@ -1 +1 @@
-1798
+1504
diff --git a/.forge-snapshots/SwapMath_zeroForOne_exactOutPartial.snap b/.forge-snapshots/SwapMath_zeroForOne_exactOutPartial.snap
@@ -1 +1 @@
-2458
+1850
diff --git a/.forge-snapshots/poolManager bytecode size.snap b/.forge-snapshots/poolManager bytecode size.snap
@@ -1 +1 @@
-19986
+19871
diff --git a/.forge-snapshots/simple swap with native.snap b/.forge-snapshots/simple swap with native.snap
@@ -1 +1 @@
-114924
+114348
diff --git a/.forge-snapshots/simple swap.snap b/.forge-snapshots/simple swap.snap
@@ -1 +1 @@
-130089
+129513
diff --git a/.forge-snapshots/swap CA fee on unspecified.snap b/.forge-snapshots/swap CA fee on unspecified.snap
@@ -1 +1 @@
-179527
+178951
diff --git a/.forge-snapshots/swap against liquidity with native token.snap b/.forge-snapshots/swap against liquidity with native token.snap
@@ -1 +1 @@
-110279
+110009
diff --git a/.forge-snapshots/swap against liquidity.snap b/.forge-snapshots/swap against liquidity.snap
@@ -1 +1 @@
-121626
+121356
diff --git a/.forge-snapshots/swap burn 6909 for input.snap b/.forge-snapshots/swap burn 6909 for input.snap
@@ -1 +1 @@
-133455
+132744
diff --git a/.forge-snapshots/swap burn native 6909 for input.snap b/.forge-snapshots/swap burn native 6909 for input.snap
@@ -1 +1 @@
-122743
+122135
diff --git a/.forge-snapshots/swap mint native output as 6909.snap b/.forge-snapshots/swap mint native output as 6909.snap
@@ -1 +1 @@
-144531
+144293
diff --git a/.forge-snapshots/swap mint output as 6909.snap b/.forge-snapshots/swap mint output as 6909.snap
@@ -1 +1 @@
-161161
+160585
diff --git a/.forge-snapshots/swap skips hook call if hook is caller.snap b/.forge-snapshots/swap skips hook call if hook is caller.snap
@@ -1 +1 @@
-217571
+216725
diff --git a/.forge-snapshots/swap with dynamic fee.snap b/.forge-snapshots/swap with dynamic fee.snap
@@ -1 +1 @@
-145362
+144786
diff --git a/.forge-snapshots/swap with hooks.snap b/.forge-snapshots/swap with hooks.snap
@@ -1 +1 @@
-138626
+138356
diff --git a/.forge-snapshots/swap with lp fee and protocol fee.snap b/.forge-snapshots/swap with lp fee and protocol fee.snap
@@ -1 +1 @@
-176624
+175722
diff --git a/.forge-snapshots/swap with return dynamic fee.snap b/.forge-snapshots/swap with return dynamic fee.snap
@@ -1 +1 @@
-152398
+151822
diff --git a/.forge-snapshots/update dynamic fee in before swap.snap b/.forge-snapshots/update dynamic fee in before swap.snap
@@ -1 +1 @@
-155040
+154464
diff --git a/src/libraries/SwapMath.sol b/src/libraries/SwapMath.sol
@@ -7,6 +7,8 @@ import {SqrtPriceMath} from "./SqrtPriceMath.sol";
 /// @title Computes the result of a swap within ticks
 /// @notice Contains methods for computing the result of a swap within a single tick price range, i.e., a single tick.
 library SwapMath {
+    uint256 internal constant MAX_FEE_PIPS = 1e6;
+
     /// @notice Computes the sqrt price target for the next swap step
     /// @param zeroForOne The direction of the swap, true for currency0 to currency1, false for currency1 to currency0
     /// @param sqrtPriceNextX96 The Q64.96 sqrt price for the next initialized tick
@@ -49,63 +51,53 @@ library SwapMath {
         uint24 feePips
     ) internal pure returns (uint160 sqrtPriceNextX96, uint256 amountIn, uint256 amountOut, uint256 feeAmount) {
         unchecked {
+            uint256 _feePips = feePips; // upcast once and cache
             bool zeroForOne = sqrtPriceCurrentX96 >= sqrtPriceTargetX96;
             bool exactIn = amountRemaining < 0;
 
             if (exactIn) {
-                uint256 amountRemainingLessFee = FullMath.mulDiv(uint256(-amountRemaining), 1e6 - feePips, 1e6);
+                uint256 amountRemainingLessFee =
+                    FullMath.mulDiv(uint256(-amountRemaining), MAX_FEE_PIPS - _feePips, MAX_FEE_PIPS);
                 amountIn = zeroForOne
                     ? SqrtPriceMath.getAmount0Delta(sqrtPriceTargetX96, sqrtPriceCurrentX96, liquidity, true)
                     : SqrtPriceMath.getAmount1Delta(sqrtPriceCurrentX96, sqrtPriceTargetX96, liquidity, true);
                 if (amountRemainingLessFee >= amountIn) {
+                    // `amountIn` is capped by the target price
                     sqrtPriceNextX96 = sqrtPriceTargetX96;
+                    feeAmount = _feePips == MAX_FEE_PIPS
+                        ? amountIn
+                        : FullMath.mulDivRoundingUp(amountIn, _feePips, MAX_FEE_PIPS - _feePips);
                 } else {
                     sqrtPriceNextX96 = SqrtPriceMath.getNextSqrtPriceFromInput(
                         sqrtPriceCurrentX96, liquidity, amountRemainingLessFee, zeroForOne
                     );
+                    amountIn = zeroForOne
+                        ? SqrtPriceMath.getAmount0Delta(sqrtPriceNextX96, sqrtPriceCurrentX96, liquidity, true)
+                        : SqrtPriceMath.getAmount1Delta(sqrtPriceCurrentX96, sqrtPriceNextX96, liquidity, true);
+                    // we didn't reach the target, so take the remainder of the maximum input as fee
+                    feeAmount = uint256(-amountRemaining) - amountIn;
                 }
+                amountOut = zeroForOne
+                    ? SqrtPriceMath.getAmount1Delta(sqrtPriceNextX96, sqrtPriceCurrentX96, liquidity, false)
+                    : SqrtPriceMath.getAmount0Delta(sqrtPriceCurrentX96, sqrtPriceNextX96, liquidity, false);
             } else {
                 amountOut = zeroForOne
                     ? SqrtPriceMath.getAmount1Delta(sqrtPriceTargetX96, sqrtPriceCurrentX96, liquidity, false)
                     : SqrtPriceMath.getAmount0Delta(sqrtPriceCurrentX96, sqrtPriceTargetX96, liquidity, false);
                 if (uint256(amountRemaining) >= amountOut) {
+                    // `amountOut` is capped by the target price
                     sqrtPriceNextX96 = sqrtPriceTargetX96;
                 } else {
-                    sqrtPriceNextX96 = SqrtPriceMath.getNextSqrtPriceFromOutput(
-                        sqrtPriceCurrentX96, liquidity, uint256(amountRemaining), zeroForOne
-                    );
+                    // cap the output amount to not exceed the remaining output amount
+                    amountOut = uint256(amountRemaining);
+                    sqrtPriceNextX96 =
+                        SqrtPriceMath.getNextSqrtPriceFromOutput(sqrtPriceCurrentX96, liquidity, amountOut, zeroForOne);
                 }
-            }
-
-            bool max = sqrtPriceTargetX96 == sqrtPriceNextX96;
-
-            // get the input/output amounts
-            if (zeroForOne) {
-                amountIn = max && exactIn
-                    ? amountIn
-                    : SqrtPriceMath.getAmount0Delta(sqrtPriceNextX96, sqrtPriceCurrentX96, liquidity, true);
-                amountOut = max && !exactIn
-                    ? amountOut
-                    : SqrtPriceMath.getAmount1Delta(sqrtPriceNextX96, sqrtPriceCurrentX96, liquidity, false);
-            } else {
-                amountIn = max && exactIn
-                    ? amountIn
+                amountIn = zeroForOne
+                    ? SqrtPriceMath.getAmount0Delta(sqrtPriceNextX96, sqrtPriceCurrentX96, liquidity, true)
                     : SqrtPriceMath.getAmount1Delta(sqrtPriceCurrentX96, sqrtPriceNextX96, liquidity, true);
-                amountOut = max && !exactIn
-                    ? amountOut
-                    : SqrtPriceMath.getAmount0Delta(sqrtPriceCurrentX96, sqrtPriceNextX96, liquidity, false);
-            }
-
-            // cap the output amount to not exceed the remaining output amount
-            if (!exactIn && amountOut > uint256(amountRemaining)) {
-                amountOut = uint256(amountRemaining);
-            }
-
-            if (exactIn && sqrtPriceNextX96 != sqrtPriceTargetX96) {
-                // we didn't reach the target, so take the remainder of the maximum input as fee
-                feeAmount = uint256(-amountRemaining) - amountIn;
-            } else {
-                feeAmount = feePips == 1e6 ? amountIn : FullMath.mulDivRoundingUp(amountIn, feePips, 1e6 - feePips);
+                // `feePips` cannot be `MAX_FEE_PIPS` for exact out
+                feeAmount = FullMath.mulDivRoundingUp(amountIn, _feePips, MAX_FEE_PIPS - _feePips);
             }
         }
     }