From 8c36e2fe38c5fd41dc1d88b1c3454661c547447a Mon Sep 17 00:00:00 2001 From: Shuhui Luo <107524008+shuhuiluo@users.noreply.github.com> Date: Fri, 5 Apr 2024 15:58:17 -0400 Subject: [PATCH] Refactor calculation and assembly block in `TickBitmap` This commit refactors the calculation of mask and the arrangement of assembly block in TickBitmap.sol for clarity and easier maintainability. Now, the bitPos utilizes a consistent hex representation across the file. The bitPos overflow scenario is handled taking into consideration that 1 << 256 equals 0. --- .gas-snapshot | 14 ++++++------ src/SqrtPriceMath.sol | 2 +- src/TickBitmap.sol | 50 +++++++++++++++++++++---------------------- 3 files changed, 33 insertions(+), 33 deletions(-) diff --git a/.gas-snapshot b/.gas-snapshot index af8cc8c..bd41a05 100644 --- a/.gas-snapshot +++ b/.gas-snapshot @@ -154,24 +154,24 @@ SwapMathTest:testGas_ComputeSwapStepExactOut() (gas: 370755) SwapMathTest:testGas_ComputeSwapStepExactOut_Og() (gas: 510266) SwapMathTest:testGas_ComputeSwapStep_Og() (gas: 526558) TickBitmapPCSTest:testFuzz_Compress(int24,int24) (runs: 65540, μ: 8624, ~: 8658) -TickBitmapPCSTest:testFuzz_FlipTick(int24) (runs: 65540, μ: 111577, ~: 111849) -TickBitmapPCSTest:testFuzz_NextInitializedTickWithinOneWord(int24,uint8,bool) (runs: 65540, μ: 111543, ~: 111448) +TickBitmapPCSTest:testFuzz_FlipTick(int24) (runs: 65540, μ: 111578, ~: 111849) +TickBitmapPCSTest:testFuzz_NextInitializedTickWithinOneWord(int24,uint8,bool) (runs: 65540, μ: 111540, ~: 111446) TickBitmapPCSTest:testFuzz_Position(int24) (runs: 65540, μ: 3714, ~: 3714) TickBitmapPCSTest:testGas_NextInitializedTickWithinOneWord() (gas: 12286448) TickBitmapPCSTest:testGas_NextInitializedTickWithinOneWord_Og() (gas: 12567608) -TickBitmapPCSTest:test_NextInitializedTickWithinOneWord_GT() (gas: 105811) +TickBitmapPCSTest:test_NextInitializedTickWithinOneWord_GT() (gas: 105771) TickBitmapPCSTest:test_NextInitializedTickWithinOneWord_LTE() (gas: 105789) -TickBitmapPCSTest:test_NextInitializedTick_GT() (gas: 114408) +TickBitmapPCSTest:test_NextInitializedTick_GT() (gas: 114328) TickBitmapPCSTest:test_NextInitializedTick_LTE() (gas: 115996) TickBitmapTest:testFuzz_Compress(int24,int24) (runs: 65540, μ: 8624, ~: 8669) TickBitmapTest:testFuzz_FlipTick(int24) (runs: 65540, μ: 111578, ~: 111849) -TickBitmapTest:testFuzz_NextInitializedTickWithinOneWord(int24,uint8,bool) (runs: 65540, μ: 111544, ~: 111448) +TickBitmapTest:testFuzz_NextInitializedTickWithinOneWord(int24,uint8,bool) (runs: 65540, μ: 111541, ~: 111446) TickBitmapTest:testFuzz_Position(int24) (runs: 65540, μ: 3714, ~: 3714) TickBitmapTest:testGas_NextInitializedTickWithinOneWord() (gas: 12286448) TickBitmapTest:testGas_NextInitializedTickWithinOneWord_Og() (gas: 12567608) -TickBitmapTest:test_NextInitializedTickWithinOneWord_GT() (gas: 231106) +TickBitmapTest:test_NextInitializedTickWithinOneWord_GT() (gas: 230082) TickBitmapTest:test_NextInitializedTickWithinOneWord_LTE() (gas: 232560) -TickBitmapTest:test_NextInitializedTick_GT() (gas: 375838) +TickBitmapTest:test_NextInitializedTick_GT() (gas: 373790) TickBitmapTest:test_NextInitializedTick_LTE() (gas: 420968) TickMathTest:testFuzz_GetSqrtRatioAtTick(int24) (runs: 65540, μ: 17222, ~: 17433) TickMathTest:testFuzz_GetTickAtSqrtRatio(uint160) (runs: 65540, μ: 16021, ~: 16211) diff --git a/src/SqrtPriceMath.sol b/src/SqrtPriceMath.sol index a71c18a..2f55c69 100644 --- a/src/SqrtPriceMath.sol +++ b/src/SqrtPriceMath.sol @@ -93,7 +93,7 @@ library SqrtPriceMath { : FullMath.mulDiv(amount, FixedPoint96.Q96, liquidity) ); - nextSqrtPrice = (sqrtPX96 + quotient).toUint160(); + nextSqrtPrice = (uint256(sqrtPX96) + quotient).toUint160(); } else { uint256 quotient = ( amount <= type(uint160).max diff --git a/src/TickBitmap.sol b/src/TickBitmap.sol index 1909388..401fca3 100644 --- a/src/TickBitmap.sol +++ b/src/TickBitmap.sol @@ -23,15 +23,15 @@ library TickBitmap { } } - /// @notice Computes the word position and the bit position given a tick. + /// @notice Computes the position in the mapping where the initialized bit for a tick lives /// @param tick The tick for which to compute the position /// @return wordPos The key in the mapping containing the word in which the bit is stored /// @return bitPos The bit position in the word where the flag is stored - function position(int24 tick) private pure returns (int16 wordPos, uint8 bitPos) { + function position(int24 tick) internal pure returns (int16 wordPos, uint8 bitPos) { assembly { // signed arithmetic shift right wordPos := sar(8, tick) - bitPos := and(tick, 255) + bitPos := and(tick, 0xff) } } @@ -54,7 +54,7 @@ library TickBitmap { let slot := keccak256(0, 0x40) // mask = 1 << bitPos = 1 << (tick % 256) // self[wordPos] ^= mask - sstore(slot, xor(sload(slot), shl(and(tick, 255), 1))) + sstore(slot, xor(sload(slot), shl(and(tick, 0xff), 1))) } } @@ -77,14 +77,14 @@ library TickBitmap { if (lte) { (int16 wordPos, uint8 bitPos) = position(compressed); - // all the 1s at or to the right of the current bitPos - // mask = (1 << (bitPos + 1)) - 1 - // (bitPos + 1) may be 256 but fine - // masked = self[wordPos] & mask assembly ("memory-safe") { + // all the 1s at or to the right of the current bitPos + // mask = (1 << (bitPos + 1)) - 1 + // (bitPos + 1) may overflow but fine since 1 << 256 = 0 + let mask := sub(shl(add(bitPos, 1), 1), 1) + // masked = self[wordPos] & mask mstore(0, wordPos) mstore(0x20, self.slot) - let mask := sub(shl(add(bitPos, 1), 1), 1) masked := and(sload(keccak256(0, 0x40)), mask) initialized := gt(masked, 0) } @@ -107,13 +107,13 @@ library TickBitmap { ++compressed; } (int16 wordPos, uint8 bitPos) = position(compressed); - // all the 1s at or to the left of the bitPos - // mask = ~((1 << bitPos) - 1) - // masked = self[wordPos] & mask assembly ("memory-safe") { + // all the 1s at or to the left of the bitPos + // mask = ~((1 << bitPos) - 1) = -((1 << bitPos) - 1) - 1 = -(1 << bitPos) + let mask := sub(0, shl(bitPos, 1)) + // masked = self[wordPos] & mask mstore(0, wordPos) mstore(0x20, self.slot) - let mask := not(sub(shl(bitPos, 1), 1)) masked := and(sload(keccak256(0, 0x40)), mask) initialized := gt(masked, 0) } @@ -161,10 +161,10 @@ library TickBitmap { (wordPos, bitPos) = position(compressed); // Reuse the same word if the position doesn't change tickWord = wordPos == lastWordPos ? lastWord : pool.tickBitmap(wordPos); - // all the 1s at or to the right of the current bitPos - // mask = (1 << (bitPos + 1)) - 1 - // (bitPos + 1) may be 256 but fine assembly { + // all the 1s at or to the right of the current bitPos + // mask = (1 << (bitPos + 1)) - 1 + // (bitPos + 1) may overflow but fine since 1 << 256 = 0 let mask := sub(shl(add(bitPos, 1), 1), 1) masked := and(tickWord, mask) initialized := gt(masked, 0) @@ -189,10 +189,10 @@ library TickBitmap { } // Reuse the same word if the position doesn't change tickWord = wordPos == lastWordPos ? lastWord : pool.tickBitmap(wordPos); - // all the 1s at or to the left of the bitPos - // mask = ~((1 << bitPos) - 1) assembly { - let mask := not(sub(shl(bitPos, 1), 1)) + // all the 1s at or to the left of the bitPos + // mask = ~((1 << bitPos) - 1) = -((1 << bitPos) - 1) - 1 = -(1 << bitPos) + let mask := sub(0, shl(bitPos, 1)) masked := and(tickWord, mask) initialized := gt(masked, 0) } @@ -240,10 +240,10 @@ library TickBitmap { (wordPos, bitPos) = position(compressed); // Reuse the same word if the position doesn't change tickWord = wordPos == lastWordPos ? lastWord : pool.tickBitmap(wordPos); - // all the 1s at or to the right of the current bitPos - // mask = (1 << (bitPos + 1)) - 1 - // (bitPos + 1) may be 256 but fine assembly { + // all the 1s at or to the right of the current bitPos + // mask = (1 << (bitPos + 1)) - 1 + // (bitPos + 1) may overflow but fine since 1 << 256 = 0 let mask := sub(shl(add(bitPos, 1), 1), 1) masked := and(tickWord, mask) } @@ -262,10 +262,10 @@ library TickBitmap { } // Reuse the same word if the position doesn't change tickWord = wordPos == lastWordPos ? lastWord : pool.tickBitmap(wordPos); - // all the 1s at or to the left of the bitPos - // mask = ~((1 << bitPos) - 1) assembly { - let mask := not(sub(shl(bitPos, 1), 1)) + // all the 1s at or to the left of the bitPos + // mask = ~((1 << bitPos) - 1) = -((1 << bitPos) - 1) - 1 = -(1 << bitPos) + let mask := sub(0, shl(bitPos, 1)) masked := and(tickWord, mask) } while (masked == 0) {