✨ Add new calldata optimal execution mode for ERC7821 (#1485)

* feat: add new commonTo execution mode, for calldata optimization

* feat: add address(0) replacement to new execution mode

* chore: use memory-safe-assembly tag

* chore: natspec

* fix: clean to address, before checking if it is 0

* feat: replace commonTo mode with calldata optimal mode

* chore: replace datas with dataArr

Author: legion2002

src/accounts/ERC7821.sol

@@ -55,6 +55,7 @@ contract ERC7821 is Receiver {
     /// - `0x01000000000000000000...`: Single batch. Does not support optional `opData`.
     /// - `0x01000000000078210001...`: Single batch. Supports optional `opData`.
     /// - `0x01000000000078210002...`: Batch of batches.
+    /// - `0x01000000000078210003...`: Single batch with common `to` address and optional `opData`.
     ///
     /// For the "batch of batches" mode, each batch will be recursively passed into
     /// `execute` internally with mode `0x01000000000078210001...`.
@@ -73,8 +74,10 @@ contract ERC7821 is Receiver {
     function execute(bytes32 mode, bytes calldata executionData) public payable virtual {
         uint256 id = _executionModeId(mode);
         if (id == 3) return _executeBatchOfBatches(mode, executionData);
+        if (id == 4) return _executeCalldataOptimal(mode, executionData);
         Call[] calldata calls;
         bytes calldata opData;
+
         /// @solidity memory-safe-assembly
         assembly {
             if iszero(id) {
@@ -126,7 +129,40 @@ contract ERC7821 is Receiver {
             id := eq(m, 0x01000000000000000000) // 1.
             id := or(shl(1, eq(m, 0x01000000000078210001)), id) // 2.
             id := or(mul(3, eq(m, 0x01000000000078210002)), id) // 3.
+            id := or(mul(4, eq(m, 0x01000000000078210003)), id) // 4.
+        }
+    }
+
+    /// @dev For execution of a batch of batches with a common `to` address.
+    /// @dev if to == address(0), it will be replaced with address(this)
+    /// Execution Data: abi.encode(address to, bytes[] dataArr, bytes opData)
+    function _executeCalldataOptimal(bytes32 mode, bytes calldata executionData) internal virtual {
+        address to;
+        bytes[] calldata dataArr;
+        bytes calldata opData;
+
+        /// @solidity memory-safe-assembly
+        assembly {
+            to := calldataload(executionData.offset)
+
+            let dataOffset :=
+                add(executionData.offset, calldataload(add(0x20, executionData.offset)))
+            dataArr.offset := add(dataOffset, 0x20)
+            dataArr.length := calldataload(dataOffset)
+
+            // This line is needed to ensure that opdata is valid in all code paths.
+            // Otherwise the compiler complains.
+            opData.length := 0
+            // If the offset of `executionData` allows for `opData`, and the mode supports it.
+            if gt(calldataload(add(0x20, executionData.offset)), 0x40) {
+                let opDataOffset :=
+                    add(executionData.offset, calldataload(add(0x40, executionData.offset)))
+                opData.offset := add(opDataOffset, 0x20)
+                opData.length := calldataload(opDataOffset)
+            }
         }
+
+        _execute(mode, executionData, to, dataArr, opData);
     }
 
     /// @dev For execution of a batch of batches.
@@ -190,6 +226,29 @@ contract ERC7821 is Receiver {
         revert(); // In your override, replace this with logic to operate on `opData`.
     }
 
+    /// @dev Executes the calls.
+    /// Reverts and bubbles up error if any call fails.
+    /// The `mode` and `executionData` are passed along in case there's a need to use them.
+    function _execute(
+        bytes32 mode,
+        bytes calldata executionData,
+        address to,
+        bytes[] calldata dataArr,
+        bytes calldata opData
+    ) internal virtual {
+        // Silence compiler warning on unused variables.
+        mode = mode;
+        executionData = executionData;
+        // Very basic auth to only allow this contract to be called by itself.
+        // Override this function to perform more complex auth with `opData`.
+        if (opData.length == uint256(0)) {
+            require(msg.sender == address(this));
+            // Remember to return `_execute(calls, extraData)` when you override this function.
+            return _execute(dataArr, to, bytes32(0));
+        }
+        revert(); // In your override, replace this with logic to operate on `opData`.
+    }
+
     /// @dev Executes the calls.
     /// Reverts and bubbles up error if any call fails.
     /// `extraData` can be any supplementary data (e.g. a memory pointer, some hash).
@@ -204,6 +263,28 @@ contract ERC7821 is Receiver {
         }
     }
 
+    /// @dev Executes the dataArr, with a common `to` address.
+    /// @dev if to == address(0), it will be replaced with address(this)
+    /// @dev value for all calls is set to 0
+    /// Reverts and bubbles up error if any call fails.
+    /// `extraData` can be any supplementary data (e.g. a memory pointer, some hash).
+    function _execute(bytes[] calldata dataArr, address to, bytes32 extraData) internal virtual {
+        unchecked {
+            uint256 i;
+            // If `to` is address(0), it will be replaced with address(this)
+            /// @solidity memory-safe-assembly
+            assembly {
+                let t := shr(96, shl(96, to))
+                to := or(mul(address(), iszero(t)), t)
+            }
+            if (dataArr.length == uint256(0)) return;
+            do {
+
+                _execute(to, 0, dataArr[i], extraData);
+            } while (++i != dataArr.length);
+        }
+    }
+
     /// @dev Executes the call.
     /// Reverts and bubbles up error if any call fails.
     /// `extraData` can be any supplementary data (e.g. a memory pointer, some hash).

test/ERC7821.t.sol

@@ -13,7 +13,7 @@ contract ERC7821Test is SoladyTest {
     address target;
 
     bytes32 internal constant _SUPPORTED_MODE = bytes10(0x01000000000078210001);
-
+    bytes32 internal constant _CALLDATA_OPTIMAL_MODE = bytes10(0x01000000000078210003);
     bytes[] internal _bytes;
 
     function setUp() public {
@@ -54,6 +54,21 @@ contract ERC7821Test is SoladyTest {
         mbe.execute{value: _totalValue(calls)}(_SUPPORTED_MODE, data);
     }
 
+    function testERC7821CalldataOptimalGas() public {
+        vm.pauseGasMetering();
+        vm.deal(address(this), 1 ether);
+
+        bytes[] memory dataArr = new bytes[](2);
+
+        dataArr[0] = abi.encodeWithSignature("returnsBytes(bytes)", "hehe");
+        dataArr[1] = abi.encodeWithSignature("returnsHash(bytes)", "lol");
+
+        bytes memory data = abi.encode(target, dataArr);
+        vm.resumeGasMetering();
+
+        mbe.execute(_CALLDATA_OPTIMAL_MODE, data);
+    }
+
     function testERC7821(bytes memory opData) public {
         vm.deal(address(this), 1 ether);
 
@@ -72,6 +87,18 @@ contract ERC7821Test is SoladyTest {
         assertEq(mbe.lastOpData(), opData);
     }
 
+    function testERC7821CalldataOptimal(bytes memory opData) public {
+        vm.deal(address(this), 1 ether);
+
+        bytes[] memory dataArr = new bytes[](2);
+        dataArr[0] = abi.encodeWithSignature("returnsBytes(bytes)", "hehe");
+        dataArr[1] = abi.encodeWithSignature("returnsHash(bytes)", "lol");
+
+        mbe.execute(_CALLDATA_OPTIMAL_MODE, _encodeCalldataOptimal(target, dataArr, opData));
+
+        assertEq(mbe.lastOpData(), opData);
+    }
+
     function testERC7821ForRevert() public {
         ERC7821.Call[] memory calls = new ERC7821.Call[](1);
         calls[0].to = target;
@@ -82,6 +109,14 @@ contract ERC7821Test is SoladyTest {
         mbe.execute{value: _totalValue(calls)}(_SUPPORTED_MODE, _encode(calls, ""));
     }
 
+    function testERC7821CalldataOptimalForRevert() public {
+        bytes[] memory dataArr = new bytes[](1);
+        dataArr[0] = abi.encodeWithSignature("revertsWithCustomError()");
+
+        vm.expectRevert(CustomError.selector);
+        mbe.execute(_CALLDATA_OPTIMAL_MODE, _encodeCalldataOptimal(target, dataArr, ""));
+    }
+
     function _encode(ERC7821.Call[] memory calls, bytes memory opData)
         internal
         returns (bytes memory)
@@ -90,6 +125,14 @@ contract ERC7821Test is SoladyTest {
         return abi.encode(calls, opData);
     }
 
+    function _encodeCalldataOptimal(address to, bytes[] memory dataArr, bytes memory opData)
+        internal
+        returns (bytes memory)
+    {
+        if (_randomChance(2) && opData.length == 0) return abi.encode(to, dataArr);
+        return abi.encode(to, dataArr, opData);
+    }
+
     struct Payload {
         bytes data;
         uint256 mode;
@@ -125,6 +168,34 @@ contract ERC7821Test is SoladyTest {
         }
     }
 
+    function testERC7821CalldataOptimal(bytes32) public {
+        vm.deal(address(this), 1 ether);
+
+        bytes[] memory dataArr = new bytes[](_randomUniform() & 3);
+        Payload[] memory payloads = new Payload[](dataArr.length);
+
+        for (uint256 i; i < dataArr.length; ++i) {
+            bytes memory data = _truncateBytes(_randomBytes(), 0x1ff);
+            payloads[i].data = data;
+            if (_randomChance(2)) {
+                payloads[i].mode = 0;
+                dataArr[i] = abi.encodeWithSignature("returnsBytes(bytes)", data);
+            } else {
+                payloads[i].mode = 1;
+                dataArr[i] = abi.encodeWithSignature("returnsHash(bytes)", data);
+            }
+        }
+
+        mbe.executeDirect(dataArr, target);
+
+        if (dataArr.length != 0 && _randomChance(32)) {
+            dataArr[_randomUniform() % dataArr.length] =
+                abi.encodeWithSignature("revertsWithCustomError()");
+            vm.expectRevert(CustomError.selector);
+            mbe.executeDirect(dataArr, target);
+        }
+    }
+
     function _totalValue(ERC7821.Call[] memory calls) internal pure returns (uint256 result) {
         unchecked {
             for (uint256 i; i < calls.length; ++i) {
@@ -133,6 +204,7 @@ contract ERC7821Test is SoladyTest {
         }
     }
 
+
     function testERC7821ExecuteBatchOfBatches() public {
         bytes32 mode = bytes32(0x0100000000007821000200000000000000000000000000000000000000000000);
         bytes[] memory batchBytes = new bytes[](3);
@@ -182,4 +254,17 @@ contract ERC7821Test is SoladyTest {
     function pushBytes(bytes memory x) public {
         _bytes.push(x);
     }
+
+    function testERC7821CalldataOptimalWithZeroAddress() public {
+        // Test that when to=address(0), it gets replaced with address(this) (the MockERC7821 contract)
+        // We'll call executeDirect which directly calls the internal _execute function
+        bytes[] memory dataArr = new bytes[](1);
+        dataArr[0] = abi.encodeWithSignature("setAuthorizedCaller(address,bool)", address(0x123), true);
+
+        // This should replace address(0) with address(mbe) and call setAuthorizedCaller on itself
+        mbe.executeDirect(dataArr, address(0));
+
+        // Verify the call succeeded by checking that address(0x123) is now authorized
+        assertTrue(mbe.isAuthorizedCaller(address(0x123)));
+    }
 }

test/utils/mocks/MockERC7821.sol

@@ -22,6 +22,17 @@ contract MockERC7821 is ERC7821, Brutalizer {
         _execute(calls, bytes32(0));
     }
 
+    function _execute(
+        bytes32,
+        bytes calldata,
+        address to,
+        bytes[] calldata dataArr,
+        bytes calldata opData
+    ) internal virtual override {
+        lastOpData = opData;
+        _execute(dataArr, to, bytes32(0));
+    }
+
     function execute(bytes32 mode, bytes calldata executionData) public payable virtual override {
         if (!isAuthorizedCaller[msg.sender]) revert Unauthorized();
         super.execute(mode, executionData);
@@ -34,6 +45,13 @@ contract MockERC7821 is ERC7821, Brutalizer {
         _checkMemory();
     }
 
+    function executeDirect(bytes[] calldata dataArr, address to) public payable virtual {
+        _misalignFreeMemoryPointer();
+        _brutalizeMemory();
+        _execute(dataArr, to, bytes32(0));
+        _checkMemory();
+    }
+
     function setAuthorizedCaller(address target, bool status) public {
         isAuthorizedCaller[target] = status;
     }