Add warp around and tests

Vectorized · Feb 10, 2024 · 4db08cf · 4db08cf
1 parent 3fe655d
commit 4db08cf
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 17 deletions.
diff --git a/.gas-snapshot b/.gas-snapshot
@@ -1,10 +1,11 @@
-DN404Test:testBurnOnTransfer(uint32,address,address) (runs: 256, μ: 277892, ~: 279525)
-DN404Test:testInitialize(uint32,address) (runs: 256, μ: 98159, ~: 112410)
-DN404Test:testMintOnTransfer(uint32,address,address) (runs: 256, μ: 255154, ~: 256709)
-DN404Test:testNameAndSymbol(string,string) (runs: 256, μ: 206444, ~: 206785)
+DN404Test:testBurnOnTransfer(uint32,address,address) (runs: 256, μ: 278399, ~: 280264)
+DN404Test:testInitialize(uint32,address) (runs: 256, μ: 100139, ~: 112410)
+DN404Test:testMintOnTransfer(uint32,address,address) (runs: 256, μ: 255841, ~: 257085)
+DN404Test:testNameAndSymbol(string,string) (runs: 256, μ: 206466, ~: 206807)
 DN404Test:testRegisterAndResolveAlias(address,address) (runs: 256, μ: 120543, ~: 120621)
-DN404Test:testSetAndGetOperatorApprovals(address,address,bool) (runs: 256, μ: 129145, ~: 120210)
+DN404Test:testSetAndGetOperatorApprovals(address,address,bool) (runs: 256, μ: 129167, ~: 120232)
 DN404Test:testTokenURI(string,uint256) (runs: 256, μ: 157297, ~: 134998)
-DN404Test:test__codesize() (gas: 20772)
+DN404Test:testWrapAround(uint32,uint256) (runs: 256, μ: 395047, ~: 400286)
+DN404Test:test__codesize() (gas: 22531)
 SoladyTest:test__codesize() (gas: 1075)
 TestPlus:test__codesize() (gas: 398)
diff --git a/src/DN404.sol b/src/DN404.sol
@@ -238,6 +238,14 @@ abstract contract DN404 {
         return addressAlias;
     }
 
+    struct _TransferTemps {
+        address sister;
+        uint256 fromBalanceBefore;
+        uint256 toBalanceBefore;
+        uint256 toBalance;
+        uint256 fromBalance;
+    }
+
     function _transfer(address from, address to, uint256 amount) internal returns (bool) {
         if (to == address(0)) revert TransferToZeroAddress();
 
@@ -246,18 +254,20 @@ abstract contract DN404 {
         AddressData storage fromAddressData = $.addressData[from];
         AddressData storage toAddressData = $.addressData[to];
 
-        uint256 fromBalanceBefore = fromAddressData.balance;
-        fromAddressData.balance = uint96(fromBalanceBefore - amount);
+        _TransferTemps memory t;
+        t.sister = $.sisterERC721;
+
+        t.fromBalanceBefore = fromAddressData.balance;
+        fromAddressData.balance = uint96(t.fromBalance = t.fromBalanceBefore - amount);
 
         unchecked {
-            uint256 toBalanceBefore = toAddressData.balance;
-            toAddressData.balance = uint96(toBalanceBefore + amount);
+            t.toBalanceBefore = toAddressData.balance;
+            toAddressData.balance = uint96(t.toBalance = t.toBalanceBefore + amount);
 
             if (!$.whitelist[from]) {
                 LibMap.Uint32Map storage fromOwned = $.owned[from];
                 uint256 i = fromAddressData.ownedLength;
-                uint256 end =
-                    i - ((fromBalanceBefore / _WAD) - ((fromBalanceBefore - amount) / _WAD));
+                uint256 end = i - ((t.fromBalanceBefore / _WAD) - (t.fromBalance / _WAD));
                 // Burn loop.
                 if (i != end) {
                     do {
@@ -266,7 +276,7 @@ abstract contract DN404 {
                         $.ownerships.set(id, 0);
                         delete $.tokenApprovals[id];
 
-                        _logNFTTransfer($.sisterERC721, from, address(0), id);
+                        _logNFTTransfer(t.sister, from, address(0), id);
                     } while (i != end);
                     fromAddressData.ownedLength = uint32(i);
                 }
@@ -275,22 +285,23 @@ abstract contract DN404 {
             if (!$.whitelist[to]) {
                 LibMap.Uint32Map storage toOwned = $.owned[to];
                 uint256 i = toAddressData.ownedLength;
-                uint256 end = i + (((toBalanceBefore + amount) / _WAD) - (toBalanceBefore / _WAD));
+                uint256 end = i + ((t.toBalance / _WAD) - (t.toBalanceBefore / _WAD));
                 uint256 id = $.nextTokenId;
                 uint32 toAlias = _registerAndResolveAlias(to);
+                uint256 totalNFTSupply = $.totalNFTSupply;
                 // Mint loop.
                 if (i != end) {
                     do {
-                        while ($.ownerships.get(id) != 0) if (++id > _MAX_TOKEN_ID) id = 1;
+                        while ($.ownerships.get(id) != 0) if (++id > totalNFTSupply) id = 1;
 
                         toOwned.set(i, uint32(id));
                         $.ownerships.set(id, toAlias);
                         $.ownedIndex.set(id, uint32(i++));
 
-                        _logNFTTransfer($.sisterERC721, address(0), to, id);
+                        _logNFTTransfer(t.sister, address(0), to, id);
 
                         // todo: ensure we don't overwrite ownership of early tokens that weren't burned
-                        if (++id > _MAX_TOKEN_ID) id = 1;
+                        if (++id > totalNFTSupply) id = 1;
                     } while (i != end);
                     toAddressData.ownedLength = uint32(i);
                     $.nextTokenId = uint32(id);

diff --git a/test/DN404.t.sol b/test/DN404.t.sol
@@ -55,6 +55,25 @@ contract DN404Test is SoladyTest {
         }
     }
 
+    function testWrapAround(uint32 totalNFTSupply, uint256 r) public {
+        address alice = address(111);
+        address bob = address(222);
+        totalNFTSupply = uint32(_bound(totalNFTSupply, 1, 5));
+        dn.initializeDN404(totalNFTSupply, address(this), address(shadow));
+        dn.transfer(alice, _WAD * uint256(totalNFTSupply));
+        for (uint256 t; t != 2; ++t) {
+            uint256 id = _bound(r, 1, totalNFTSupply);
+            vm.prank(alice);
+            shadow.transferFrom(alice, bob, id);
+            vm.prank(bob);
+            shadow.transferFrom(bob, alice, id);
+            vm.prank(alice);
+            dn.transfer(bob, _WAD);
+            vm.prank(bob);
+            dn.transfer(alice, _WAD);
+        }
+    }
+
     function testSetAndGetOperatorApprovals(address owner, address operator, bool approved)
         public
     {