Refactor and optimizations

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@aperture_finance/uni-v3-lib",
   "description": "A suite of Solidity libraries that have been imported and rewritten from Uniswap's v3-core and v3-periphery",
-  "version": "2.0.3",
+  "version": "2.1.0",
   "author": "Aperture Finance",
   "homepage": "https://aperture.finance/",
   "license": "GPL-2.0-or-later",
@@ -26,8 +26,8 @@
     "build": "forge build",
     "clean": "forge clean",
     "test": "forge build && forge test",
-    "snapshot": "forge build && forge snapshot --via-ir --evm-version cancun --isolate",
-    "snapshot:diff": "forge build && forge snapshot --diff --via-ir --evm-version cancun --isolate",
+    "snapshot": "forge build && forge snapshot --isolate",
+    "snapshot:diff": "forge build && forge snapshot --diff --isolate",
     "prettier": "prettier -c src/*.sol {src,test}/**/*.sol",
     "prettier:fix": "prettier -w src/*.sol {src,test}/**/*.sol"
   },

src/FullMath.sol

@@ -33,18 +33,21 @@ library FullMath {
         return FixedPointMathLib.fullMulDivUp(a, b, denominator);
     }
 
-    /// @notice Calculates x * y / 2^96 with full precision.
-    function mulDiv96(uint256 x, uint256 y) internal pure returns (uint256 result) {
+    /// @notice Calculates a * b / 2^96 with full precision.
+    /// @param a The multiplicand
+    /// @param b The multiplier
+    /// @return result The 256-bit result
+    function mulDivQ96(uint256 a, uint256 b) internal pure returns (uint256 result) {
         assembly ("memory-safe") {
-            // 512-bit multiply `[prod1 prod0] = x * y`.
+            // 512-bit multiply `[prod1 prod0] = a * b`.
             // Compute the product mod `2**256` and mod `2**256 - 1`
             // then use the Chinese Remainder Theorem to reconstruct
             // the 512 bit result. The result is stored in two 256
             // variables such that `product = prod1 * 2**256 + prod0`.
 
             // Least significant 256 bits of the product.
-            let prod0 := mul(x, y)
-            let mm := mulmod(x, y, not(0))
+            let prod0 := mul(a, b)
+            let mm := mulmod(a, b, not(0))
             // Most significant 256 bits of the product.
             let prod1 := sub(mm, add(prod0, lt(mm, prod0)))
 
@@ -61,6 +64,37 @@ library FullMath {
         }
     }
 
+    /// @notice Calculates a * b / 2^128 with full precision.
+    /// @param a The multiplicand
+    /// @param b The multiplier
+    /// @return result The 256-bit result
+    function mulDivQ128(uint256 a, uint256 b) internal pure returns (uint256 result) {
+        assembly ("memory-safe") {
+            // 512-bit multiply `[prod1 prod0] = a * b`.
+            // Compute the product mod `2**256` and mod `2**256 - 1`
+            // then use the Chinese Remainder Theorem to reconstruct
+            // the 512 bit result. The result is stored in two 256
+            // variables such that `product = prod1 * 2**256 + prod0`.
+
+            // Least significant 256 bits of the product.
+            let prod0 := mul(a, b)
+            let mm := mulmod(a, b, not(0))
+            // Most significant 256 bits of the product.
+            let prod1 := sub(mm, add(prod0, lt(mm, prod0)))
+
+            // Make sure the result is less than `2**256`.
+            if iszero(gt(0x100000000000000000000000000000000, prod1)) {
+                // Store the function selector of `FullMulDivFailed()`.
+                mstore(0x00, 0xae47f702)
+                // Revert with (offset, size).
+                revert(0x1c, 0x04)
+            }
+
+            // Divide [prod1 prod0] by 2^128.
+            result := or(shr(128, prod0), shl(128, prod1))
+        }
+    }
+
     /// @dev Returns the square root of a number. If the number is not a perfect square, the value is rounded down.
     function sqrt(uint256 x) internal pure returns (uint256) {
         return FixedPointMathLib.sqrt(x);

src/LiquidityAmounts.sol

@@ -26,9 +26,25 @@ library LiquidityAmounts {
         uint160 sqrtRatioBX96,
         uint256 amount0
     ) internal pure returns (uint128 liquidity) {
-        (sqrtRatioAX96, sqrtRatioBX96) = TernaryLib.sort2(sqrtRatioAX96, sqrtRatioBX96);
-        uint256 intermediate = FullMath.mulDiv96(sqrtRatioAX96, sqrtRatioBX96);
-        return FullMath.mulDiv(amount0, intermediate, sqrtRatioBX96.sub(sqrtRatioAX96)).toUint128();
+        uint256 intermediate = FullMath.mulDivQ96(sqrtRatioAX96, sqrtRatioBX96);
+        return FullMath.mulDiv(amount0, intermediate, TernaryLib.absDiffU160(sqrtRatioAX96, sqrtRatioBX96)).toUint128();
+    }
+
+    /// @notice Computes the amount of liquidity received for a given amount of token0 and price range
+    /// @dev Calculates amount0 * (sqrt(upper) * sqrt(lower)) / (sqrt(upper) - sqrt(lower))
+    /// @param sqrtRatioAX96 A sqrt price representing the lower tick boundary
+    /// @param sqrtRatioBX96 A sqrt price representing the upper tick boundary
+    /// @param amount0 The amount0 being sent in
+    /// @return liquidity The amount of returned liquidity
+    function getLiquidityForAmount0Sorted(
+        uint160 sqrtRatioAX96,
+        uint160 sqrtRatioBX96,
+        uint256 amount0
+    ) internal pure returns (uint128 liquidity) {
+        unchecked {
+            uint256 intermediate = FullMath.mulDivQ96(sqrtRatioAX96, sqrtRatioBX96);
+            return FullMath.mulDiv(amount0, intermediate, sqrtRatioBX96 - sqrtRatioAX96).toUint128();
+        }
     }
 
     /// @notice Computes the amount of liquidity received for a given amount of token1 and price range
@@ -42,8 +58,26 @@ library LiquidityAmounts {
         uint160 sqrtRatioBX96,
         uint256 amount1
     ) internal pure returns (uint128 liquidity) {
-        (sqrtRatioAX96, sqrtRatioBX96) = TernaryLib.sort2(sqrtRatioAX96, sqrtRatioBX96);
-        return FullMath.mulDiv(amount1, FixedPoint96.Q96, sqrtRatioBX96.sub(sqrtRatioAX96)).toUint128();
+        return
+            FullMath
+                .mulDiv(amount1, FixedPoint96.Q96, TernaryLib.absDiffU160(sqrtRatioAX96, sqrtRatioBX96))
+                .toUint128();
+    }
+
+    /// @notice Computes the amount of liquidity received for a given amount of token1 and price range
+    /// @dev Calculates amount1 / (sqrt(upper) - sqrt(lower)).
+    /// @param sqrtRatioAX96 A sqrt price representing the lower tick boundary
+    /// @param sqrtRatioBX96 A sqrt price representing the upper tick boundary
+    /// @param amount1 The amount1 being sent in
+    /// @return liquidity The amount of returned liquidity
+    function getLiquidityForAmount1Sorted(
+        uint160 sqrtRatioAX96,
+        uint160 sqrtRatioBX96,
+        uint256 amount1
+    ) internal pure returns (uint128 liquidity) {
+        unchecked {
+            return FullMath.mulDiv(amount1, FixedPoint96.Q96, sqrtRatioBX96 - sqrtRatioAX96).toUint128();
+        }
     }
 
     /// @notice Computes the maximum amount of liquidity received for a given amount of token0, token1, the current
@@ -61,19 +95,19 @@ library LiquidityAmounts {
         uint256 amount0,
         uint256 amount1
     ) internal pure returns (uint128 liquidity) {
-        (sqrtRatioAX96, sqrtRatioBX96) = TernaryLib.sort2(sqrtRatioAX96, sqrtRatioBX96);
+        (sqrtRatioAX96, sqrtRatioBX96) = TernaryLib.sort2U160(sqrtRatioAX96, sqrtRatioBX96);
 
         if (sqrtRatioX96 <= sqrtRatioAX96) {
-            liquidity = getLiquidityForAmount0(sqrtRatioAX96, sqrtRatioBX96, amount0);
+            liquidity = getLiquidityForAmount0Sorted(sqrtRatioAX96, sqrtRatioBX96, amount0);
         } else if (sqrtRatioX96 < sqrtRatioBX96) {
-            uint128 liquidity0 = getLiquidityForAmount0(sqrtRatioX96, sqrtRatioBX96, amount0);
-            uint128 liquidity1 = getLiquidityForAmount1(sqrtRatioAX96, sqrtRatioX96, amount1);
+            uint128 liquidity0 = getLiquidityForAmount0Sorted(sqrtRatioX96, sqrtRatioBX96, amount0);
+            uint128 liquidity1 = getLiquidityForAmount1Sorted(sqrtRatioAX96, sqrtRatioX96, amount1);
             // liquidity = min(liquidity0, liquidity1);
             assembly {
                 liquidity := xor(liquidity0, mul(xor(liquidity0, liquidity1), lt(liquidity1, liquidity0)))
             }
         } else {
-            liquidity = getLiquidityForAmount1(sqrtRatioAX96, sqrtRatioBX96, amount1);
+            liquidity = getLiquidityForAmount1Sorted(sqrtRatioAX96, sqrtRatioBX96, amount1);
         }
     }
 
@@ -87,11 +121,31 @@ library LiquidityAmounts {
         uint160 sqrtRatioBX96,
         uint128 liquidity
     ) internal pure returns (uint256 amount0) {
-        (sqrtRatioAX96, sqrtRatioBX96) = TernaryLib.sort2(sqrtRatioAX96, sqrtRatioBX96);
-        return
-            FullMath
-                .mulDiv(uint256(liquidity) << FixedPoint96.RESOLUTION, sqrtRatioBX96.sub(sqrtRatioAX96), sqrtRatioBX96)
-                .div(sqrtRatioAX96);
+        unchecked {
+            (sqrtRatioAX96, sqrtRatioBX96) = TernaryLib.sort2U160(sqrtRatioAX96, sqrtRatioBX96);
+            return
+                FullMath
+                    .mulDiv(uint256(liquidity) << FixedPoint96.RESOLUTION, sqrtRatioBX96 - sqrtRatioAX96, sqrtRatioBX96)
+                    .div(sqrtRatioAX96);
+        }
+    }
+
+    /// @notice Computes the amount of token0 for a given amount of liquidity and a price range
+    /// @param sqrtRatioAX96 A sqrt price representing the lower tick boundary
+    /// @param sqrtRatioBX96 A sqrt price representing the upper tick boundary
+    /// @param liquidity The liquidity being valued
+    /// @return amount0 The amount of token0
+    function getAmount0ForLiquiditySorted(
+        uint160 sqrtRatioAX96,
+        uint160 sqrtRatioBX96,
+        uint128 liquidity
+    ) internal pure returns (uint256 amount0) {
+        unchecked {
+            return
+                FullMath
+                    .mulDiv(uint256(liquidity) << FixedPoint96.RESOLUTION, sqrtRatioBX96 - sqrtRatioAX96, sqrtRatioBX96)
+                    .div(sqrtRatioAX96);
+        }
     }
 
     /// @notice Computes the amount of token1 for a given amount of liquidity and a price range
@@ -104,8 +158,22 @@ library LiquidityAmounts {
         uint160 sqrtRatioBX96,
         uint128 liquidity
     ) internal pure returns (uint256 amount1) {
-        (sqrtRatioAX96, sqrtRatioBX96) = TernaryLib.sort2(sqrtRatioAX96, sqrtRatioBX96);
-        return FullMath.mulDiv96(liquidity, sqrtRatioBX96.sub(sqrtRatioAX96));
+        return FullMath.mulDivQ96(liquidity, TernaryLib.absDiffU160(sqrtRatioAX96, sqrtRatioBX96));
+    }
+
+    /// @notice Computes the amount of token1 for a given amount of liquidity and a price range
+    /// @param sqrtRatioAX96 A sqrt price representing the lower tick boundary
+    /// @param sqrtRatioBX96 A sqrt price representing the upper tick boundary
+    /// @param liquidity The liquidity being valued
+    /// @return amount1 The amount of token1
+    function getAmount1ForLiquiditySorted(
+        uint160 sqrtRatioAX96,
+        uint160 sqrtRatioBX96,
+        uint128 liquidity
+    ) internal pure returns (uint256 amount1) {
+        unchecked {
+            return FullMath.mulDivQ96(liquidity, sqrtRatioBX96 - sqrtRatioAX96);
+        }
     }
 
     /// @notice Computes the token0 and token1 value for a given amount of liquidity, the current
@@ -122,15 +190,15 @@ library LiquidityAmounts {
         uint160 sqrtRatioBX96,
         uint128 liquidity
     ) internal pure returns (uint256 amount0, uint256 amount1) {
-        (sqrtRatioAX96, sqrtRatioBX96) = TernaryLib.sort2(sqrtRatioAX96, sqrtRatioBX96);
+        (sqrtRatioAX96, sqrtRatioBX96) = TernaryLib.sort2U160(sqrtRatioAX96, sqrtRatioBX96);
 
         if (sqrtRatioX96 <= sqrtRatioAX96) {
-            amount0 = getAmount0ForLiquidity(sqrtRatioAX96, sqrtRatioBX96, liquidity);
-        } else if (sqrtRatioX96 < sqrtRatioBX96) {
-            amount0 = getAmount0ForLiquidity(sqrtRatioX96, sqrtRatioBX96, liquidity);
-            amount1 = getAmount1ForLiquidity(sqrtRatioAX96, sqrtRatioX96, liquidity);
+            amount0 = getAmount0ForLiquiditySorted(sqrtRatioAX96, sqrtRatioBX96, liquidity);
+        } else if (sqrtRatioX96 <= sqrtRatioBX96) {
+            amount0 = getAmount0ForLiquiditySorted(sqrtRatioX96, sqrtRatioBX96, liquidity);
+            amount1 = getAmount1ForLiquiditySorted(sqrtRatioAX96, sqrtRatioX96, liquidity);
         } else {
-            amount1 = getAmount1ForLiquidity(sqrtRatioAX96, sqrtRatioBX96, liquidity);
+            amount1 = getAmount1ForLiquiditySorted(sqrtRatioAX96, sqrtRatioBX96, liquidity);
         }
     }
 }

src/LiquidityMath.sol

@@ -8,7 +8,7 @@ library LiquidityMath {
     /// @notice Add a signed liquidity delta to liquidity and revert if it overflows or underflows
     /// @param x The liquidity before change
     /// @param y The delta by which liquidity should be changed
-    /// @return z The liquidity delta
+    /// @return z The liquidity after
     function addDelta(uint128 x, int128 y) internal pure returns (uint128 z) {
         /// @solidity memory-safe-assembly
         assembly {

src/PoolCaller.sol

@@ -321,6 +321,7 @@ library PoolCaller {
             mstore(add(fmp, 0x84), 0xa0)
             // length = data.length + 32
             let length := add(mload(data), 0x20)
+            // TODO: Use `mcopy`
             // Call the identity precompile 0x04 to copy `data` into calldata.
             pop(staticcall(gas(), 0x04, data, length, add(fmp, 0xa4), length))
             // We use `196 + data.length` for the length of our calldata.

src/SafeCast.sol

@@ -51,7 +51,7 @@ library SafeCast {
     function toInt256(uint256 y) internal pure returns (int256 z) {
         /// @solidity memory-safe-assembly
         assembly {
-            if shr(255, y) {
+            if slt(y, 0) {
                 revert(0, 0)
             }
             z := y