Optimize TickBitmap.flipTick using assembly (#653)

* Fix gas snapshot tests

* Optimize `TickBitmap.flipTick` using assembly

* Add comment

* Add fuzz test for `flipTick` and update `isInitialized` function

A new fuzz test function `test_fuzz_flipTick` has been added to the TickBitmap library for improved testing accuracy. Additionally, the `isInitialized` function has been re-implemented to enhance its precision and correctness.

* Make `flipTick` test more robust

* Refactor `TickBitmap` test and add a new fuzz test case

A refactoring was done in the TickBitmap test file which includes changing import paths and renaming a test function. A new fuzz test case on testing the next initialized tick within one word, but on an empty bitmap, was also added. This contributes to verifying the accuracy of tick operations under more extensive scenarios.

* Update `tickSpacing` bounds in `TickBitmap` test

.forge-snapshots/addLiquidity CA fee.snap

@@ -1 +1 @@
-331619
+331309

.forge-snapshots/addLiquidity with empty hook.snap

@@ -1 +1 @@
-286612
+286302

.forge-snapshots/create new liquidity to a position with salt.snap

@@ -1 +1 @@
-302131
+301821

.forge-snapshots/flipTick_flippingATickThatResultsInDeletingAWord.snap

@@ -1 +1 @@
-33
+5244

.forge-snapshots/flipTick_flippingFirstTickInWordToInitialized.snap

@@ -1 +1 @@
-22496
+22341

.forge-snapshots/flipTick_flippingSecondTickInWordToInitialized.snap

@@ -1 +1 @@
-5505
+5350

.forge-snapshots/nextInitializedTickWithinOneWord_lteFalse_onBoundary.snap

@@ -1 +1 @@
-33
+2488

.forge-snapshots/poolManager bytecode size.snap

@@ -1 +1 @@
-21796
+21709

.forge-snapshots/removeLiquidity CA fee.snap

@@ -1 +1 @@
-187085
+186775

.forge-snapshots/removeLiquidity with empty hook.snap

@@ -1 +1 @@
-123453
+123143

.forge-snapshots/removeLiquidity with native token.snap

@@ -1 +1 @@
-120240
+119930

.forge-snapshots/removeLiquidity.snap

@@ -0,0 +1 @@
+120637

.forge-snapshots/simple addLiquidity.snap

@@ -1 +1 @@
-167734
+167424

.forge-snapshots/simple removeLiquidity.snap

@@ -1 +1 @@
-90879
+90569

src/libraries/TickBitmap.sol

@@ -43,11 +43,32 @@ library TickBitmap {
     /// @param tick The tick to flip
     /// @param tickSpacing The spacing between usable ticks
     function flipTick(mapping(int16 => uint256) storage self, int24 tick, int24 tickSpacing) internal {
-        unchecked {
-            if (tick % tickSpacing != 0) revert TickMisaligned(tick, tickSpacing); // ensure that the tick is spaced
-            (int16 wordPos, uint8 bitPos) = position(tick / tickSpacing);
-            uint256 mask = 1 << bitPos;
-            self[wordPos] ^= mask;
+        /**
+         * Equivalent to the following Solidity:
+         *     if (tick % tickSpacing != 0) revert TickMisaligned(tick, tickSpacing);
+         *     (int16 wordPos, uint8 bitPos) = position(tick / tickSpacing);
+         *     uint256 mask = 1 << bitPos;
+         *     self[wordPos] ^= mask;
+         */
+        /// @solidity memory-safe-assembly
+        assembly {
+            // ensure that the tick is spaced
+            if smod(tick, tickSpacing) {
+                mstore(0, 0xd4d8f3e6) // selector for TickMisaligned(int24,int24)
+                mstore(0x20, tick)
+                mstore(0x40, tickSpacing)
+                revert(0x1c, 0x44)
+            }
+            tick := sdiv(tick, tickSpacing)
+            // calculate the storage slot corresponding to the tick
+            // wordPos = tick >> 8
+            mstore(0, sar(8, tick))
+            mstore(0x20, self.slot)
+            // the slot of self[wordPos] is keccak256(abi.encode(wordPos, self.slot))
+            let slot := keccak256(0, 0x40)
+            // mask = 1 << bitPos = 1 << (tick % 256)
+            // self[wordPos] ^= mask
+            sstore(slot, xor(sload(slot), shl(and(tick, 0xff), 1)))
         }
     }
 

test/libraries/TickBitmap.t.sol

@@ -4,7 +4,8 @@ pragma solidity ^0.8.20;
 import {Test} from "forge-std/Test.sol";
 import {Vm} from "forge-std/Vm.sol";
 import {GasSnapshot} from "lib/forge-gas-snapshot/src/GasSnapshot.sol";
-import {TickBitmap} from "src/libraries/TickBitmap.sol";
+import {TickBitmap} from "../../src/libraries/TickBitmap.sol";
+import {TickMath} from "../../src/libraries/TickMath.sol";
 
 contract TickBitmapTest is Test, GasSnapshot {
     using TickBitmap for mapping(int16 => uint256);
@@ -14,6 +15,7 @@ contract TickBitmapTest is Test, GasSnapshot {
     int24 constant SOLO_INITIALIZED_TICK_IN_WORD = -10000;
 
     mapping(int16 => uint256) public bitmap;
+    mapping(int16 => uint256) internal emptyBitmap;
 
     function setUp() public {
         // set dirty slots beforehand for certain gas tests
@@ -105,11 +107,27 @@ contract TickBitmapTest is Test, GasSnapshot {
     }
 
     function test_flipTick_flippingATickThatResultsInDeletingAWord_gas() public {
-        flipTick(SOLO_INITIALIZED_TICK_IN_WORD);
         snapStart("flipTick_flippingATickThatResultsInDeletingAWord");
+        flipTick(SOLO_INITIALIZED_TICK_IN_WORD);
         snapEnd();
     }
 
+    function test_fuzz_flipTick(int24 tick, int24 tickSpacing) public {
+        tickSpacing = int24(bound(tickSpacing, 1, type(int24).max));
+
+        if (tick % tickSpacing != 0) {
+            vm.expectRevert(abi.encodeWithSelector(TickBitmap.TickMisaligned.selector, tick, tickSpacing));
+            bitmap.flipTick(tick, tickSpacing);
+        } else {
+            bool initialized = isInitialized(tick, tickSpacing);
+            bitmap.flipTick(tick, tickSpacing);
+            assertEq(isInitialized(tick, tickSpacing), !initialized);
+            // flip again
+            bitmap.flipTick(tick, tickSpacing);
+            assertEq(isInitialized(tick, tickSpacing), initialized);
+        }
+    }
+
     function test_nextInitializedTickWithinOneWord_lteFalse_returnsTickToRightIfAtInitializedTick() public view {
         (int24 next, bool initialized) = bitmap.nextInitializedTickWithinOneWord(78, 1, false);
         assertEq(next, 84);
@@ -180,8 +198,8 @@ contract TickBitmapTest is Test, GasSnapshot {
     }
 
     function test_nextInitializedTickWithinOneWord_lteFalse_onBoundary_gas() public {
-        bitmap.nextInitializedTickWithinOneWord(255, 1, false);
         snapStart("nextInitializedTickWithinOneWord_lteFalse_onBoundary");
+        bitmap.nextInitializedTickWithinOneWord(255, 1, false);
         snapEnd();
     }
 
@@ -276,7 +294,7 @@ contract TickBitmapTest is Test, GasSnapshot {
         snapEnd();
     }
 
-    function test_nextInitializedTickWithinOneWord_fuzz(int24 tick, bool lte) public view {
+    function test_fuzz_nextInitializedTickWithinOneWord(int24 tick, bool lte) public view {
         // assume tick is at least one word inside type(int24).(max | min)
         vm.assume(lte ? tick >= -8388352 : tick < 8388351);
 
@@ -301,9 +319,41 @@ contract TickBitmapTest is Test, GasSnapshot {
         }
     }
 
+    function test_fuzz_nextInitializedTickWithinOneWord_onEmptyBitmap(
+        int24 tick,
+        int24 tickSpacing,
+        uint8 nextBitPos,
+        bool lte
+    ) public {
+        tick = int24(bound(tick, TickMath.MIN_TICK, TickMath.MAX_TICK));
+        tickSpacing = int24(bound(tickSpacing, TickMath.MIN_TICK_SPACING, TickMath.MAX_TICK_SPACING));
+        int24 compressed = TickBitmap.compress(tick, tickSpacing);
+        if (!lte) ++compressed;
+        (int16 wordPos, uint8 bitPos) = TickBitmap.position(compressed);
+
+        if (lte) {
+            nextBitPos = uint8(bound(nextBitPos, 0, bitPos));
+        } else {
+            nextBitPos = uint8(bound(nextBitPos, bitPos, 255));
+        }
+        // Choose the next initialized tick within one word at random and flip it.
+        int24 nextInitializedTick = ((int24(wordPos) << 8) + int24(uint24(nextBitPos))) * tickSpacing;
+        emptyBitmap.flipTick(nextInitializedTick, tickSpacing);
+        (int24 next, bool initialized) = emptyBitmap.nextInitializedTickWithinOneWord(tick, tickSpacing, lte);
+        assertEq(initialized, true);
+        assertEq(next, nextInitializedTick);
+    }
+
+    function isInitialized(int24 tick, int24 tickSpacing) internal view returns (bool) {
+        unchecked {
+            if (tick % tickSpacing != 0) return false;
+            (int16 wordPos, uint8 bitPos) = TickBitmap.position(tick / tickSpacing);
+            return bitmap[wordPos] & (1 << bitPos) != 0;
+        }
+    }
+
     function isInitialized(int24 tick) internal view returns (bool) {
-        (int24 next, bool initialized) = bitmap.nextInitializedTickWithinOneWord(tick, 1, true);
-        return next == tick ? initialized : false;
+        return isInitialized(tick, 1);
     }
 
     function flipTick(int24 tick) internal {