Optimize Extsload and Exttload, add comments and fuzz test (#674)

* Optimize `extsload`, add comments and fuzz test Refactored the assembly block in Extsload.sol to improve code readability and optimize gas usage. Changed variable initialization for more efficient looping and calculation of offsets. Also, added a new function for fuzz testing in Extsload.t.sol to ensure the robustness of the `extsload` function with random inputs. * Optimize contiguous `extsload` variant * Optimize single slot `extsload` variant * Optimize `extsload` again * Add dirty bits to fuzzing test of `extsload` Added a new parameter called `dirtyBits` to the function `test_fuzz_extsload` in the test file Extsload.t.sol. With this update, we're now able to simulate the presence of dirty bits during fuzz testing. The test checks for successful execution and correct return data length which should be a multiple of 32. * Make `extsload` more robust * Optimize `Exttload` too
Uniswap · May 23, 2024 · 620fab1 · 620fab1
1 parent b7e8c40
commit 620fab1
Show file tree

Hide file tree

Showing 39 changed files with 90 additions and 58 deletions.
diff --git a/.forge-snapshots/add liquidity to already existing position with salt.snap b/.forge-snapshots/add liquidity to already existing position with salt.snap
@@ -1 +1 @@
-150588
+149414
diff --git a/.forge-snapshots/addLiquidity CA fee.snap b/.forge-snapshots/addLiquidity CA fee.snap
@@ -1 +1 @@
-327334
+326160
diff --git a/.forge-snapshots/addLiquidity with empty hook.snap b/.forge-snapshots/addLiquidity with empty hook.snap
@@ -1 +1 @@
-281448
+280274
diff --git a/.forge-snapshots/addLiquidity with native token.snap b/.forge-snapshots/addLiquidity with native token.snap
@@ -1 +1 @@
-140794
+139620
diff --git a/.forge-snapshots/create new liquidity to a position with salt.snap b/.forge-snapshots/create new liquidity to a position with salt.snap
@@ -1 +1 @@
-298766
+297592
diff --git a/.forge-snapshots/donate gas with 1 token.snap b/.forge-snapshots/donate gas with 1 token.snap
@@ -1 +1 @@
-106753
+106629
diff --git a/.forge-snapshots/donate gas with 2 tokens.snap b/.forge-snapshots/donate gas with 2 tokens.snap
@@ -1 +1 @@
-147247
+147124
diff --git a/.forge-snapshots/extsload getFeeGrowthGlobals.snap b/.forge-snapshots/extsload getFeeGrowthGlobals.snap
@@ -1 +1 @@
-1175
+661
diff --git a/.forge-snapshots/extsload getFeeGrowthInside.snap b/.forge-snapshots/extsload getFeeGrowthInside.snap
@@ -1 +1 @@
-446
+415
diff --git a/.forge-snapshots/extsload getLiquidity.snap b/.forge-snapshots/extsload getLiquidity.snap
@@ -1 +1 @@
-446
+415
diff --git a/.forge-snapshots/extsload getPositionInfo.snap b/.forge-snapshots/extsload getPositionInfo.snap
@@ -1 +1 @@
-1373
+817
diff --git a/.forge-snapshots/extsload getPositionLiquidity.snap b/.forge-snapshots/extsload getPositionLiquidity.snap
@@ -1 +1 @@
-446
+415
diff --git a/.forge-snapshots/extsload getSlot0.snap b/.forge-snapshots/extsload getSlot0.snap
@@ -1 +1 @@
-446
+415
diff --git a/.forge-snapshots/extsload getTickBitmap.snap b/.forge-snapshots/extsload getTickBitmap.snap
@@ -1 +1 @@
-446
+415
diff --git a/.forge-snapshots/extsload getTickFeeGrowthOutside.snap b/.forge-snapshots/extsload getTickFeeGrowthOutside.snap
@@ -1 +1 @@
-1175
+661
diff --git a/.forge-snapshots/extsload getTickInfo.snap b/.forge-snapshots/extsload getTickInfo.snap
@@ -1 +1 @@
-1373
+817
diff --git a/.forge-snapshots/extsload getTickLiquidity.snap b/.forge-snapshots/extsload getTickLiquidity.snap
@@ -1 +1 @@
-446
+415
diff --git a/.forge-snapshots/getReserves.snap b/.forge-snapshots/getReserves.snap
@@ -1 +1 @@
-1198
+1167
diff --git a/.forge-snapshots/poolManager bytecode size.snap b/.forge-snapshots/poolManager bytecode size.snap
@@ -1 +1 @@
-20580
+20498
diff --git a/.forge-snapshots/removeLiquidity CA fee.snap b/.forge-snapshots/removeLiquidity CA fee.snap
@@ -1 +1 @@
-182927
+181753
diff --git a/.forge-snapshots/removeLiquidity with empty hook.snap b/.forge-snapshots/removeLiquidity with empty hook.snap
@@ -1 +1 @@
-137216
+136042
diff --git a/.forge-snapshots/removeLiquidity with native token.snap b/.forge-snapshots/removeLiquidity with native token.snap
@@ -1 +1 @@
-117002
+115828
diff --git a/.forge-snapshots/swap CA custom curve + swap noop.snap b/.forge-snapshots/swap CA custom curve + swap noop.snap
@@ -1 +1 @@
-132247
+132123
diff --git a/.forge-snapshots/swap CA fee on unspecified.snap b/.forge-snapshots/swap CA fee on unspecified.snap
@@ -1 +1 @@
-179756
+179632
diff --git a/.forge-snapshots/swap against liquidity with native token.snap b/.forge-snapshots/swap against liquidity with native token.snap
@@ -1 +1 @@
-110466
+110342
diff --git a/.forge-snapshots/swap against liquidity.snap b/.forge-snapshots/swap against liquidity.snap
@@ -1 +1 @@
-121813
+121689
diff --git a/.forge-snapshots/swap burn 6909 for input.snap b/.forge-snapshots/swap burn 6909 for input.snap
@@ -1 +1 @@
-133642
+133518
diff --git a/.forge-snapshots/swap burn native 6909 for input.snap b/.forge-snapshots/swap burn native 6909 for input.snap
@@ -1 +1 @@
-122930
+122806
diff --git a/.forge-snapshots/swap mint native output as 6909.snap b/.forge-snapshots/swap mint native output as 6909.snap
@@ -1 +1 @@
-144718
+144594
diff --git a/.forge-snapshots/swap mint output as 6909.snap b/.forge-snapshots/swap mint output as 6909.snap
@@ -1 +1 @@
-161348
+161224
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 @@
-217898
+217712
diff --git a/.forge-snapshots/swap with dynamic fee.snap b/.forge-snapshots/swap with dynamic fee.snap
@@ -1 +1 @@
-145549
+145425
diff --git a/.forge-snapshots/swap with hooks.snap b/.forge-snapshots/swap with hooks.snap
@@ -1 +1 @@
-138870
+138746
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 @@
-176786
+176662
diff --git a/.forge-snapshots/swap with return dynamic fee.snap b/.forge-snapshots/swap with return dynamic fee.snap
@@ -1 +1 @@
-152600
+152476
diff --git a/.forge-snapshots/update dynamic fee in before swap.snap b/.forge-snapshots/update dynamic fee in before swap.snap
@@ -1 +1 @@
-155202
+155078
diff --git a/src/Extsload.sol b/src/Extsload.sol
@@ -7,45 +7,54 @@ import {IExtsload} from "./interfaces/IExtsload.sol";
 /// https://eips.ethereum.org/EIPS/eip-2330#rationale
 abstract contract Extsload is IExtsload {
     /// @inheritdoc IExtsload
-    function extsload(bytes32 slot) external view returns (bytes32 value) {
+    function extsload(bytes32 slot) external view returns (bytes32) {
         /// @solidity memory-safe-assembly
         assembly {
-            value := sload(slot)
+            mstore(0, sload(slot))
+            return(0, 0x20)
         }
     }
 
     /// @inheritdoc IExtsload
     function extsload(bytes32 startSlot, uint256 nSlots) external view returns (bytes memory) {
-        bytes memory value = new bytes(32 * nSlots);
-
         /// @solidity memory-safe-assembly
         assembly {
-            for { let i := 0 } lt(i, nSlots) { i := add(i, 1) } {
-                mstore(add(value, mul(add(i, 1), 32)), sload(add(startSlot, i)))
+            // The abi offset of dynamic array in the returndata is 32.
+            mstore(0, 0x20)
+            // A left bit-shift of 5 is equivalent to multiplying by 32 but costs less gas.
+            mstore(0x20, shl(5, nSlots))
+            let end := add(0x40, shl(5, nSlots))
+            for { let memptr := 0x40 } 1 {} {
+                mstore(memptr, sload(startSlot))
+                memptr := add(memptr, 0x20)
+                startSlot := add(startSlot, 1)
+                if iszero(lt(memptr, end)) { break }
             }
+            // The end offset is also the length of the returndata.
+            return(0, end)
         }
-
-        return value;
     }
 
     /// @inheritdoc IExtsload
     function extsload(bytes32[] calldata slots) external view returns (bytes32[] memory) {
         // since the function is external and enters a new call context and exits right
         // after execution, Solidity's memory management convention can be disregarded
         // and a direct slice of memory can be returned
-        assembly ("memory-safe") {
-            // abi offset for dynamic array
-            mstore(0, 0x20)
-            mstore(0x20, slots.length)
+        /// @solidity memory-safe-assembly
+        assembly {
+            // Copy the abi offset of dynamic array and the length of the array to memory.
+            calldatacopy(0, 0x04, 0x40)
+            // A left bit-shift of 5 is equivalent to multiplying by 32 but costs less gas.
             let end := add(0x40, shl(5, slots.length))
-            let memptr := 0x40
             let calldataptr := slots.offset
-            for {} 1 {} {
+            // Return values will start at 64 while calldata offset is 68.
+            for { let memptr := 0x40 } 1 {} {
                 mstore(memptr, sload(calldataload(calldataptr)))
                 memptr := add(memptr, 0x20)
                 calldataptr := add(calldataptr, 0x20)
                 if iszero(lt(memptr, end)) { break }
             }
+            // The end offset is also the length of the returndata.
             return(0, end)
         }
     }

diff --git a/src/Exttload.sol b/src/Exttload.sol
@@ -7,10 +7,11 @@ import {IExttload} from "./interfaces/IExttload.sol";
 /// https://eips.ethereum.org/EIPS/eip-2330#rationale
 abstract contract Exttload is IExttload {
     /// @inheritdoc IExttload
-    function exttload(bytes32 slot) external view returns (bytes32 value) {
+    function exttload(bytes32 slot) external view returns (bytes32) {
         /// @solidity memory-safe-assembly
         assembly {
-            value := tload(slot)
+            mstore(0, tload(slot))
+            return(0, 0x20)
         }
     }
 
@@ -19,19 +20,21 @@ abstract contract Exttload is IExttload {
         // since the function is external and enters a new call context and exits right
         // after execution, Solidity's memory management convention can be disregarded
         // and a direct slice of memory can be returned
-        assembly ("memory-safe") {
-            // abi offset for dynamic array
-            mstore(0, 0x20)
-            mstore(0x20, slots.length)
+        /// @solidity memory-safe-assembly
+        assembly {
+            // Copy the abi offset of dynamic array and the length of the array to memory.
+            calldatacopy(0, 0x04, 0x40)
+            // A left bit-shift of 5 is equivalent to multiplying by 32 but costs less gas.
             let end := add(0x40, shl(5, slots.length))
-            let memptr := 0x40
             let calldataptr := slots.offset
-            for {} 1 {} {
+            // Return values will start at 64 while calldata offset is 68.
+            for { let memptr := 0x40 } 1 {} {
                 mstore(memptr, tload(calldataload(calldataptr)))
                 memptr := add(memptr, 0x20)
                 calldataptr := add(calldataptr, 0x20)
                 if iszero(lt(memptr, end)) { break }
             }
+            // The end offset is also the length of the returndata.
             return(0, end)
         }
     }

diff --git a/test/Extsload.t.sol b/test/Extsload.t.sol
@@ -25,4 +25,24 @@ contract ExtsloadTest is Test, GasSnapshot {
             assertEq(values[i], bytes32(i));
         }
     }
+
+    function test_fuzz_extsload(uint256 length, uint256 seed, bytes memory dirtyBits) public {
+        length = bound(length, 0, 1000);
+        bytes32[] memory slots = new bytes32[](length);
+        bytes32[] memory expected = new bytes32[](length);
+        for (uint256 i; i < length; ++i) {
+            slots[i] = keccak256(abi.encode(i, seed));
+            expected[i] = keccak256(abi.encode(slots[i]));
+            vm.store(address(loadable), slots[i], expected[i]);
+        }
+        bytes32[] memory values = loadable.extsload(slots);
+        assertEq(values, expected);
+        // test with dirty bits
+        bytes memory data = abi.encodeWithSignature("extsload(bytes32[])", (slots));
+        bytes memory malformedData = bytes.concat(data, dirtyBits);
+        (bool success, bytes memory returnData) = address(loadable).staticcall(malformedData);
+        assertTrue(success, "extsload failed");
+        assertEq(returnData.length % 0x20, 0, "return data length is not a multiple of 32");
+        assertEq(abi.decode(returnData, (bytes32[])), expected);
+    }
 }