From 4d22ea3d159b435a860a2663ca8e8dc41e9d13ce Mon Sep 17 00:00:00 2001
From: Vectorized <webby1111@hotmail.com>
Date: Mon, 6 Feb 2023 12:03:52 +0000
Subject: [PATCH 1/7] Add changes to prepare for bulk burns

---
 contracts/ERC721A.sol                     | 48 +++++++++++------------
 contracts/extensions/ERC721AQueryable.sol | 28 +++++++++----
 2 files changed, 45 insertions(+), 31 deletions(-)

diff --git a/contracts/ERC721A.sol b/contracts/ERC721A.sol
index d35d97141..974c9acd5 100644
--- a/contracts/ERC721A.sol
+++ b/contracts/ERC721A.sol
@@ -346,33 +346,33 @@ contract ERC721A is IERC721A {
     function _packedOwnershipOf(uint256 tokenId) private view returns (uint256 packed) {
         if (_startTokenId() <= tokenId) {
             packed = _packedOwnerships[tokenId];
-            // If not burned.
-            if (packed & _BITMASK_BURNED == 0) {
-                // If the data at the starting slot does not exist, start the scan.
-                if (packed == 0) {
-                    if (tokenId >= _currentIndex) _revert(OwnerQueryForNonexistentToken.selector);
-                    // Invariant:
-                    // There will always be an initialized ownership slot
-                    // (i.e. `ownership.addr != address(0) && ownership.burned == false`)
-                    // before an unintialized ownership slot
-                    // (i.e. `ownership.addr == address(0) && ownership.burned == false`)
-                    // Hence, `tokenId` will not underflow.
-                    //
-                    // We can directly compare the packed value.
-                    // If the address is zero, packed will be zero.
-                    for (;;) {
-                        unchecked {
-                            packed = _packedOwnerships[--tokenId];
-                        }
-                        if (packed == 0) continue;
-                        return packed;
+            // If the data at the starting slot does not exist, start the scan.
+            if (packed == 0) {
+                if (tokenId >= _currentIndex) _revert(OwnerQueryForNonexistentToken.selector);
+                // Invariant:
+                // There will always be an initialized ownership slot
+                // (i.e. `ownership.addr != address(0) && ownership.burned == false`)
+                // before an unintialized ownership slot
+                // (i.e. `ownership.addr == address(0) && ownership.burned == false`)
+                // Hence, `tokenId` will not underflow.
+                //
+                // We can directly compare the packed value.
+                // If the address is zero, packed will be zero.
+                for (;;) {
+                    unchecked {
+                        packed = _packedOwnerships[--tokenId];
                     }
+                    if (packed == 0) continue;
+                    if (packed & _BITMASK_BURNED == 0) return packed;
+                    // Otherwise, the token is burned, and we must revert.
+                    _revert(OwnerQueryForNonexistentToken.selector);
                 }
-                // Otherwise, the data exists and is not burned. We can skip the scan.
-                // This is possible because we have already achieved the target condition.
-                // This saves 2143 gas on transfers of initialized tokens.
-                return packed;
             }
+            // Otherwise, the data exists and is not burned. We can skip the scan.
+            // This is possible because we have already achieved the target condition.
+            // This saves 2143 gas on transfers of initialized tokens.
+            if (packed & _BITMASK_BURNED == 0) return packed;
+            // Otherwise, the token is burned, and we must revert.
         }
         _revert(OwnerQueryForNonexistentToken.selector);
     }
diff --git a/contracts/extensions/ERC721AQueryable.sol b/contracts/extensions/ERC721AQueryable.sol
index 3b1813273..6b0e3d2dd 100644
--- a/contracts/extensions/ERC721AQueryable.sol
+++ b/contracts/extensions/ERC721AQueryable.sol
@@ -44,11 +44,15 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
         override
         returns (TokenOwnership memory ownership)
     {
-        if (tokenId >= _startTokenId()) {
-            if (tokenId < _nextTokenId()) {
-                ownership = _ownershipAt(tokenId);
-                if (!ownership.burned) {
-                    ownership = _ownershipOf(tokenId);
+        unchecked {
+            if (tokenId >= _startTokenId()) {
+                if (tokenId < _nextTokenId()) {
+                    // Simply scan all the way backwards until a non-zero slot.
+                    for (;;) {
+                        ownership = _ownershipAt(tokenId--);
+                        if (ownership.addr != address(0)) 
+                            return ownership;
+                    }
                 }
             }
         }
@@ -157,8 +161,9 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
                 do {
                     ownership = _ownershipAt(start);
                     assembly {
+                        switch mload(add(ownership, 0x40))
                         // if `ownership.burned == false`.
-                        if iszero(mload(add(ownership, 0x40))) {
+                        case 0 {
                             // if `ownership.addr != address(0)`.
                             // The `addr` already has it's upper 96 bits clearned,
                             // since it is written to memory with regular Solidity.
@@ -173,6 +178,10 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
                                 mstore(add(tokenIds, shl(5, tokenIdsIdx)), start)
                             }
                         }
+                        // Otherwise, reset `currOwnershipAddr`.
+                        default {
+                            currOwnershipAddr := 0
+                        }
                         start := add(start, 1)
                     }
                 } while (!(start == stop || tokenIdsIdx == tokenIdsMaxLength));
@@ -212,8 +221,9 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
         for (uint256 i = _startTokenId(); tokenIdsIdx != tokenIdsLength; ) {
             TokenOwnership memory ownership = _ownershipAt(i);
             assembly {
+                switch mload(add(ownership, 0x40))
                 // if `ownership.burned == false`.
-                if iszero(mload(add(ownership, 0x40))) {
+                case 0 {
                     // if `ownership.addr != address(0)`.
                     // The `addr` already has it's upper 96 bits clearned,
                     // since it is written to memory with regular Solidity.
@@ -228,6 +238,10 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
                         mstore(add(tokenIds, shl(5, tokenIdsIdx)), i)
                     }
                 }
+                // Otherwise, reset `currOwnershipAddr`.
+                default {
+                    currOwnershipAddr := 0
+                }
                 i := add(i, 1)
             }
         }

From 5bd4f28c68c82a5f18360f9210d7ac0e5c604f8e Mon Sep 17 00:00:00 2001
From: Vectorized <webby1111@hotmail.com>
Date: Mon, 6 Feb 2023 15:03:38 +0000
Subject: [PATCH 2/7] Add tests and edit Queryable

---
 contracts/ERC721A.sol                         |  2 +-
 contracts/extensions/ERC721AQueryable.sol     | 93 +++++++------------
 contracts/mocks/DirectBurnBitSetterHelper.sol | 21 +++++
 .../mocks/ERC721ABurnableStartTokenIdMock.sol |  2 +-
 contracts/mocks/ERC721AMock.sol               |  3 +-
 contracts/mocks/ERC721AQueryableMock.sol      |  3 +-
 .../ERC721AQueryableStartTokenIdMock.sol      |  2 +-
 contracts/mocks/ERC721AStartTokenIdMock.sol   |  2 +-
 contracts/mocks/StartTokenIdHelper.sol        | 15 ++-
 test/ERC721A.test.js                          | 21 +++++
 test/extensions/ERC721AQueryable.test.js      | 33 +++++++
 test/helpers.js                               |  3 +
 12 files changed, 133 insertions(+), 67 deletions(-)
 create mode 100644 contracts/mocks/DirectBurnBitSetterHelper.sol

diff --git a/contracts/ERC721A.sol b/contracts/ERC721A.sol
index 974c9acd5..5ae5afead 100644
--- a/contracts/ERC721A.sol
+++ b/contracts/ERC721A.sol
@@ -1133,4 +1133,4 @@ contract ERC721A is IERC721A {
             revert(0x00, 0x04)
         }
     }
-}
\ No newline at end of file
+}
diff --git a/contracts/extensions/ERC721AQueryable.sol b/contracts/extensions/ERC721AQueryable.sol
index 6b0e3d2dd..20ace4fea 100644
--- a/contracts/extensions/ERC721AQueryable.sol
+++ b/contracts/extensions/ERC721AQueryable.sol
@@ -47,12 +47,9 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
         unchecked {
             if (tokenId >= _startTokenId()) {
                 if (tokenId < _nextTokenId()) {
-                    // Simply scan all the way backwards until a non-zero slot.
-                    for (;;) {
+                    do {
                         ownership = _ownershipAt(tokenId--);
-                        if (ownership.addr != address(0)) 
-                            return ownership;
-                    }
+                    } while (ownership.addr == address(0));
                 }
             }
         }
@@ -113,6 +110,38 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
         uint256 start,
         uint256 stop
     ) external view virtual override returns (uint256[] memory) {
+        return _tokensOfOwnerIn(owner, start, stop);
+    }
+
+    /**
+     * @dev Returns an array of token IDs owned by `owner`.
+     *
+     * This function scans the ownership mapping and is O(`totalSupply`) in complexity.
+     * It is meant to be called off-chain.
+     *
+     * See {ERC721AQueryable-tokensOfOwnerIn} for splitting the scan into
+     * multiple smaller scans if the collection is large enough to cause
+     * an out-of-gas error (10K collections should be fine).
+     */
+    function tokensOfOwner(address owner) external view virtual override returns (uint256[] memory) {
+        uint256 start = _startTokenId();
+        uint256 stop = _nextTokenId();
+        uint256[] memory tokenIds;
+        if (start != stop) tokenIds = _tokensOfOwnerIn(owner, start, stop);
+        return tokenIds;
+    }
+
+    /**
+     * @dev Helper function for returning an array of token IDs owned by `owner`.
+     *
+     * Note that this function is optimized for smaller bytecode size over runtime gas,
+     * since it is meant to be called off-chain.
+     */
+    function _tokensOfOwnerIn(
+        address owner,
+        uint256 start,
+        uint256 stop
+    ) private view returns (uint256[] memory) {
         unchecked {
             if (start >= stop) _revert(InvalidQueryRange.selector);
             // Set `start = max(start, _startTokenId())`.
@@ -193,58 +222,4 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
             return tokenIds;
         }
     }
-
-    /**
-     * @dev Returns an array of token IDs owned by `owner`.
-     *
-     * This function scans the ownership mapping and is O(`totalSupply`) in complexity.
-     * It is meant to be called off-chain.
-     *
-     * See {ERC721AQueryable-tokensOfOwnerIn} for splitting the scan into
-     * multiple smaller scans if the collection is large enough to cause
-     * an out-of-gas error (10K collections should be fine).
-     */
-    function tokensOfOwner(address owner) external view virtual override returns (uint256[] memory) {
-        uint256 tokenIdsLength = balanceOf(owner);
-        uint256[] memory tokenIds;
-        assembly {
-            // Grab the free memory pointer.
-            tokenIds := mload(0x40)
-            // Allocate one word for the length, and `tokenIdsMaxLength` words
-            // for the data. `shl(5, x)` is equivalent to `mul(32, x)`.
-            mstore(0x40, add(tokenIds, shl(5, add(tokenIdsLength, 1))))
-            // Store the length of `tokenIds`.
-            mstore(tokenIds, tokenIdsLength)
-        }
-        address currOwnershipAddr;
-        uint256 tokenIdsIdx;
-        for (uint256 i = _startTokenId(); tokenIdsIdx != tokenIdsLength; ) {
-            TokenOwnership memory ownership = _ownershipAt(i);
-            assembly {
-                switch mload(add(ownership, 0x40))
-                // if `ownership.burned == false`.
-                case 0 {
-                    // if `ownership.addr != address(0)`.
-                    // The `addr` already has it's upper 96 bits clearned,
-                    // since it is written to memory with regular Solidity.
-                    if mload(ownership) {
-                        currOwnershipAddr := mload(ownership)
-                    }
-                    // if `currOwnershipAddr == owner`.
-                    // The `shl(96, x)` is to make the comparison agnostic to any
-                    // dirty upper 96 bits in `owner`.
-                    if iszero(shl(96, xor(currOwnershipAddr, owner))) {
-                        tokenIdsIdx := add(tokenIdsIdx, 1)
-                        mstore(add(tokenIds, shl(5, tokenIdsIdx)), i)
-                    }
-                }
-                // Otherwise, reset `currOwnershipAddr`.
-                default {
-                    currOwnershipAddr := 0
-                }
-                i := add(i, 1)
-            }
-        }
-        return tokenIds;
-    }
 }
diff --git a/contracts/mocks/DirectBurnBitSetterHelper.sol b/contracts/mocks/DirectBurnBitSetterHelper.sol
new file mode 100644
index 000000000..c9f9efb96
--- /dev/null
+++ b/contracts/mocks/DirectBurnBitSetterHelper.sol
@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: MIT
+// ERC721A Contracts v4.2.3
+// Creators: Chiru Labs
+
+pragma solidity ^0.8.4;
+
+contract DirectBurnBitSetterHelper {
+    function directSetBurnBit(uint256 index) public virtual {
+        uint256 bitmaskBurned = 1 << 224;
+        // We use assembly to directly access the private mapping.
+        // Note that we cannot use this method to test the Upgradeable variant
+        // as it uses Diamond storage layout.
+        assembly {
+            // The `_packedOwnerships` mapping is at slot 4.
+            mstore(0x20, 4)
+            mstore(0x00, index)
+            let ownershipStorageSlot := keccak256(0x00, 0x40)
+            sstore(ownershipStorageSlot, or(sload(ownershipStorageSlot), bitmaskBurned))
+        }
+    }
+}
diff --git a/contracts/mocks/ERC721ABurnableStartTokenIdMock.sol b/contracts/mocks/ERC721ABurnableStartTokenIdMock.sol
index 212d29bd3..f37f6bcec 100644
--- a/contracts/mocks/ERC721ABurnableStartTokenIdMock.sol
+++ b/contracts/mocks/ERC721ABurnableStartTokenIdMock.sol
@@ -15,6 +15,6 @@ contract ERC721ABurnableStartTokenIdMock is StartTokenIdHelper, ERC721ABurnableM
     ) StartTokenIdHelper(startTokenId_) ERC721ABurnableMock(name_, symbol_) {}
 
     function _startTokenId() internal view override returns (uint256) {
-        return startTokenId;
+        return startTokenId();
     }
 }
diff --git a/contracts/mocks/ERC721AMock.sol b/contracts/mocks/ERC721AMock.sol
index ca5e53523..89beacc6a 100644
--- a/contracts/mocks/ERC721AMock.sol
+++ b/contracts/mocks/ERC721AMock.sol
@@ -5,8 +5,9 @@
 pragma solidity ^0.8.4;
 
 import '../ERC721A.sol';
+import './DirectBurnBitSetterHelper.sol';
 
-contract ERC721AMock is ERC721A {
+contract ERC721AMock is ERC721A, DirectBurnBitSetterHelper {
     constructor(string memory name_, string memory symbol_) ERC721A(name_, symbol_) {}
 
     function numberMinted(address owner) public view returns (uint256) {
diff --git a/contracts/mocks/ERC721AQueryableMock.sol b/contracts/mocks/ERC721AQueryableMock.sol
index 74d18b9bf..80f885760 100644
--- a/contracts/mocks/ERC721AQueryableMock.sol
+++ b/contracts/mocks/ERC721AQueryableMock.sol
@@ -6,8 +6,9 @@ pragma solidity ^0.8.4;
 
 import '../extensions/ERC721AQueryable.sol';
 import '../extensions/ERC721ABurnable.sol';
+import './DirectBurnBitSetterHelper.sol';
 
-contract ERC721AQueryableMock is ERC721AQueryable, ERC721ABurnable {
+contract ERC721AQueryableMock is ERC721AQueryable, ERC721ABurnable, DirectBurnBitSetterHelper {
     constructor(string memory name_, string memory symbol_) ERC721A(name_, symbol_) {}
 
     function safeMint(address to, uint256 quantity) public {
diff --git a/contracts/mocks/ERC721AQueryableStartTokenIdMock.sol b/contracts/mocks/ERC721AQueryableStartTokenIdMock.sol
index 3b01084d3..61e9e9ac8 100644
--- a/contracts/mocks/ERC721AQueryableStartTokenIdMock.sol
+++ b/contracts/mocks/ERC721AQueryableStartTokenIdMock.sol
@@ -15,6 +15,6 @@ contract ERC721AQueryableStartTokenIdMock is StartTokenIdHelper, ERC721AQueryabl
     ) StartTokenIdHelper(startTokenId_) ERC721AQueryableMock(name_, symbol_) {}
 
     function _startTokenId() internal view override returns (uint256) {
-        return startTokenId;
+        return startTokenId();
     }
 }
diff --git a/contracts/mocks/ERC721AStartTokenIdMock.sol b/contracts/mocks/ERC721AStartTokenIdMock.sol
index fa4acbf70..cba545e7e 100644
--- a/contracts/mocks/ERC721AStartTokenIdMock.sol
+++ b/contracts/mocks/ERC721AStartTokenIdMock.sol
@@ -15,6 +15,6 @@ contract ERC721AStartTokenIdMock is StartTokenIdHelper, ERC721AMock {
     ) StartTokenIdHelper(startTokenId_) ERC721AMock(name_, symbol_) {}
 
     function _startTokenId() internal view override returns (uint256) {
-        return startTokenId;
+        return startTokenId();
     }
 }
diff --git a/contracts/mocks/StartTokenIdHelper.sol b/contracts/mocks/StartTokenIdHelper.sol
index 22c63b8a2..1a80ebe50 100644
--- a/contracts/mocks/StartTokenIdHelper.sol
+++ b/contracts/mocks/StartTokenIdHelper.sol
@@ -10,9 +10,20 @@ pragma solidity ^0.8.4;
  * to be returned by the overridden `_startTokenId()` function of ERC721A in the ERC721AStartTokenId mocks.
  */
 contract StartTokenIdHelper {
-    uint256 public startTokenId;
+    // `bytes4(keccak256('startTokenId'))`.
+    uint256 private constant _START_TOKEN_ID_STORAGE_SLOT = 0x28f75032;
 
     constructor(uint256 startTokenId_) {
-        startTokenId = startTokenId_;
+        // We use assembly to directly set the `startTokenId` in storage so that
+        // inheriting this class won't affect the layout of other storage slots.
+        assembly {
+            sstore(_START_TOKEN_ID_STORAGE_SLOT, startTokenId_)
+        }
+    }
+
+    function startTokenId() public view returns (uint256 result) {
+        assembly {
+            result := sload(_START_TOKEN_ID_STORAGE_SLOT)
+        }
     }
 }
diff --git a/test/ERC721A.test.js b/test/ERC721A.test.js
index 43e1cdc4a..07b62fb24 100644
--- a/test/ERC721A.test.js
+++ b/test/ERC721A.test.js
@@ -701,6 +701,27 @@ const createTestSuite = ({ contract, constructorArgs }) =>
         });
       });
 
