♻️ Make receiverFallback more customizable (#1065)

Vectorized · Sep 4, 2024 · 3f38b6f · 3f38b6f
1 parent 4363564
commit 3f38b6f
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 21 deletions.
diff --git a/.gas-snapshot b/.gas-snapshot
@@ -202,7 +202,7 @@ ERC1155Test:testSafeTransferFromToZeroReverts(uint256) (runs: 317, μ: 137183, ~
 ERC1155Test:test__codesize() (gas: 42015)
 ERC1271Test:testBasefeeBytecodeContract() (gas: 45430)
 ERC1271Test:testIsValidSignature() (gas: 2390984)
-ERC1271Test:testIsValidSignature(uint256) (runs: 317, μ: 247947, ~: 210835)
+ERC1271Test:testIsValidSignature(uint256) (runs: 317, μ: 238238, ~: 210953)
 ERC1271Test:testSupportsNestedTypedDataSign() (gas: 165249)
 ERC1271Test:test__codesize() (gas: 29885)
 ERC1967FactoryTest:testChangeAdmin() (gas: 266261)
@@ -256,24 +256,24 @@ ERC2981Test:testRoyaltyOverflowCheckDifferential(uint256,uint256) (runs: 317, μ
 ERC2981Test:testSetAndGetRoyaltyInfo(uint256) (runs: 317, μ: 108776, ~: 105307)
 ERC2981Test:test__codesize() (gas: 7229)
 ERC4337FactoryTest:testCreateAccountRepeatedDeployment() (gas: 149753)
-ERC4337FactoryTest:testCreateAccountRepeatedDeployment(uint256) (runs: 317, μ: 171339, ~: 171396)
-ERC4337FactoryTest:testDeployDeterministic(uint256) (runs: 317, μ: 132044, ~: 139611)
+ERC4337FactoryTest:testCreateAccountRepeatedDeployment(uint256) (runs: 317, μ: 171337, ~: 171295)
+ERC4337FactoryTest:testDeployDeterministic(uint256) (runs: 317, μ: 132988, ~: 139611)
 ERC4337FactoryTest:test__codesize() (gas: 13112)
 ERC4337Test:testCdFallback() (gas: 438041)
 ERC4337Test:testCdFallback2() (gas: 1135303)
 ERC4337Test:testDelegateExecute() (gas: 366407)
-ERC4337Test:testDelegateExecute(uint256) (runs: 317, μ: 355559, ~: 345250)
+ERC4337Test:testDelegateExecute(uint256) (runs: 317, μ: 355000, ~: 345250)
 ERC4337Test:testDelegateExecuteRevertsIfOwnerSlotValueChanged() (gas: 320067)
 ERC4337Test:testDepositFunctions() (gas: 497221)
 ERC4337Test:testDirectStorage() (gas: 70503)
 ERC4337Test:testDisableInitializerForImplementation() (gas: 1348481)
 ERC4337Test:testETHReceived() (gas: 16584)
 ERC4337Test:testExecute() (gas: 383460)
 ERC4337Test:testExecuteBatch() (gas: 694187)
-ERC4337Test:testExecuteBatch(uint256) (runs: 317, μ: 543702, ~: 670564)
+ERC4337Test:testExecuteBatch(uint256) (runs: 317, μ: 538052, ~: 669950)
 ERC4337Test:testInitializer() (gas: 285469)
 ERC4337Test:testIsValidSignature() (gas: 120476)
-ERC4337Test:testIsValidSignaturePersonalSign(bytes32) (runs: 317, μ: 103745, ~: 103879)
+ERC4337Test:testIsValidSignaturePersonalSign(bytes32) (runs: 317, μ: 103739, ~: 103879)
 ERC4337Test:testIsValidSignatureViaRPC() (gas: 96881551)
 ERC4337Test:testIsValidSignatureWrapped() (gas: 446377)
 ERC4337Test:testOnERC1155BatchReceived() (gas: 1467437)
@@ -303,21 +303,21 @@ ERC4626Test:testWithdrawZero() (gas: 52056)
 ERC4626Test:test__codesize() (gas: 36907)
 ERC6551Test:testBaseFeeMini() (gas: 39514)
 ERC6551Test:testCdFallback() (gas: 895614)
-ERC6551Test:testDeployERC6551(uint256) (runs: 317, μ: 170018, ~: 168568)
+ERC6551Test:testDeployERC6551(uint256) (runs: 317, μ: 170385, ~: 168569)
 ERC6551Test:testDeployERC6551Proxy() (gas: 80395)
 ERC6551Test:testExecute() (gas: 507328)
 ERC6551Test:testExecuteBatch() (gas: 817415)
-ERC6551Test:testExecuteBatch(uint256) (runs: 317, μ: 644333, ~: 762189)
+ERC6551Test:testExecuteBatch(uint256) (runs: 317, μ: 658798, ~: 768700)
 ERC6551Test:testInitializeERC6551ProxyImplementation() (gas: 189976)
-ERC6551Test:testIsValidSigner(address) (runs: 317, μ: 167335, ~: 167323)
+ERC6551Test:testIsValidSigner(address) (runs: 317, μ: 167332, ~: 167320)
 ERC6551Test:testOnERC1155BatchReceived() (gas: 1602799)
 ERC6551Test:testOnERC1155Received() (gas: 1600016)
 ERC6551Test:testOnERC721Received() (gas: 1523963)
 ERC6551Test:testOnERC721ReceivedCycles() (gas: 2975498)
-ERC6551Test:testOnERC721ReceivedCyclesWithDifferentChainIds(uint256) (runs: 317, μ: 452554, ~: 457628)
-ERC6551Test:testOwnerWorksWithChainIdChange(uint256,uint256) (runs: 317, μ: 1386677, ~: 1386687)
+ERC6551Test:testOnERC721ReceivedCyclesWithDifferentChainIds(uint256) (runs: 317, μ: 453067, ~: 457730)
+ERC6551Test:testOwnerWorksWithChainIdChange(uint256,uint256) (runs: 317, μ: 1386684, ~: 1386687)
 ERC6551Test:testSupportsInterface() (gas: 168994)
-ERC6551Test:testUpdateState(uint256) (runs: 317, μ: 234735, ~: 234661)
+ERC6551Test:testUpdateState(uint256) (runs: 317, μ: 234723, ~: 234661)
 ERC6551Test:testUpgrade() (gas: 248410)
 ERC6551Test:test__codesize() (gas: 50997)
 ERC6909Test:testApprove() (gas: 37066)

diff --git a/src/accounts/Receiver.sol b/src/accounts/Receiver.sol
@@ -17,17 +17,32 @@ abstract contract Receiver {
 
     /// @dev Modifier for the fallback function to handle token callbacks.
     modifier receiverFallback() virtual {
-        /// @solidity memory-safe-assembly
-        assembly {
-            let s := shr(224, calldataload(0))
-            // 0x150b7a02: `onERC721Received(address,address,uint256,bytes)`.
-            // 0xf23a6e61: `onERC1155Received(address,address,uint256,uint256,bytes)`.
-            // 0xbc197c81: `onERC1155BatchReceived(address,address,uint256[],uint256[],bytes)`.
-            if or(eq(s, 0x150b7a02), or(eq(s, 0xf23a6e61), eq(s, 0xbc197c81))) {
-                mstore(0x20, s) // Store `msg.sig`.
-                return(0x3c, 0x20) // Return `msg.sig`.
+        _beforeReceiverFallbackBody();
+        if (_useReceiverFallbackBody()) {
+            /// @solidity memory-safe-assembly
+            assembly {
+                let s := shr(224, calldataload(0))
+                // 0x150b7a02: `onERC721Received(address,address,uint256,bytes)`.
+                // 0xf23a6e61: `onERC1155Received(address,address,uint256,uint256,bytes)`.
+                // 0xbc197c81: `onERC1155BatchReceived(address,address,uint256[],uint256[],bytes)`.
+                if or(eq(s, 0x150b7a02), or(eq(s, 0xf23a6e61), eq(s, 0xbc197c81))) {
+                    mstore(0x20, s) // Store `msg.sig`.
+                    return(0x3c, 0x20) // Return `msg.sig`.
+                }
             }
         }
+        _afterReceiverFallbackBody();
         _;
     }
+
+    /// @dev Whether we want to use the body of the `receiverFallback` modifier.
+    function _useReceiverFallbackBody() internal view virtual returns (bool) {
+        return true;
+    }
+
+    /// @dev Called before the body of the `receiverFallback` modifier.
+    function _beforeReceiverFallbackBody() internal virtual {}
+
+    /// @dev Called after the body of the `receiverFallback` modifier.
+    function _afterReceiverFallbackBody() internal virtual {}
 }