✨ LibBytes calldata extractors (#1361)

Vectorized · Feb 20, 2025 · ca4b9af · ca4b9af
1 parent 1eb89b9
commit ca4b9af
Show file tree

Hide file tree

Showing 4 changed files with 251 additions and 8 deletions.
diff --git a/docs/utils/libbytes.md b/docs/utils/libbytes.md
@@ -366,8 +366,7 @@ function load(bytes memory a, uint256 offset)
     returns (bytes32 result)
 ```
 
-Returns the word at `offset`, without any bounds checks.   
-To load an address, you can use `address(bytes20(load(a, offset)))`.
+Returns the word at `offset`, without any bounds checks.
 
 ### loadCalldata(bytes,uint256)
 
@@ -378,8 +377,40 @@ function loadCalldata(bytes calldata a, uint256 offset)
     returns (bytes32 result)
 ```
 
-Returns the word at `offset`, without any bounds checks.   
-To load an address, you can use `address(bytes20(loadCalldata(a, offset)))`.
+Returns the word at `offset`, without any bounds checks.
+
+### staticStructInCalldata(bytes,uint256)
+
+```solidity
+function staticStructInCalldata(bytes calldata a, uint256 offset)
+    internal
+    pure
+    returns (bytes calldata result)
+```
+
+Returns a slice representing a static struct in the calldata. Performs bounds checks.
+
+### dynamicStructInCalldata(bytes,uint256)
+
+```solidity
+function dynamicStructInCalldata(bytes calldata a, uint256 offset)
+    internal
+    pure
+    returns (bytes calldata result)
+```
+
+Returns a slice representing a dynamic struct in the calldata. Performs bounds checks.
+
+### bytesInCalldata(bytes,uint256)
+
+```solidity
+function bytesInCalldata(bytes calldata a, uint256 offset)
+    internal
+    pure
+    returns (bytes calldata result)
+```
+
+Returns bytes in calldata. Performs bounds checks.
 
 ### emptyCalldata()
 

diff --git a/src/utils/LibBytes.sol b/src/utils/LibBytes.sol
@@ -698,7 +698,6 @@ library LibBytes {
     }
 
     /// @dev Returns the word at `offset`, without any bounds checks.
-    /// To load an address, you can use `address(bytes20(load(a, offset)))`.
     function load(bytes memory a, uint256 offset) internal pure returns (bytes32 result) {
         /// @solidity memory-safe-assembly
         assembly {
@@ -707,7 +706,6 @@ library LibBytes {
     }
 
     /// @dev Returns the word at `offset`, without any bounds checks.
-    /// To load an address, you can use `address(bytes20(loadCalldata(a, offset)))`.
     function loadCalldata(bytes calldata a, uint256 offset)
         internal
         pure
@@ -719,6 +717,55 @@ library LibBytes {
         }
     }
 
+    /// @dev Returns a slice representing a static struct in the calldata. Performs bounds checks.
+    function staticStructInCalldata(bytes calldata a, uint256 offset)
+        internal
+        pure
+        returns (bytes calldata result)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let l := sub(a.length, 0x20)
+            result.offset := add(a.offset, offset)
+            result.length := sub(a.length, offset)
+            if or(shr(64, or(l, a.offset)), gt(offset, l)) { revert(l, 0x00) }
+        }
+    }
+
+    /// @dev Returns a slice representing a dynamic struct in the calldata. Performs bounds checks.
+    function dynamicStructInCalldata(bytes calldata a, uint256 offset)
+        internal
+        pure
+        returns (bytes calldata result)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let l := sub(a.length, 0x20)
+            let s := calldataload(add(a.offset, offset)) // Relative offset of `result` from `a.offset`.
+            result.offset := add(a.offset, s)
+            result.length := sub(a.length, s)
+            if or(shr(64, or(s, or(l, a.offset))), gt(offset, l)) { revert(l, 0x00) }
+        }
+    }
+
+    /// @dev Returns bytes in calldata. Performs bounds checks.
+    function bytesInCalldata(bytes calldata a, uint256 offset)
+        internal
+        pure
+        returns (bytes calldata result)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let l := sub(a.length, 0x20)
+            let s := calldataload(add(a.offset, offset)) // Relative offset of `result` from `a.offset`.
+            result.offset := add(add(a.offset, s), 0x20)
+            result.length := calldataload(add(a.offset, s))
+            // forgefmt: disable-next-item
+            if or(shr(64, or(result.length, or(s, or(l, a.offset)))),
+                or(gt(add(s, result.length), l), gt(offset, l))) { revert(l, 0x00) }
+        }
+    }
+
     /// @dev Returns empty calldata bytes. For silencing the compiler.
     function emptyCalldata() internal pure returns (bytes calldata result) {
         /// @solidity memory-safe-assembly

diff --git a/src/utils/g/LibBytes.sol b/src/utils/g/LibBytes.sol
@@ -702,7 +702,6 @@ library LibBytes {
     }
 
     /// @dev Returns the word at `offset`, without any bounds checks.
-    /// To load an address, you can use `address(bytes20(load(a, offset)))`.
     function load(bytes memory a, uint256 offset) internal pure returns (bytes32 result) {
         /// @solidity memory-safe-assembly
         assembly {
@@ -711,7 +710,6 @@ library LibBytes {
     }
 
     /// @dev Returns the word at `offset`, without any bounds checks.
-    /// To load an address, you can use `address(bytes20(loadCalldata(a, offset)))`.
     function loadCalldata(bytes calldata a, uint256 offset)
         internal
         pure
@@ -723,6 +721,55 @@ library LibBytes {
         }
     }
 
+    /// @dev Returns a slice representing a static struct in the calldata. Performs bounds checks.
+    function staticStructInCalldata(bytes calldata a, uint256 offset)
+        internal
+        pure
+        returns (bytes calldata result)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let l := sub(a.length, 0x20)
+            result.offset := add(a.offset, offset)
+            result.length := sub(a.length, offset)
+            if or(shr(64, or(l, a.offset)), gt(offset, l)) { revert(l, 0x00) }
+        }
+    }
+
+    /// @dev Returns a slice representing a dynamic struct in the calldata. Performs bounds checks.
+    function dynamicStructInCalldata(bytes calldata a, uint256 offset)
+        internal
+        pure
+        returns (bytes calldata result)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let l := sub(a.length, 0x20)
+            let s := calldataload(add(a.offset, offset)) // Relative offset of `result` from `a.offset`.
+            result.offset := add(a.offset, s)
+            result.length := sub(a.length, s)
+            if or(shr(64, or(s, or(l, a.offset))), gt(offset, l)) { revert(l, 0x00) }
+        }
+    }
+
+    /// @dev Returns bytes in calldata. Performs bounds checks.
+    function bytesInCalldata(bytes calldata a, uint256 offset)
+        internal
+        pure
+        returns (bytes calldata result)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let l := sub(a.length, 0x20)
+            let s := calldataload(add(a.offset, offset)) // Relative offset of `result` from `a.offset`.
+            result.offset := add(add(a.offset, s), 0x20)
+            result.length := calldataload(add(a.offset, s))
+            // forgefmt: disable-next-item
+            if or(shr(64, or(result.length, or(s, or(l, a.offset)))),
+                or(gt(add(s, result.length), l), gt(offset, l))) { revert(l, 0x00) }
+        }
+    }
+
     /// @dev Returns empty calldata bytes. For silencing the compiler.
     function emptyCalldata() internal pure returns (bytes calldata result) {
         /// @solidity memory-safe-assembly

diff --git a/test/LibBytes.t.sol b/test/LibBytes.t.sol
@@ -181,6 +181,124 @@ contract LibBytesTest is SoladyTest {
         assertEq(keccak256(b), bHash);
     }
 
+    struct SampleDynamicStruct {
+        address target;
+        uint256 value;
+        bytes data;
+    }
+
+    struct SampleStaticSubStruct {
+        uint256 x;
+        uint256 y;
+    }
+
+    struct SampleStaticStruct {
+        SampleStaticSubStruct a;
+        SampleStaticSubStruct b;
+    }
+
+    function testStaticStructInCalldata() public {
+        SampleStaticStruct memory s;
+        s.a.x = 1;
+        s.a.y = 2;
+        s.b.x = 3;
+        s.b.y = 4;
+
+        SampleDynamicStruct memory u;
+        u.target = address(0xaaa);
+        u.value = 123;
+        u.data = "hehe";
+
+        this._testStaticStructInCalldata(abi.encode(s, u), 0x20 * 0, s);
+        this._testStaticStructInCalldata(abi.encode(u, s, u), 0x20 * 1, s);
+        this._testStaticStructInCalldata(abi.encode(u, u, s, u, s), 0x20 * 2, s);
+        this._testStaticStructInCalldata(abi.encode(u, u, s, u, s), 0x20 * 2 + 0x20 * 4 + 0x20, s);
+    }
+
+    function _testStaticStructInCalldata(
+        bytes calldata encoded,
+        uint256 offset,
+        SampleStaticStruct memory expected
+    ) public {
+        bytes calldata p = LibBytes.staticStructInCalldata(encoded, offset);
+        assertEq(uint256(LibBytes.loadCalldata(p, 0x00)), expected.a.x);
+        assertEq(uint256(LibBytes.loadCalldata(p, 0x20)), expected.a.y);
+        assertEq(uint256(LibBytes.loadCalldata(p, 0x40)), expected.b.x);
+        assertEq(uint256(LibBytes.loadCalldata(p, 0x60)), expected.b.y);
+    }
+
+    function testDynamicStructInCalldata() public {
+        SampleDynamicStruct memory u;
+        u.target = address(1);
+        u.value = 123;
+        u.data = "hehe";
+        bytes memory encoded = abi.encode(u);
+        this._testDynamicStructInCalldata(encoded, 0x00, u);
+    }
+
+    function testDynamicStructInCalldata2() public {
+        SampleDynamicStruct memory u;
+        u.target = address(1);
+        u.value = 123;
+        u.data = "hehe";
+
+        SampleStaticStruct memory s;
+        s.a.x = _random();
+        s.a.y = _random();
+        s.b.x = _random();
+        s.b.y = _random();
+        this._testDynamicStructInCalldata(abi.encode(s, u), 0x80, u);
+        this._testDynamicStructInCalldata(abi.encode(s, u, s), 0x80, u);
+        this._testDynamicStructInCalldata(abi.encode(s, s, u), 0x80 * 2, u);
+    }
+
+    function testDynamicStructInCalldata(bytes32) public {
+        SampleDynamicStruct memory u;
+        u.target = _randomHashedAddress();
+        u.value = _randomUniform();
+        u.data = _truncateBytes(_randomBytes(), 100);
+        bytes memory encoded;
+        encoded = abi.encode(u);
+        this._testDynamicStructInCalldata(encoded, 0x00, u);
+        encoded = abi.encode(uint256(1), u);
+        this._testDynamicStructInCalldata(encoded, 0x20, u);
+        encoded = abi.encode(uint256(1), uint256(2), u);
+        if (_randomChance(32)) encoded = abi.encodePacked(encoded, _randomBytes());
+        this._testDynamicStructInCalldata(encoded, 0x40, u);
+    }
+
+    function _testDynamicStructInCalldata(
+        bytes calldata encoded,
+        uint256 offset,
+        SampleDynamicStruct memory expected
+    ) public {
+        bytes calldata p = LibBytes.dynamicStructInCalldata(encoded, offset);
+        assertEq(uint256(LibBytes.loadCalldata(p, 0x00)), uint160(expected.target));
+        assertEq(uint256(LibBytes.loadCalldata(p, 0x20)), expected.value);
+        assertEq(LibBytes.bytesInCalldata(p, 0x40), expected.data);
+    }
+
+    function testBytesInCalldata() public {
+        this._testBytesInCalldata(abi.encode("hello"), 0x00, "hello");
+    }
+
+    function testBytesInCalldata(bytes32) public {
+        bytes memory u = _truncateBytes(_randomBytes(), 100);
+        this._testBytesInCalldata(abi.encode(u), 0x00, u);
+        this._testBytesInCalldata(abi.encode(uint256(1), u), 0x20, u);
+        if (_randomChance(16)) {
+            bytes memory encoded = abi.encode(uint256(1), uint256(2), u);
+            if (_randomChance(32)) encoded = abi.encodePacked(encoded, _randomBytes());
+            this._testBytesInCalldata(encoded, 0x40, u);
+        }
+    }
+
+    function _testBytesInCalldata(bytes calldata encoded, uint256 offset, bytes memory expected)
+        public
+    {
+        assertEq(LibBytes.bytesInCalldata(encoded, offset), expected);
+    }
+
     function _brutalizeRightPadding(bytes memory s) internal returns (bytes memory result) {
         uint256 n = s.length;
         result = abi.encodePacked(s, _randomUniform(), _randomUniform());