+      context('with direct set burn bit', async function () {
+        it('ownerOf reverts for an uninitialized burnt token', async function () {
+          const [owner] = await ethers.getSigners();
+          await this.erc721a['safeMint(address,uint256)'](owner.address, 3);
+          await this.erc721a['safeMint(address,uint256)'](owner.address, 2);
+          await this.erc721a['safeMint(address,uint256)'](owner.address, 1);
+          for (let i = 0; i < 3 + 2 + 1; ++i) {
+            expect(await this.erc721a.ownerOf(this.startTokenId + i)).to.eq(owner.address);
+          }
+          await this.erc721a.directSetBurnBit(this.startTokenId + 3);
+          for (let i = 0; i < 3 + 2 + 1; ++i) {
+            if (3 <= i && i < 3 + 2) {
+              await expect(this.erc721a.ownerOf(this.startTokenId + i))
+                .to.be.revertedWith('OwnerQueryForNonexistentToken');
+            } else {
+              expect(await this.erc721a.ownerOf(this.startTokenId + i)).to.eq(owner.address);  
+            }
+          }
+        });
+      });
+
       context('_toString', async function () {
         it('returns correct value', async function () {
           expect(await this.erc721a['toString(uint256)']('0')).to.eq('0');
diff --git a/test/extensions/ERC721AQueryable.test.js b/test/extensions/ERC721AQueryable.test.js
index b5b9a39c6..de9e1fe63 100644
--- a/test/extensions/ERC721AQueryable.test.js
+++ b/test/extensions/ERC721AQueryable.test.js
@@ -161,6 +161,18 @@ const createTestSuite = ({ contract, constructorArgs }) =>
             // Verify the function can still read the correct token ids
             expect(ownerTokens).to.eql(offsetted(6, 8));
           });
+
+          it('with direct set burn bit', async function () {
+            if (this.erc721aQueryable.isUpgradeable) return;
+            await this.erc721aQueryable['safeMint(address,uint256)'](this.addr3.address, 3);
+            await this.erc721aQueryable['safeMint(address,uint256)'](this.addr3.address, 2);
+            const nextTokenId = this.owner.expected.tokens[this.owner.expected.tokens.length - 1].add(1);
+            expect(await this.erc721aQueryable.tokensOfOwner(this.addr3.address))
+              .to.eql(this.addr3.expected.tokens.concat(Array.from({length: 5}, (_, i) => nextTokenId.add(i))));
+            await this.erc721aQueryable.directSetBurnBit(nextTokenId);
+            expect(await this.erc721aQueryable.tokensOfOwner(this.addr3.address))
+              .to.eql(this.addr3.expected.tokens.concat(Array.from({length: 2}, (_, i) => nextTokenId.add(3 + i))));
+          });
         });
 
         describe('tokensOfOwnerIn', async function () {
@@ -218,6 +230,18 @@ const createTestSuite = ({ contract, constructorArgs }) =>
           subTests('after a token burn', async function () {
             await this.erc721aQueryable.burn(offsetted(7));
           });
+
+          it('with direct set burn bit', async function () {
+            if (this.erc721aQueryable.isUpgradeable) return;
+            await this.erc721aQueryable['safeMint(address,uint256)'](this.addr3.address, 3);
+            await this.erc721aQueryable['safeMint(address,uint256)'](this.addr3.address, 2);
+            const nextTokenId = this.owner.expected.tokens[this.owner.expected.tokens.length - 1].add(1);
+            expect(await this.erc721aQueryable.tokensOfOwnerIn(this.addr3.address, 0, 99))
+              .to.eql(this.addr3.expected.tokens.concat(Array.from({length: 5}, (_, i) => nextTokenId.add(i))));
+            await this.erc721aQueryable.directSetBurnBit(nextTokenId);
+            expect(await this.erc721aQueryable.tokensOfOwnerIn(this.addr3.address, 0, 99))
+              .to.eql(this.addr3.expected.tokens.concat(Array.from({length: 2}, (_, i) => nextTokenId.add(3 + i))));
+          })
         });
 
         describe('explicitOwnershipOf', async function () {
@@ -245,6 +269,15 @@ const createTestSuite = ({ contract, constructorArgs }) =>
             const explicitOwnership = await this.erc721aQueryable.explicitOwnershipOf(this.currentIndex);
             expectExplicitOwnershipNotExists(explicitOwnership);
           });
