Preparation for efficient consecutive burns

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.
      */
@@ -346,33 +354,35 @@ 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.
+                    // 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);
                 }
-                // 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 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;
         }
         _revert(OwnerQueryForNonexistentToken.selector);
     }
@@ -1133,4 +1143,4 @@ contract ERC721A is IERC721A {
             revert(0x00, 0x04)
         }
     }
-}
+}

contracts/extensions/ERC721AQueryable.sol

@@ -44,11 +44,13 @@ 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()) {
+                    // If the `tokenId` is within bounds,
+                    // scan backwards for the initialized ownership slot.
+                    while (!_ownershipIsInitialized(tokenId)) --tokenId;
+                    return _ownershipAt(tokenId);
                 }
             }
         }
@@ -109,6 +111,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())`.
@@ -157,8 +191,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 +208,12 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
                                 mstore(add(tokenIds, shl(5, tokenIdsIdx)), start)
                             }
                         }
+                        // Otherwise, reset `currOwnershipAddr`.
+                        // This handles the case of batch burned tokens
+                        // (burned bit of first slot set, remaining slots left uninitialized).
+                        default {
+                            currOwnershipAddr := 0
+                        }
                         start := add(start, 1)
                     }
                 } while (!(start == stop || tokenIdsIdx == tokenIdsMaxLength));
@@ -184,53 +225,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 {
-                // if `ownership.burned == false`.
-                if iszero(mload(add(ownership, 0x40))) {
-                    // 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)
-                    }
-                }
-                i := add(i, 1)
-            }
-        }
-        return tokenIds;
-    }
 }

contracts/mocks/DirectBurnBitSetterHelper.sol

@@ -0,0 +1,28 @@
+// 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 {
+        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.
+        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))
+        }
+    }
+}

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();
     }
 }

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) {

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 {

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();
     }
 }

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();
     }
 }

contracts/mocks/StartTokenIdHelper.sol

@@ -10,9 +10,24 @@ 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_;
+        _initializeStartTokenId(startTokenId_);
+    }
+
+    function startTokenId() public view returns (uint256 result) {
+        assembly {
+            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)
+        }
     }
 }

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');

test/extensions/ERC721AQueryable.test.js

@@ -161,6 +161,17 @@ 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 () {
+            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 +229,17 @@ 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 () {
+            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 +267,14 @@ const createTestSuite = ({ contract, constructorArgs }) =>
             const explicitOwnership = await this.erc721aQueryable.explicitOwnershipOf(this.currentIndex);
             expectExplicitOwnershipNotExists(explicitOwnership);
           });
+
+          it('with direct set burn bit', async function () {
+            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 () {