Refactor Multicallable

.gas-snapshot

@@ -1040,18 +1040,18 @@ MinHeapLibTest:testMemHeapSmallestGas() (gas: 2985537)
 MinHeapLibTest:testMemHeapWriteAndReadFromStorage() (gas: 67756)
 MinHeapLibTest:testMemHeapWriteAndReadFromStorage2() (gas: 67774)
 MinHeapLibTest:test__codesize() (gas: 14576)
-MulticallableTest:testMulticallableBenchmark() (gas: 29642)
-MulticallableTest:testMulticallableOriginalBenchmark() (gas: 38935)
-MulticallableTest:testMulticallablePreservesMsgSender() (gas: 11166)
-MulticallableTest:testMulticallableReturnDataIsProperlyEncoded() (gas: 11665)
-MulticallableTest:testMulticallableReturnDataIsProperlyEncoded(string,string,uint256) (runs: 317, μ: 12173, ~: 7437)
-MulticallableTest:testMulticallableReturnDataIsProperlyEncoded(uint256,uint256,uint256,uint256) (runs: 317, μ: 11811, ~: 11811)
-MulticallableTest:testMulticallableRevertWithCustomError() (gas: 11811)
-MulticallableTest:testMulticallableRevertWithMessage() (gas: 13483)
-MulticallableTest:testMulticallableRevertWithMessage(string) (runs: 317, μ: 14095, ~: 13944)
-MulticallableTest:testMulticallableRevertWithNothing() (gas: 11673)
-MulticallableTest:testMulticallableWithNoData() (gas: 6286)
-MulticallableTest:test__codesize() (gas: 8867)
+MulticallableTest:testMulticallableBenchmark() (gas: 30291)
+MulticallableTest:testMulticallableOriginalBenchmark() (gas: 38715)
+MulticallableTest:testMulticallablePreservesMsgSender() (gas: 11413)
+MulticallableTest:testMulticallableReturnDataIsProperlyEncoded() (gas: 397027)
+MulticallableTest:testMulticallableReturnDataIsProperlyEncoded(string,string,uint256) (runs: 317, μ: 241178, ~: 43841)
+MulticallableTest:testMulticallableReturnDataIsProperlyEncoded(uint256,uint256,uint256,uint256) (runs: 317, μ: 250280, ~: 393867)
+MulticallableTest:testMulticallableRevertWithCustomError() (gas: 11859)
+MulticallableTest:testMulticallableRevertWithMessage() (gas: 13509)
+MulticallableTest:testMulticallableRevertWithMessage(string) (runs: 317, μ: 14121, ~: 13970)
+MulticallableTest:testMulticallableRevertWithNothing() (gas: 11766)
+MulticallableTest:testMulticallableWithNoData() (gas: 6361)
+MulticallableTest:test__codesize() (gas: 9846)
 OwnableRolesTest:testBytecodeSize() (gas: 350635)
 OwnableRolesTest:testGrantAndRemoveRolesDirect(address,uint256,uint256) (runs: 317, μ: 41493, ~: 42162)
 OwnableRolesTest:testGrantAndRevokeOrRenounceRoles(address,bool,bool,bool,uint256,uint256) (runs: 317, μ: 27815, ~: 20899)

src/utils/Multicallable.sol

