✨ Add _hasOwnershipCycle to ERC6551

.gas-snapshot

@@ -176,8 +176,8 @@ ERC1155Test:testSafeTransferFromToZeroReverts(uint256) (runs: 278, μ: 71146, ~:
 ERC1155Test:test__codesize() (gas: 42341)
 ERC1271Test:testBasefeeBytecodeContract() (gas: 45430)
 ERC1271Test:testIsValidSignature() (gas: 2464705)
-ERC1271Test:testIsValidSignature(uint256) (runs: 278, μ: 245649, ~: 213657)
-ERC1271Test:test__codesize() (gas: 30392)
+ERC1271Test:testIsValidSignature(uint256) (runs: 278, μ: 251606, ~: 214580)
+ERC1271Test:test__codesize() (gas: 30438)
 ERC1967FactoryTest:testChangeAdmin() (gas: 266356)
 ERC1967FactoryTest:testChangeAdminUnauthorized() (gas: 257653)
 ERC1967FactoryTest:testDeploy() (gas: 257202)
@@ -194,7 +194,7 @@ ERC1967FactoryTest:testUpgradeAndCallWithRevert() (gas: 265824)
 ERC1967FactoryTest:testUpgradeUnauthorized() (gas: 270340)
 ERC1967FactoryTest:testUpgradeWithCorruptedProxy() (gas: 263294)
 ERC1967FactoryTest:test__codesize() (gas: 32011)
-ERC20Invariants:invariantBalanceSum() (runs: 256, calls: 128000, reverts: 78216)
+ERC20Invariants:invariantBalanceSum() (runs: 256, calls: 128000, reverts: 78398)
 ERC20Invariants:test__codesize() (gas: 7532)
 ERC20Test:testApprove() (gas: 35730)
 ERC20Test:testApprove(address,uint256) (runs: 278, μ: 30751, ~: 31181)
@@ -275,24 +275,24 @@ ERC4626Test:testWithdrawWithNotEnoughUnderlyingAmountReverts() (gas: 144343)
 ERC4626Test:testWithdrawZero() (gas: 51874)
 ERC4626Test:test__codesize() (gas: 37175)
 ERC6551Test:testBaseFeeMini() (gas: 39514)
-ERC6551Test:testCdFallback() (gas: 895599)
-ERC6551Test:testDeployERC6551(uint256) (runs: 278, μ: 170165, ~: 168891)
+ERC6551Test:testCdFallback() (gas: 895594)
+ERC6551Test:testDeployERC6551(uint256) (runs: 278, μ: 170456, ~: 168890)
 ERC6551Test:testDeployERC6551Proxy() (gas: 80395)
 ERC6551Test:testExecute() (gas: 507518)
 ERC6551Test:testExecuteBatch() (gas: 817605)
-ERC6551Test:testExecuteBatch(uint256) (runs: 278, μ: 639638, ~: 763255)
+ERC6551Test:testExecuteBatch(uint256) (runs: 278, μ: 612344, ~: 762410)
 ERC6551Test:testInitializeERC6551ProxyImplementation() (gas: 189914)
-ERC6551Test:testIsValidSigner(address) (runs: 278, μ: 167471, ~: 167467)
-ERC6551Test:testOnERC1155BatchReceived() (gas: 1702522)
-ERC6551Test:testOnERC1155Received() (gas: 1699885)
-ERC6551Test:testOnERC721Received() (gas: 1738813)
-ERC6551Test:testOnERC721ReceivedCycles() (gas: 1728089)
-ERC6551Test:testOnERC721ReceivedCyclesWithDifferentChainIds(uint256) (runs: 278, μ: 450735, ~: 455448)
-ERC6551Test:testOwnerWorksWithChainIdChange(uint256,uint256) (runs: 278, μ: 1364052, ~: 1364067)
+ERC6551Test:testIsValidSigner(address) (runs: 278, μ: 167472, ~: 167467)
+ERC6551Test:testOnERC1155BatchReceived() (gas: 1702519)
+ERC6551Test:testOnERC1155Received() (gas: 1699882)
+ERC6551Test:testOnERC721Received() (gas: 1738866)
+ERC6551Test:testOnERC721ReceivedCycles() (gas: 2914176)
+ERC6551Test:testOnERC721ReceivedCyclesWithDifferentChainIds(uint256) (runs: 278, μ: 450409, ~: 455548)
+ERC6551Test:testOwnerWorksWithChainIdChange(uint256,uint256) (runs: 278, μ: 1368677, ~: 1368689)
 ERC6551Test:testSupportsInterface() (gas: 169450)
-ERC6551Test:testUpdateState(uint256) (runs: 278, μ: 235212, ~: 235137)
+ERC6551Test:testUpdateState(uint256) (runs: 278, μ: 235193, ~: 235130)
 ERC6551Test:testUpgrade() (gas: 248450)
-ERC6551Test:test__codesize() (gas: 51546)
+ERC6551Test:test__codesize() (gas: 52354)
 ERC6909Test:testApprove() (gas: 36872)
 ERC6909Test:testApprove(address,uint256,uint256) (runs: 278, μ: 37442, ~: 37514)
 ERC6909Test:testBurn() (gas: 40870)
@@ -1191,7 +1191,7 @@ UpgradeableBeaconTest:testUpgradeableSolidityBeaconOnlyOwnerFunctions() (gas: 26
 UpgradeableBeaconTest:testUpgradeableYulBeaconOnlyFnSelectorNotRecognised() (gas: 172796)
 UpgradeableBeaconTest:testUpgradeableYulBeaconOnlyOwnerFunctions() (gas: 198398)
 UpgradeableBeaconTest:test__codesize() (gas: 8808)
-WETHInvariants:invariantTotalSupplyEqualsBalance() (runs: 256, calls: 128000, reverts: 59410)
+WETHInvariants:invariantTotalSupplyEqualsBalance() (runs: 256, calls: 128000, reverts: 59795)
 WETHInvariants:test__codesize() (gas: 5264)
 WETHTest:testDeposit() (gas: 68090)
 WETHTest:testDeposit(uint256) (runs: 278, μ: 67950, ~: 68384)

src/accounts/ERC6551.sol

@@ -299,6 +299,46 @@ abstract contract ERC6551 is UUPSUpgradeable, Receiver, ERC1271 {
         }
     }
 
+    /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+    /*                      INTERNAL HELPERS                      */
+    /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+    /// @dev Returns whether there is an ownership cycle.
+    function _hasOwnershipCycle() internal view returns (bool result) {
+        uint256 cachedChainId = _cachedChainId;
+        /// @solidity memory-safe-assembly
+        assembly {
+            let m := mload(0x40) // Cache the free memory pointer.
+            extcodecopy(address(), 0x00, 0x4d, 0x60) // `(chainId, tokenContract, tokenId)`.
+            mstore(0x60, 0xfc0c546a) // `token()`.
+            for {} 1 {} {
+                let tokenContract := mload(0x20)
+                mstore(0x20, 0x6352211e) // `ownerOf(uint256)`.
+                let currentOwner :=
+                    mul( // `chainId == cachedChainId ? tokenContract.ownerOf(tokenId) : address(0)`.
+                        mload(0x20),
+                        and(
+                            and(gt(returndatasize(), 0x1f), eq(mload(0x00), cachedChainId)),
+                            staticcall(gas(), tokenContract, 0x3c, 0x24, 0x20, 0x20)
+                        )
+                    )
+                if iszero(eq(currentOwner, address())) {
+                    if iszero(
+                        and( // `(chainId, tokenContract, tokenId) = currentOwner.token()`.
+                            gt(returndatasize(), 0x5f),
+                            staticcall(gas(), currentOwner, 0x7c, 0x04, 0x00, 0x60)
+                        )
+                    ) { break }
+                    continue
+                }
+                result := 1
+                break
+            }
+            mstore(0x40, m) // Restore the free memory pointer.
+            mstore(0x60, 0) // Restore the zero pointer.
+        }
+    }
+
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
     /*                         OVERRIDES                          */
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
@@ -322,51 +362,24 @@ abstract contract ERC6551 is UUPSUpgradeable, Receiver, ERC1271 {
     /// @dev For handling token callbacks.
     /// Safe-transferred ERC721 tokens will trigger a ownership cycle check.
     modifier receiverFallback() override(Receiver) {
-        uint256 cachedChainId = _cachedChainId;
-        /// @solidity memory-safe-assembly
-        assembly {
-            let s := shr(224, calldataload(0x00))
-            // 0x150b7a02: `onERC721Received(address,address,uint256,bytes)`.
-            if eq(s, 0x150b7a02) {
-                extcodecopy(address(), 0x00, 0x4d, 0x60) // `chainId`, `tokenContract`, `tokenId`.
-                mstore(0x60, 0xfc0c546a) // `token()`.
-                for {} 1 {} {
-                    let tokenContract := mload(0x20)
-                    // `tokenId` is already at 0x40.
-                    mstore(0x20, 0x6352211e) // `ownerOf(uint256)`.
-                    let chainsEq := eq(mload(0x00), cachedChainId)
-                    let currentOwner :=
-                        mul(
-                            mload(0x20),
-                            and(
-                                and(gt(returndatasize(), 0x1f), chainsEq),
-                                staticcall(gas(), tokenContract, 0x3c, 0x24, 0x20, 0x20)
-                            )
-                        )
-                    if iszero(
-                        or(
-                            eq(currentOwner, address()),
-                            and(
-                                and(chainsEq, eq(tokenContract, caller())),
-                                eq(mload(0x40), calldataload(0x44))
-                            )
-                        )
-                    ) {
-                        if iszero(
-                            and(
-                                gt(returndatasize(), 0x5f),
-                                staticcall(gas(), currentOwner, 0x7c, 0x04, 0x00, 0x60)
-                            )
-                        ) {
-                            mstore(0x40, s) // Store `msg.sig`.
-                            return(0x5c, 0x20) // Return `msg.sig`.
-                        }
-                        continue
-                    }
-                    mstore(0x00, 0xaed146d3) // `SelfOwnDetected()`.
-                    revert(0x1c, 0x04)
+        uint256 s = uint256(bytes32(msg.sig)) >> 224;
+        // 0x150b7a02: `onERC721Received(address,address,uint256,bytes)`.
+        if (s == 0x150b7a02) {
+            if (!_hasOwnershipCycle()) {
+                /// @solidity memory-safe-assembly
+                assembly {
+                    mstore(0x20, s) // Store `msg.sig`.
+                    return(0x3c, 0x20) // Return `msg.sig`.
                 }
             }
+            /// @solidity memory-safe-assembly
+            assembly {
+                mstore(0x00, 0xaed146d3) // `SelfOwnDetected()`.
+                revert(0x1c, 0x04)
+            }
+        }
+        /// @solidity memory-safe-assembly
+        assembly {
             // 0xf23a6e61: `onERC1155Received(address,address,uint256,uint256,bytes)`.
             // 0xbc197c81: `onERC1155BatchReceived(address,address,uint256[],uint256[],bytes)`.
             if or(eq(s, 0xf23a6e61), eq(s, 0xbc197c81)) {

test/ERC6551.t.sol

@@ -172,6 +172,16 @@ contract ERC6551Test is SoladyTest {
                         t[i].owner, address(t[j].account), t[i].tokenId
                     );
                 }
+                for (uint256 j; j != i; ++j) {
+                    vm.prank(t[i].owner);
+                    MockERC721(_erc721).safeTransferFrom(
+                        t[i].owner, address(t[j].account), t[i].tokenId
+                    );
+                    vm.prank(address(t[j].account));
+                    MockERC721(_erc721).safeTransferFrom(
+                        address(t[j].account), t[i].owner, t[i].tokenId
+                    );
+                }
             }
 
             _TestTemps memory u = _testTemps();