+
+          it('with direct set burn bit', async function () {
+            if (this.erc721aQueryable.isUpgradeable) return;
+            await this.erc721aQueryable.directSetBurnBit(this.addr3.expected.tokens[0]);
+            for (let i = 0; i < this.addr3.expected.tokens.length; ++i) {
+              const explicitOwnership = await this.erc721aQueryable.explicitOwnershipOf(this.addr3.expected.tokens[i]);
+              expectExplicitOwnershipBurned(explicitOwnership, this.addr3.address);
+            }
+          });
         });
 
         describe('explicitOwnershipsOf', async function () {
diff --git a/test/helpers.js b/test/helpers.js
index a3d697254..b726532e1 100644
--- a/test/helpers.js
+++ b/test/helpers.js
@@ -3,13 +3,16 @@ const { BigNumber } = require('ethers');
 
 const deployContract = async function (contractName, constructorArgs) {
   let factory;
+  let isUpgradeable = false;
   try {
     factory = await ethers.getContractFactory(contractName);
   } catch (e) {
     factory = await ethers.getContractFactory(contractName + 'UpgradeableWithInit');
+    isUpgradeable = true;
   }
   let contract = await factory.deploy(...(constructorArgs || []));
   await contract.deployed();
+  contract.isUpgradeable = isUpgradeable;
   return contract;
 };
 

From a47f4de53605e9236cff00e935dd264a7026976b Mon Sep 17 00:00:00 2001
From: Vectorized <webby1111@hotmail.com>
Date: Mon, 6 Feb 2023 15:06:54 +0000
Subject: [PATCH 3/7] Tidy

---
 contracts/mocks/DirectBurnBitSetterHelper.sol | 4 ++--
 test/ERC721A.test.js                          | 2 ++
 test/extensions/ERC721AQueryable.test.js      | 1 +
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/contracts/mocks/DirectBurnBitSetterHelper.sol b/contracts/mocks/DirectBurnBitSetterHelper.sol
index c9f9efb96..768d167dc 100644
--- a/contracts/mocks/DirectBurnBitSetterHelper.sol
+++ b/contracts/mocks/DirectBurnBitSetterHelper.sol
@@ -8,8 +8,8 @@ contract DirectBurnBitSetterHelper {
     function directSetBurnBit(uint256 index) public virtual {
         uint256 bitmaskBurned = 1 << 224;
         // We use assembly to directly access the private mapping.
-        // Note that we cannot use this method to test the Upgradeable variant
-        // as it uses Diamond storage layout.
+        // Note that we cannot use this method to test the upgradeable variant
+        // as it uses EIP-2535 diamond storage layout.
         assembly {
             // The `_packedOwnerships` mapping is at slot 4.
             mstore(0x20, 4)
diff --git a/test/ERC721A.test.js b/test/ERC721A.test.js
index 07b62fb24..2d709ec7b 100644
--- a/test/ERC721A.test.js
+++ b/test/ERC721A.test.js
@@ -703,6 +703,8 @@ const createTestSuite = ({ contract, constructorArgs }) =>
 
       context('with direct set burn bit', async function () {
         it('ownerOf reverts for an uninitialized burnt token', async function () {
+          // This test is only for the non-upgradeable variant.
+          if (this.erc721a.isUpgradeable) return;
           const [owner] = await ethers.getSigners();
           await this.erc721a['safeMint(address,uint256)'](owner.address, 3);
           await this.erc721a['safeMint(address,uint256)'](owner.address, 2);
diff --git a/test/extensions/ERC721AQueryable.test.js b/test/extensions/ERC721AQueryable.test.js
index de9e1fe63..82ac6062a 100644
--- a/test/extensions/ERC721AQueryable.test.js
+++ b/test/extensions/ERC721AQueryable.test.js
@@ -163,6 +163,7 @@ const createTestSuite = ({ contract, constructorArgs }) =>
           });
 
           it('with direct set burn bit', async function () {
+            // This test is only for the non-upgradeable variant.
             if (this.erc721aQueryable.isUpgradeable) return;
             await this.erc721aQueryable['safeMint(address,uint256)'](this.addr3.address, 3);
             await this.erc721aQueryable['safeMint(address,uint256)'](this.addr3.address, 2);

From 26f65a2d2b0dc4d9366fb1962d281094d9bd8b0a Mon Sep 17 00:00:00 2001
From: Vectorized <webby1111@hotmail.com>
Date: Mon, 6 Feb 2023 15:35:05 +0000
Subject: [PATCH 4/7] Add _ownershipIsInitialized

---
 contracts/ERC721A.sol                     | 8 ++++++++
 contracts/extensions/ERC721AQueryable.sol | 7 ++++---
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/contracts/ERC721A.sol b/contracts/ERC721A.sol
index 5ae5afead..2f40f6262 100644
--- a/contracts/ERC721A.sol
+++ b/contracts/ERC721A.sol
@@ -331,6 +331,14 @@ contract ERC721A is IERC721A {
         return _unpackedOwnership(_packedOwnerships[index]);
     }
 
+    /**
+     * @dev Returns whether the ownership slot at `index` is initialized.
+     * An uninitialized slot does not necessarily mean that the slot has no owner.
+     */
+    function _ownershipIsInitialized(uint256 index) internal view virtual returns (bool) {
+        return _packedOwnerships[index] != 0;
+    }
+
     /**
      * @dev Initializes the ownership slot minted at `index` for efficiency purposes.
      */
diff --git a/contracts/extensions/ERC721AQueryable.sol b/contracts/extensions/ERC721AQueryable.sol
index 20ace4fea..bc9283611 100644
--- a/contracts/extensions/ERC721AQueryable.sol
+++ b/contracts/extensions/ERC721AQueryable.sol
@@ -47,9 +47,10 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
         unchecked {
             if (tokenId >= _startTokenId()) {
                 if (tokenId < _nextTokenId()) {
-                    do {
-                        ownership = _ownershipAt(tokenId--);
-                    } while (ownership.addr == address(0));
+                    // If the `tokenId` is within bounds,
+                    // scan backwards for the initialized ownership slot.
+                    while (!_ownershipIsInitialized(tokenId)) --tokenId;
+                    return _ownershipAt(tokenId);
                 }
             }
         }

From 432c71a752278415cda9254564ae1a52f26eb1c9 Mon Sep 17 00:00:00 2001
From: Vectorized <webby1111@hotmail.com>
Date: Mon, 6 Feb 2023 15:47:27 +0000
Subject: [PATCH 5/7] Add comments

---
 contracts/ERC721A.sol                         | 2 ++
 contracts/extensions/ERC721AQueryable.sol     | 2 ++
 contracts/mocks/DirectBurnBitSetterHelper.sol | 1 +
 3 files changed, 5 insertions(+)

diff --git a/contracts/ERC721A.sol b/contracts/ERC721A.sol
index 2f40f6262..d4a30897a 100644
--- a/contracts/ERC721A.sol
+++ b/contracts/ERC721A.sol
@@ -373,6 +373,8 @@ contract ERC721A is IERC721A {
                     if (packed == 0) continue;
                     if (packed & _BITMASK_BURNED == 0) return packed;
                     // Otherwise, the token is burned, and we must revert.
+                    // This handles the case of batch burned tokens, where only the burned bit
+                    // of the starting slot is set, and remaining slots are left uninitialized.
                     _revert(OwnerQueryForNonexistentToken.selector);
                 }
             }
diff --git a/contracts/extensions/ERC721AQueryable.sol b/contracts/extensions/ERC721AQueryable.sol
index bc9283611..081815bd3 100644
--- a/contracts/extensions/ERC721AQueryable.sol
+++ b/contracts/extensions/ERC721AQueryable.sol
@@ -209,6 +209,8 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
                             }
                         }
                         // Otherwise, reset `currOwnershipAddr`.
+                        // This handles the case of batch burned tokens
+                        // (burned bit of first slot set, remaining slots left uninitialized).
                         default {
                             currOwnershipAddr := 0
                         }
diff --git a/contracts/mocks/DirectBurnBitSetterHelper.sol b/contracts/mocks/DirectBurnBitSetterHelper.sol
index 768d167dc..0dabe247f 100644
--- a/contracts/mocks/DirectBurnBitSetterHelper.sol
+++ b/contracts/mocks/DirectBurnBitSetterHelper.sol
@@ -6,6 +6,7 @@ pragma solidity ^0.8.4;
 
 contract DirectBurnBitSetterHelper {
     function directSetBurnBit(uint256 index) public virtual {
+        // This is `_BITMASK_BURNED` from ERC721A.
         uint256 bitmaskBurned = 1 << 224;
         // We use assembly to directly access the private mapping.
         // Note that we cannot use this method to test the upgradeable variant

From 1a3bd70bb985fe72476c3f04cbe857da136a0e37 Mon Sep 17 00:00:00 2001
From: Vectorized <webby1111@hotmail.com>
Date: Mon, 6 Feb 2023 16:02:11 +0000
Subject: [PATCH 6/7] Edit comments

---
 contracts/ERC721A.sol | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/contracts/ERC721A.sol b/contracts/ERC721A.sol
index d4a30897a..580f480e0 100644
--- a/contracts/ERC721A.sol
+++ b/contracts/ERC721A.sol
@@ -378,11 +378,11 @@ contract ERC721A is IERC721A {
                     _revert(OwnerQueryForNonexistentToken.selector);
                 }
             }
-            // Otherwise, the data exists and is not burned. We can skip the scan.
+            // Otherwise, the data exists and we can skip the scan.
             // This is possible because we have already achieved the target condition.
             // This saves 2143 gas on transfers of initialized tokens.
+            // If the token is not burned, return `packed`. Otherwise, revert.
             if (packed & _BITMASK_BURNED == 0) return packed;
-            // Otherwise, the token is burned, and we must revert.
         }
         _revert(OwnerQueryForNonexistentToken.selector);
     }

From c069d18d2ff347b57352c64434ad0291f57bbed8 Mon Sep 17 00:00:00 2001
From: Vectorized <webby1111@hotmail.com>
Date: Fri, 10 Feb 2023 23:18:22 +0000
Subject: [PATCH 7/7] Make tests compatible with diamond storage

---
 contracts/mocks/DirectBurnBitSetterHelper.sol | 10 ++++++++--
 contracts/mocks/StartTokenIdHelper.sol        | 14 +++++++++-----
 test/ERC721A.test.js                          |  2 --
 test/extensions/ERC721AQueryable.test.js      |  4 ----
 test/helpers.js                               |  3 ---
 5 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/contracts/mocks/DirectBurnBitSetterHelper.sol b/contracts/mocks/DirectBurnBitSetterHelper.sol
index 0dabe247f..088802014 100644
--- a/contracts/mocks/DirectBurnBitSetterHelper.sol
+++ b/contracts/mocks/DirectBurnBitSetterHelper.sol
@@ -6,17 +6,23 @@ pragma solidity ^0.8.4;
 
 contract DirectBurnBitSetterHelper {
     function directSetBurnBit(uint256 index) public virtual {
+        bytes32 erc721aDiamondStorageSlot = keccak256('openzepplin.contracts.storage.ERC721A');
+
         // This is `_BITMASK_BURNED` from ERC721A.
         uint256 bitmaskBurned = 1 << 224;
         // We use assembly to directly access the private mapping.
-        // Note that we cannot use this method to test the upgradeable variant
-        // as it uses EIP-2535 diamond storage layout.
         assembly {
             // The `_packedOwnerships` mapping is at slot 4.
             mstore(0x20, 4)
             mstore(0x00, index)
             let ownershipStorageSlot := keccak256(0x00, 0x40)
             sstore(ownershipStorageSlot, or(sload(ownershipStorageSlot), bitmaskBurned))
+
+            // For diamond storage, we'll simply add the offset of the layout struct.
+            mstore(0x20, add(erc721aDiamondStorageSlot, 4))
+            mstore(0x00, index)
+            ownershipStorageSlot := keccak256(0x00, 0x40)
+            sstore(ownershipStorageSlot, or(sload(ownershipStorageSlot), bitmaskBurned))
         }
     }
 }
diff --git a/contracts/mocks/StartTokenIdHelper.sol b/contracts/mocks/StartTokenIdHelper.sol
index 1a80ebe50..05ba7f949 100644
--- a/contracts/mocks/StartTokenIdHelper.sol
+++ b/contracts/mocks/StartTokenIdHelper.sol
@@ -14,11 +14,7 @@ contract StartTokenIdHelper {
     uint256 private constant _START_TOKEN_ID_STORAGE_SLOT = 0x28f75032;
 
     constructor(uint256 startTokenId_) {
-        // We use assembly to directly set the `startTokenId` in storage so that
-        // inheriting this class won't affect the layout of other storage slots.
-        assembly {
-            sstore(_START_TOKEN_ID_STORAGE_SLOT, startTokenId_)
-        }
+        _initializeStartTokenId(startTokenId_);
     }
 
     function startTokenId() public view returns (uint256 result) {
@@ -26,4 +22,12 @@ contract StartTokenIdHelper {
             result := sload(_START_TOKEN_ID_STORAGE_SLOT)
         }
     }
+
+    function _initializeStartTokenId(uint256 value) private {
+        // We use assembly to directly set the `startTokenId` in storage so that
+        // inheriting this class won't affect the layout of other storage slots.
+        assembly {
+            sstore(_START_TOKEN_ID_STORAGE_SLOT, value)
+        }
+    }
 }
diff --git a/test/ERC721A.test.js b/test/ERC721A.test.js
index 2d709ec7b..07b62fb24 100644
--- a/test/ERC721A.test.js
+++ b/test/ERC721A.test.js
@@ -703,8 +703,6 @@ const createTestSuite = ({ contract, constructorArgs }) =>
 
       context('with direct set burn bit', async function () {
         it('ownerOf reverts for an uninitialized burnt token', async function () {
-          // This test is only for the non-upgradeable variant.
-          if (this.erc721a.isUpgradeable) return;
           const [owner] = await ethers.getSigners();
           await this.erc721a['safeMint(address,uint256)'](owner.address, 3);
           await this.erc721a['safeMint(address,uint256)'](owner.address, 2);
diff --git a/test/extensions/ERC721AQueryable.test.js b/test/extensions/ERC721AQueryable.test.js
index 82ac6062a..bb4007cc2 100644
--- a/test/extensions/ERC721AQueryable.test.js
+++ b/test/extensions/ERC721AQueryable.test.js
@@ -163,8 +163,6 @@ const createTestSuite = ({ contract, constructorArgs }) =>
           });
 
           it('with direct set burn bit', async function () {
-            // This test is only for the non-upgradeable variant.
-            if (this.erc721aQueryable.isUpgradeable) return;
             await this.erc721aQueryable['safeMint(address,uint256)'](this.addr3.address, 3);
             await this.erc721aQueryable['safeMint(address,uint256)'](this.addr3.address, 2);
             const nextTokenId = this.owner.expected.tokens[this.owner.expected.tokens.length - 1].add(1);
@@ -233,7 +231,6 @@ const createTestSuite = ({ contract, constructorArgs }) =>
           });
 
           it('with direct set burn bit', async function () {
-            if (this.erc721aQueryable.isUpgradeable) return;
             await this.erc721aQueryable['safeMint(address,uint256)'](this.addr3.address, 3);
             await this.erc721aQueryable['safeMint(address,uint256)'](this.addr3.address, 2);
             const nextTokenId = this.owner.expected.tokens[this.owner.expected.tokens.length - 1].add(1);
@@ -272,7 +269,6 @@ const createTestSuite = ({ contract, constructorArgs }) =>
           });
 
           it('with direct set burn bit', async function () {
-            if (this.erc721aQueryable.isUpgradeable) return;
             await this.erc721aQueryable.directSetBurnBit(this.addr3.expected.tokens[0]);
             for (let i = 0; i < this.addr3.expected.tokens.length; ++i) {
               const explicitOwnership = await this.erc721aQueryable.explicitOwnershipOf(this.addr3.expected.tokens[i]);
diff --git a/test/helpers.js b/test/helpers.js
index b726532e1..a3d697254 100644
--- a/test/helpers.js
+++ b/test/helpers.js
@@ -3,16 +3,13 @@ const { BigNumber } = require('ethers');
 
 const deployContract = async function (contractName, constructorArgs) {
   let factory;
-  let isUpgradeable = false;
   try {
     factory = await ethers.getContractFactory(contractName);
   } catch (e) {
     factory = await ethers.getContractFactory(contractName + 'UpgradeableWithInit');
-    isUpgradeable = true;
   }
   let contract = await factory.deploy(...(constructorArgs || []));
   await contract.deployed();
-  contract.isUpgradeable = isUpgradeable;
   return contract;
 };