@@ -18,62 +18,81 @@ pragma solidity ^0.8.4;
 /// https://github.com/Vectorized/multicaller
 /// which is more flexible, futureproof, and safer by default.
 abstract contract Multicallable {
-    /// @dev Apply `DELEGATECALL` with the current contract to each calldata in `data`,
-    /// and store the `abi.encode` formatted results of each `DELEGATECALL` into `results`.
-    /// If any of the `DELEGATECALL`s reverts, the entire context is reverted,
+    /// @dev Apply `delegatecall` with the current contract to each calldata in `data`,
+    /// and store the `abi.encode` formatted results of each `delegatecall` into `results`.
+    /// If any of the `delegatecall`s reverts, the entire context is reverted,
     /// and the error is bubbled up.
     ///
     /// This function is deliberately made non-payable to guard against double-spending.
     /// (See: https://www.paradigm.xyz/2021/08/two-rights-might-make-a-wrong)
     ///
-    /// For efficiency, this function will directly return the results, terminating the context.
-    /// If called internally, it must be called at the end of a function
-    /// that returns `(bytes[] memory)`.
+    /// By default, this function directly returns the results and terminates the call context.
+    /// If you need to add before and after actions to the multicall, please override this function.
     function multicall(bytes[] calldata data) public virtual returns (bytes[] memory) {
-        assembly {
-            mstore(0x00, 0x20)
-            mstore(0x20, data.length) // Store `data.length` into `results`.
-            // Early return if no data.
-            if iszero(data.length) { return(0x00, 0x40) }
-
-            let results := 0x40
-            // `shl` 5 is equivalent to multiplying by 0x20.
-            let end := shl(5, data.length)
-            // Copy the offsets from calldata into memory.
-            calldatacopy(0x40, data.offset, end)
-            // Offset into `results`.
-            let resultsOffset := end
-            // Pointer to the end of `results`.
-            end := add(results, end)
+        _multicallDirectReturn(_multicallInner(data));
+    }
 
-            for {} 1 {} {
-                // The offset of the current bytes in the calldata.
-                let o := add(data.offset, mload(results))
-                let m := add(resultsOffset, 0x40)
-                // Copy the current bytes from calldata to the memory.
-                calldatacopy(
-                    m,
-                    add(o, 0x20), // The offset of the current bytes' bytes.
-                    calldataload(o) // The length of the current bytes.
-                )
+    /// @dev The inner logic of `multicall`.
+    /// This function is included so that you can override `multicall`
+    /// to add before and after actions, and use the `_multicallDirectReturn` function.
+    function _multicallInner(bytes[] calldata data)
+        internal
+        virtual
+        returns (bytes[] memory results)
+    {
+        if (data.length == uint256(0)) return results;
+        /// @solidity memory-safe-assembly
+        assembly {
+            results := mload(0x40)
+            mstore(results, data.length)
+            let m := add(results, 0x20)
+            let p := m
+            calldatacopy(p, data.offset, shl(5, data.length))
+            let end := add(p, shl(5, data.length))
+            for { m := end } 1 {} {
+                let o := add(data.offset, mload(p))
+                calldatacopy(m, add(o, 0x20), calldataload(o))
                 if iszero(delegatecall(gas(), address(), m, calldataload(o), codesize(), 0x00)) {
                     // Bubble up the revert if the delegatecall reverts.
-                    returndatacopy(0x00, 0x00, returndatasize())
-                    revert(0x00, returndatasize())
+                    returndatacopy(results, 0x00, returndatasize())
+                    revert(results, returndatasize())
                 }
-                // Append the current `resultsOffset` into `results`.
-                mstore(results, resultsOffset)
-                results := add(results, 0x20)
+                mstore(p, m)
+                p := add(p, 0x20)
                 // Append the `returndatasize()`, and the return data.
                 mstore(m, returndatasize())
-                returndatacopy(add(m, 0x20), 0x00, returndatasize())
-                // Advance the `resultsOffset` by `returndatasize() + 0x20`,
+                o := add(m, 0x20)
+                returndatacopy(o, 0x00, returndatasize())
+                // Zeroize the slot after the returndata.
+                mstore(add(o, returndatasize()), 0x00)
+                // Advance `m` by `returndatasize() + 0x20`,
                 // rounded up to the next multiple of 32.
-                resultsOffset :=
-                    and(add(add(resultsOffset, returndatasize()), 0x3f), 0xffffffffffffffe0)
-                if iszero(lt(results, end)) { break }
+                m := and(add(add(m, returndatasize()), 0x3f), 0xffffffffffffffe0)
+                if iszero(lt(p, end)) { break }
+            }
+            mstore(0x40, m)
+        }
+    }
+
+    /// @dev Directly returns the `results` and terminates the current call context.
+    /// This is more efficient than Solidity's implicit return.
+    function _multicallDirectReturn(bytes[] memory results) internal pure virtual {
+        /// @solidity memory-safe-assembly
+        assembly {
+            if iszero(mload(results)) {
+                mstore(0x40, 0x20)
+                return(0x40, 0x40)
+            }
+            let s := add(0x20, results)
+            let m := s
+            for { let end := add(m, shl(5, mload(results))) } 1 {} {
+                mstore(m, sub(mload(m), s))
+                m := add(m, 0x20)
+                if eq(m, end) { break }
             }
-            return(0x00, add(resultsOffset, 0x40))
+            let o := sub(results, 0x20)
+            mstore(o, 0x20)
+            return(o, sub(mload(0x40), o))
         }
     }
 }

test/Multicallable.t.sol

@@ -46,7 +46,12 @@ contract MulticallableTest is SoladyTest {
         bytes[] memory data = new bytes[](2);
         data[0] = abi.encodeWithSelector(MockMulticallable.returnsTuple.selector, a0, b0);
         data[1] = abi.encodeWithSelector(MockMulticallable.returnsTuple.selector, a1, b1);
-        bytes[] memory returnedData = multicallable.multicall(data);
+        bytes[] memory returnedData;
+        if (_randomChance(2)) {
+            returnedData = multicallable.multicall(data);
+        } else {
+            returnedData = multicallable.multicallBrutalized(data);
+        }
         MockMulticallable.Tuple memory t0 = abi.decode(returnedData[0], (MockMulticallable.Tuple));
         MockMulticallable.Tuple memory t1 = abi.decode(returnedData[1], (MockMulticallable.Tuple));
         assertEq(t0.a, a0);
@@ -70,7 +75,12 @@ contract MulticallableTest is SoladyTest {
             dataIn[1] =
                 abi.encodeWithSelector(MockMulticallable.returnsRandomizedString.selector, sIn1);
         }
-        bytes[] memory dataOut = multicallable.multicall(dataIn);
+        bytes[] memory dataOut;
+        if (_randomChance(2)) {
+            dataOut = multicallable.multicall(dataIn);
+        } else {
+            dataOut = multicallable.multicallBrutalized(dataIn);
+        }
         if (n > 0) {
             assertEq(abi.decode(dataOut[0], (string)), multicallable.returnsRandomizedString(sIn0));
         }

test/utils/mocks/MockMulticallable.sol

@@ -1,11 +1,12 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.4;
 
-import "../../../src/utils/Multicallable.sol";
+import {Multicallable} from "../../../src/utils/Multicallable.sol";
+import {Brutalizer} from "../Brutalizer.sol";
 
 /// @dev WARNING! This mock is strictly intended for testing purposes only.
 /// Do NOT copy anything here into production code unless you really know what you are doing.
-contract MockMulticallable is Multicallable {
+contract MockMulticallable is Multicallable, Brutalizer {
     error CustomError();
 
     struct Tuple {
@@ -63,6 +64,12 @@ contract MockMulticallable is Multicallable {
         return msg.sender;
     }
 
+    function multicallBrutalized(bytes[] calldata data) public returns (bytes[] memory results) {
+        _brutalizeMemory();
+        results = _multicallInner(data);
+        _checkMemory();
+    }
+
     function multicallOriginal(bytes[] calldata data)
         public
         payable