Optimize ProtocolFeeLibrary (#283)

* Refactor import statements and add new gas test

Refactored import statements in ProtocolFeeLibrary.t.sol and added a new test in PoolManager.t.sol to check the gas usage when setting the protocol fee. The changes in import statements are done to follow the file path consistency and the new gas test will help in stress testing to evaluate the gas cost for a particular operation.

* Optimize `ProtocolFeeLibrary.validate`

Refactored and optimized the `validate` function in the `ProtocolFeeLibrary.sol` file. The logic for validating the fee has been simplified and compressed using assembly language. This results in a reduction in bytecode sizes in several snapshots including.poolManager, initialize, and set protocol fee.

* Add gas test for `calculateSwapFee`

This commit introduces a gas test for the `calculateSwapFee` function inside the `ProtocolFeeLibraryTest` contract. With the addition of `GasSnapshot` import, it now tracks the gas consumption of this specific function, helping us to optimize efficiency and performance.

* Refactor `calculateSwapFee` function to use assembly

The `calculateSwapFee` function in the `ProtocolFeeLibrary.sol` file has been rewritten to use assembly for performance optimization. This improves efficiency by reducing gas usage, reflected in the bytecode size reduction, which is captured by the snapshot changes. Operating at a lower level with assembly allows us to manage resources more directly.

* Change visibility of protocol fee thresholds

* Optimize `isValidProtocolFee` further

* Apply review comments

.forge-snapshots/calculateSwapFee.snap

@@ -0,0 +1 @@
+33

.forge-snapshots/initialize.snap

@@ -1 +1 @@
-61902
+61862

.forge-snapshots/poolManager bytecode size.snap

@@ -1 +1 @@
-21934
+21878

.forge-snapshots/set protocol fee.snap

@@ -0,0 +1 @@
+32415

.forge-snapshots/swap with lp fee and protocol fee.snap

@@ -1 +1 @@
-180456
+180453

src/libraries/ProtocolFeeLibrary.sol

@@ -1,42 +1,43 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 pragma solidity ^0.8.20;
 
-import "./UnsafeMath.sol";
-
 library ProtocolFeeLibrary {
     // Max protocol fee is 0.1% (1000 pips)
     uint16 public constant MAX_PROTOCOL_FEE = 1000;
 
+    // Thresholds used for optimized bounds checks on protocol fees
+    uint24 internal constant FEE_0_THRESHOLD = 1001;
+    uint24 internal constant FEE_1_THRESHOLD = 1001 << 12;
+
     // the protocol fee is represented in hundredths of a bip
     uint256 internal constant PIPS_DENOMINATOR = 1_000_000;
 
     function getZeroForOneFee(uint24 self) internal pure returns (uint16) {
-        return uint16(self & (4096 - 1));
+        return uint16(self & 0xfff);
     }
 
     function getOneForZeroFee(uint24 self) internal pure returns (uint16) {
         return uint16(self >> 12);
     }
 
-    function isValidProtocolFee(uint24 self) internal pure returns (bool) {
-        if (self != 0) {
-            uint16 fee0 = getZeroForOneFee(self);
-            uint16 fee1 = getOneForZeroFee(self);
-            // The fee is represented in pips and it cannot be greater than the MAX_PROTOCOL_FEE.
-            if ((fee0 > MAX_PROTOCOL_FEE) || (fee1 > MAX_PROTOCOL_FEE)) {
-                return false;
-            }
+    function isValidProtocolFee(uint24 self) internal pure returns (bool valid) {
+        // Equivalent to: getZeroForOneFee(self) <= MAX_PROTOCOL_FEE && getOneForZeroFee(self) <= MAX_PROTOCOL_FEE
+        assembly {
+            let isZeroForOneFeeOk := lt(and(self, 0xfff), FEE_0_THRESHOLD)
+            let isOneForZeroFeeOk := lt(self, FEE_1_THRESHOLD)
+            valid := and(isZeroForOneFeeOk, isOneForZeroFeeOk)
         }
-        return true;
     }
 
     // The protocol fee is taken from the input amount first and then the LP fee is taken from the remaining
     // The swap fee is capped at 100%
     // Equivalent to protocolFee + lpFee(1_000_000 - protocolFee) / 1_000_000
-    function calculateSwapFee(uint24 self, uint24 lpFee) internal pure returns (uint24) {
-        unchecked {
-            uint256 numerator = uint256(self) * uint256(lpFee);
-            return uint24(uint256(self) + lpFee - UnsafeMath.divRoundingUp(numerator, PIPS_DENOMINATOR));
+    function calculateSwapFee(uint24 self, uint24 lpFee) internal pure returns (uint24 swapFee) {
+        // protocolFee + lpFee - (protocolFee * lpFee / 1_000_000). Div rounds up to favor LPs over the protocol.
+        assembly {
+            let numerator := mul(self, lpFee)
+            let divRoundingUp := add(div(numerator, PIPS_DENOMINATOR), gt(mod(numerator, PIPS_DENOMINATOR), 0))
+            swapFee := sub(add(self, lpFee), divRoundingUp)
         }
     }
 }

test/PoolManager.t.sol

@@ -1282,6 +1282,12 @@ contract PoolManagerTest is Test, Deployers, GasSnapshot {
         manager.burn(address(this), key.currency0.toId(), 1);
     }
 
+    function test_setProtocolFee_gas() public {
+        vm.prank(address(feeController));
+        manager.setProtocolFee(key, MAX_PROTOCOL_FEE_BOTH_TOKENS);
+        snapLastCall("set protocol fee");
+    }
+
     function test_setProtocolFee_updatesProtocolFeeForInitializedPool(uint24 protocolFee) public {
         (,, uint24 slot0ProtocolFee,) = manager.getSlot0(key.toId());
         assertEq(slot0ProtocolFee, 0);

test/libraries/ProtocolFeeLibrary.t.sol

@@ -1,11 +1,12 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.20;
 
-import "../../src/libraries/ProtocolFeeLibrary.sol";
-import "../../src/libraries/LPFeeLibrary.sol";
-import "forge-std/Test.sol";
+import {Test} from "forge-std/Test.sol";
+import {GasSnapshot} from "forge-gas-snapshot/GasSnapshot.sol";
+import {LPFeeLibrary} from "../../src/libraries/LPFeeLibrary.sol";
+import {ProtocolFeeLibrary} from "../../src/libraries/ProtocolFeeLibrary.sol";
 
-contract ProtocolFeeLibraryTest is Test {
+contract ProtocolFeeLibraryTest is Test, GasSnapshot {
     function test_getZeroForOneFee() public pure {
         uint24 fee = uint24(ProtocolFeeLibrary.MAX_PROTOCOL_FEE - 1) << 12 | ProtocolFeeLibrary.MAX_PROTOCOL_FEE;
         assertEq(ProtocolFeeLibrary.getZeroForOneFee(fee), uint24(ProtocolFeeLibrary.MAX_PROTOCOL_FEE));
@@ -80,4 +81,12 @@ contract ProtocolFeeLibraryTest is Test {
             protocolFee + lpFee * uint256(LPFeeLibrary.MAX_LP_FEE - protocolFee) / LPFeeLibrary.MAX_LP_FEE;
         assertEq(swapFee, uint24(expectedSwapFee));
     }
+
+    function test_calculateSwapFee_gas() public {
+        uint24 self = uint24(ProtocolFeeLibrary.MAX_PROTOCOL_FEE);
+        uint24 lpFee = LPFeeLibrary.MAX_LP_FEE;
+        snapStart("calculateSwapFee");
+        ProtocolFeeLibrary.calculateSwapFee(self, lpFee);
+        snapEnd();
+    }
 }