add queryable extension to ERC721Base

contracts/base/ERC721Base.sol

@@ -3,7 +3,7 @@ pragma solidity ^0.8.0;
 
 /// @author thirdweb
 
-import { ERC721A } from "../eip/ERC721AVirtualApprove.sol";
+import "../eip/queryable/ERC721AQueryable.sol";
 
 import "../extension/ContractMetadata.sol";
 import "../extension/Multicall.sol";
@@ -30,7 +30,7 @@ import "../lib/TWStrings.sol";
  *      - EIP 2981 compliance for royalty support on NFT marketplaces.
  */
 
-contract ERC721Base is ERC721A, ContractMetadata, Multicall, Ownable, Royalty, BatchMintMetadata {
+contract ERC721Base is ERC721AQueryable, ContractMetadata, Multicall, Ownable, Royalty, BatchMintMetadata {
     using TWStrings for uint256;
 
     /*//////////////////////////////////////////////////////////////
@@ -89,7 +89,7 @@ contract ERC721Base is ERC721A, ContractMetadata, Multicall, Ownable, Royalty, B
      *
      *  @param _tokenId The tokenId of an NFT.
      */
-    function tokenURI(uint256 _tokenId) public view virtual override returns (string memory) {
+    function tokenURI(uint256 _tokenId) public view virtual override(ERC721A, IERC721Metadata) returns (string memory) {
         string memory fullUriForToken = fullURI[_tokenId];
         if (bytes(fullUriForToken).length > 0) {
             return fullUriForToken;

contracts/eip/queryable/ERC721AQueryable.sol

@@ -0,0 +1,168 @@
+// SPDX-License-Identifier: MIT
+// ERC721A Contracts v3.3.0
+// Creator: Chiru Labs
+
+pragma solidity ^0.8.4;
+
+import "./IERC721AQueryable.sol";
+import "../ERC721AVirtualApprove.sol";
+
+/**
+ * @title ERC721A Queryable
+ * @dev ERC721A subclass with convenience query functions.
+ */
+abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
+    /**
+     * @dev Returns the `TokenOwnership` struct at `tokenId` without reverting.
+     *
+     * If the `tokenId` is out of bounds:
+     *   - `addr` = `address(0)`
+     *   - `startTimestamp` = `0`
+     *   - `burned` = `false`
+     *
+     * If the `tokenId` is burned:
+     *   - `addr` = `<Address of owner before token was burned>`
+     *   - `startTimestamp` = `<Timestamp when token was burned>`
+     *   - `burned = `true`
+     *
+     * Otherwise:
+     *   - `addr` = `<Address of owner>`
+     *   - `startTimestamp` = `<Timestamp of start of ownership>`
+     *   - `burned = `false`
+     */
+    function explicitOwnershipOf(uint256 tokenId) public view override returns (TokenOwnership memory) {
+        TokenOwnership memory ownership;
+        if (tokenId < _startTokenId() || tokenId >= _currentIndex) {
+            return ownership;
+        }
+        ownership = _ownerships[tokenId];
+        if (ownership.burned) {
+            return ownership;
+        }
+        return _ownershipOf(tokenId);
+    }
+
+    /**
+     * @dev Returns an array of `TokenOwnership` structs at `tokenIds` in order.
+     * See {ERC721AQueryable-explicitOwnershipOf}
+     */
+    function explicitOwnershipsOf(uint256[] memory tokenIds) external view override returns (TokenOwnership[] memory) {
+        unchecked {
+            uint256 tokenIdsLength = tokenIds.length;
+            TokenOwnership[] memory ownerships = new TokenOwnership[](tokenIdsLength);
+            for (uint256 i; i != tokenIdsLength; ++i) {
+                ownerships[i] = explicitOwnershipOf(tokenIds[i]);
+            }
+            return ownerships;
+        }
+    }
+
+    /**
+     * @dev Returns an array of token IDs owned by `owner`,
+     * in the range [`start`, `stop`)
+     * (i.e. `start <= tokenId < stop`).
+     *
+     * This function allows for tokens to be queried if the collection
+     * grows too big for a single call of {ERC721AQueryable-tokensOfOwner}.
+     *
+     * Requirements:
+     *
+     * - `start` < `stop`
+     */
+    /* solhint-disable*/
+    function tokensOfOwnerIn(
+        address owner,
+        uint256 start,
+        uint256 stop
+    ) external view override returns (uint256[] memory) {
+        unchecked {
+            if (start >= stop) revert InvalidQueryRange();
+            uint256 tokenIdsIdx;
+            uint256 stopLimit = _currentIndex;
+            // Set `start = max(start, _startTokenId())`.
+            if (start < _startTokenId()) {
+                start = _startTokenId();
+            }
+            // Set `stop = min(stop, _currentIndex)`.
+            if (stop > stopLimit) {
+                stop = stopLimit;
+            }
+            uint256 tokenIdsMaxLength = balanceOf(owner);
+            // Set `tokenIdsMaxLength = min(balanceOf(owner), stop - start)`,
+            // to cater for cases where `balanceOf(owner)` is too big.
+            if (start < stop) {
+                uint256 rangeLength = stop - start;
+                if (rangeLength < tokenIdsMaxLength) {
+                    tokenIdsMaxLength = rangeLength;
+                }
+            } else {
+                tokenIdsMaxLength = 0;
+            }
+            uint256[] memory tokenIds = new uint256[](tokenIdsMaxLength);
+            if (tokenIdsMaxLength == 0) {
+                return tokenIds;
+            }
+            // We need to call `explicitOwnershipOf(start)`,
+            // because the slot at `start` may not be initialized.
+            TokenOwnership memory ownership = explicitOwnershipOf(start);
+            address currOwnershipAddr;
+            // If the starting slot exists (i.e. not burned), initialize `currOwnershipAddr`.
+            // `ownership.address` will not be zero, as `start` is clamped to the valid token ID range.
+            if (!ownership.burned) {
+                currOwnershipAddr = ownership.addr;
+            }
+            for (uint256 i = start; i != stop && tokenIdsIdx != tokenIdsMaxLength; ++i) {
+                ownership = _ownerships[i];
+                if (ownership.burned) {
+                    continue;
+                }
+                if (ownership.addr != address(0)) {
+                    currOwnershipAddr = ownership.addr;
+                }
+                if (currOwnershipAddr == owner) {
+                    tokenIds[tokenIdsIdx++] = i;
+                }
+            }
+            // Downsize the array to fit.
+            assembly {
+                mstore(tokenIds, tokenIdsIdx)
+            }
+            return tokenIds;
+        }
+    }
+
+    /* solhint-enable */
+
+    /**
+     * @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 pfp collections should be fine).
+     */
+    function tokensOfOwner(address owner) external view override returns (uint256[] memory) {
+        unchecked {
+            uint256 tokenIdsIdx;
+            address currOwnershipAddr;
+            uint256 tokenIdsLength = balanceOf(owner);
+            uint256[] memory tokenIds = new uint256[](tokenIdsLength);
+            TokenOwnership memory ownership;
+            for (uint256 i = _startTokenId(); tokenIdsIdx != tokenIdsLength; ++i) {
+                ownership = _ownerships[i];
+                if (ownership.burned) {
+                    continue;
+                }
+                if (ownership.addr != address(0)) {
+                    currOwnershipAddr = ownership.addr;
+                }
+                if (currOwnershipAddr == owner) {
+                    tokenIds[tokenIdsIdx++] = i;
+                }
+            }
+            return tokenIds;
+        }
+    }
+}

contracts/eip/queryable/IERC721AQueryable.sol

@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: MIT
+// ERC721A Contracts v3.3.0
+// Creator: Chiru Labs
+
+pragma solidity ^0.8.4;
+
+import "../interface/IERC721A.sol";
+
+/**
+ * @dev Interface of an ERC721AQueryable compliant contract.
+ */
+interface IERC721AQueryable is IERC721A {
+    /**
+     * Invalid query range (`start` >= `stop`).
+     */
+    error InvalidQueryRange();
+
+    /**
+     * @dev Returns the `TokenOwnership` struct at `tokenId` without reverting.
+     *
+     * If the `tokenId` is out of bounds:
+     *   - `addr` = `address(0)`
+     *   - `startTimestamp` = `0`
+     *   - `burned` = `false`
+     *
+     * If the `tokenId` is burned:
+     *   - `addr` = `<Address of owner before token was burned>`
+     *   - `startTimestamp` = `<Timestamp when token was burned>`
+     *   - `burned = `true`
+     *
+     * Otherwise:
+     *   - `addr` = `<Address of owner>`
+     *   - `startTimestamp` = `<Timestamp of start of ownership>`
+     *   - `burned = `false`
+     */
+    function explicitOwnershipOf(uint256 tokenId) external view returns (TokenOwnership memory);
+
+    /**
+     * @dev Returns an array of `TokenOwnership` structs at `tokenIds` in order.
+     * See {ERC721AQueryable-explicitOwnershipOf}
+     */
+    function explicitOwnershipsOf(uint256[] memory tokenIds) external view returns (TokenOwnership[] memory);
+
+    /**
+     * @dev Returns an array of token IDs owned by `owner`,
+     * in the range [`start`, `stop`)
+     * (i.e. `start <= tokenId < stop`).
+     *
+     * This function allows for tokens to be queried if the collection
+     * grows too big for a single call of {ERC721AQueryable-tokensOfOwner}.
+     *
+     * Requirements:
+     *
+     * - `start` < `stop`
+     */
+    function tokensOfOwnerIn(
+        address owner,
+        uint256 start,
+        uint256 stop
+    ) external view returns (uint256[] memory);
+
+    /**
+     * @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 pfp collections should be fine).
+     */
+    function tokensOfOwner(address owner) external view returns (uint256[] memory);
+}