Add a generic interface for ownership storage hitchhiking

contracts/ERC721A.sol

@@ -56,6 +56,9 @@ contract ERC721A is IERC721A {
     // The bit mask of the `nextInitialized` bit in packed ownership.
     uint256 private constant BITMASK_NEXT_INITIALIZED = 1 << 225;
 
+    // The bit position of `extraData` in packed ownership.
+    uint256 private constant BITPOS_EXTRA_DATA = 232;
+
     // The tokenId of the next token to be minted.
     uint256 private _currentIndex;
 
@@ -77,6 +80,7 @@ contract ERC721A is IERC721A {
     // - [160..223] `startTimestamp`
     // - [224]      `burned`
     // - [225]      `nextInitialized`
+    // - [232..255] `extraData`
     mapping(uint256 => uint256) private _packedOwnerships;
 
     // Mapping owner address to address data.
@@ -163,7 +167,7 @@ contract ERC721A is IERC721A {
      * @dev See {IERC721-balanceOf}.
      */
     function balanceOf(address owner) public view override returns (uint256) {
-        if (_addressToUint256(owner) == 0) revert BalanceQueryForZeroAddress();
+        if (owner == address(0)) revert BalanceQueryForZeroAddress();
         return _packedAddressData[owner] & BITMASK_ADDRESS_DATA_ENTRY;
     }
 
@@ -239,6 +243,7 @@ contract ERC721A is IERC721A {
         ownership.addr = address(uint160(packed));
         ownership.startTimestamp = uint64(packed >> BITPOS_START_TIMESTAMP);
         ownership.burned = packed & BITMASK_BURNED != 0;
+        ownership.extraData = uint24(packed >> BITPOS_EXTRA_DATA);
     }
 
     /**
@@ -270,6 +275,8 @@ contract ERC721A is IERC721A {
      */
     function _packOwnershipData(address owner, uint256 flags) private view returns (uint256 value) {
         assembly {
+            // Mask `owner` to the lower 160 bits, in case the upper bits somehow aren't clean.
+            owner := and(owner, 0xffffffffffffffffffffffffffffffffffffffff)
             // `owner | (block.timestamp << BITPOS_START_TIMESTAMP) | flags`.
             value := or(owner, or(shl(BITPOS_START_TIMESTAMP, timestamp()), flags))
         }
@@ -315,15 +322,6 @@ contract ERC721A is IERC721A {
         return '';
     }
 
-    /**
-     * @dev Casts the address to uint256 without masking.
-     */
-    function _addressToUint256(address value) private pure returns (uint256 result) {
-        assembly {
-            result := value
-        }
-    }
-
     /**
      * @dev Casts the boolean to uint256 without branching.
      */
@@ -478,7 +476,7 @@ contract ERC721A is IERC721A {
      */
     function _mint(address to, uint256 quantity) internal {
         uint256 startTokenId = _currentIndex;
-        if (_addressToUint256(to) == 0) revert MintToZeroAddress();
+        if (to == address(0)) revert MintToZeroAddress();
         if (quantity == 0) revert MintZeroQuantity();
 
         _beforeTokenTransfers(address(0), to, startTokenId, quantity);
@@ -501,7 +499,7 @@ contract ERC721A is IERC721A {
             // - `nextInitialized` to `quantity == 1`.
             _packedOwnerships[startTokenId] = _packOwnershipData(
                 to,
-                _boolToUint256(quantity == 1) << BITPOS_NEXT_INITIALIZED
+                (_boolToUint256(quantity == 1) << BITPOS_NEXT_INITIALIZED) | _nextExtraData(address(0), to, 0)
             );
 
             uint256 offset = startTokenId;
@@ -541,12 +539,12 @@ contract ERC721A is IERC721A {
             approvedAddress == _msgSenderERC721A());
 
         if (!isApprovedOrOwner) revert TransferCallerNotOwnerNorApproved();
-        if (_addressToUint256(to) == 0) revert TransferToZeroAddress();
+        if (to == address(0)) revert TransferToZeroAddress();
 
         _beforeTokenTransfers(from, to, tokenId, 1);
 
         // Clear approvals from the previous owner.
-        if (_addressToUint256(approvedAddress) != 0) {
+        if (approvedAddress != address(0)) {
             delete _tokenApprovals[tokenId];
         }
 
@@ -563,7 +561,10 @@ contract ERC721A is IERC721A {
             // - `startTimestamp` to the timestamp of transfering.
             // - `burned` to `false`.
             // - `nextInitialized` to `true`.
-            _packedOwnerships[tokenId] = _packOwnershipData(to, BITMASK_NEXT_INITIALIZED);
+            _packedOwnerships[tokenId] = _packOwnershipData(
+                to,
+                BITMASK_NEXT_INITIALIZED | _nextExtraData(from, to, prevOwnershipPacked)
+            );
 
             // If the next slot may not have been initialized (i.e. `nextInitialized == false`) .
             if (prevOwnershipPacked & BITMASK_NEXT_INITIALIZED == 0) {
@@ -617,7 +618,7 @@ contract ERC721A is IERC721A {
         _beforeTokenTransfers(from, address(0), tokenId, 1);
 
         // Clear approvals from the previous owner.
-        if (_addressToUint256(approvedAddress) != 0) {
+        if (approvedAddress != address(0)) {
             delete _tokenApprovals[tokenId];
         }
 
@@ -638,7 +639,10 @@ contract ERC721A is IERC721A {
             // - `startTimestamp` to the timestamp of burning.
             // - `burned` to `true`.
             // - `nextInitialized` to `true`.
-            _packedOwnerships[tokenId] = _packOwnershipData(from, BITMASK_BURNED | BITMASK_NEXT_INITIALIZED);
+            _packedOwnerships[tokenId] = _packOwnershipData(
+                from,
+                (BITMASK_BURNED | BITMASK_NEXT_INITIALIZED) | _nextExtraData(from, address(0), prevOwnershipPacked)
+            );
 
             // If the next slot may not have been initialized (i.e. `nextInitialized == false`) .
             if (prevOwnershipPacked & BITMASK_NEXT_INITIALIZED == 0) {
@@ -693,6 +697,41 @@ contract ERC721A is IERC721A {
         }
     }
 
+    /**
+     * @dev Gets the next extra data for the packed ownership data.
+     */
+    function _nextExtraData(
+        address from,
+        address to,
+        uint256 prevOwnershipPacked
+    ) internal view virtual returns (uint256) {
+        uint24 previousExtraData;
+        assembly {
+            previousExtraData := shr(BITPOS_EXTRA_DATA, prevOwnershipPacked)
+        }
+        return uint256(_extraData(from, to, previousExtraData)) << BITPOS_EXTRA_DATA;
+    }
+
+    /**
+     * @dev Called during each token transfer to set the 24bit `extraData` field.
+     * Intended to be overridden by the cosumer contract.
+     *
+     * `previousExtraData` - the value of `extraData` before transfer.
+     *
+     * Calling conditions:
+     *
+     * - When `from` and `to` are both non-zero, `from`'s `tokenId` will be
+     * transferred to `to`.
+     * - When `from` is zero, `tokenId` will be minted for `to`.
+     * - When `to` is zero, `tokenId` will be burned by `from`.
+     * - `from` and `to` are never both zero.
+     */
+    function _extraData(
+        address from,
+        address to,
+        uint24 previousExtraData
+    ) internal view virtual returns (uint24) {}
+
     /**
      * @dev Hook that is called before a set of serially-ordered token ids are about to be transferred. This includes minting.
      * And also called before burning one token.

contracts/IERC721A.sol

@@ -75,6 +75,8 @@ interface IERC721A {
         uint64 startTimestamp;
         // Whether the token has been burned.
         bool burned;
+        // Arbitrary data similar to `startTimestamp` that can be set through `_extraData`.
+        uint24 extraData;
     }
 
     /**

contracts/extensions/ERC721AQueryable.sol

@@ -19,16 +19,19 @@ abstract contract ERC721AQueryable is ERC721A, IERC721AQueryable {
      *   - `addr` = `address(0)`
      *   - `startTimestamp` = `0`
      *   - `burned` = `false`
+     *   - `extraData` = `0`
      *
      * If the `tokenId` is burned:
      *   - `addr` = `<Address of owner before token was burned>`
      *   - `startTimestamp` = `<Timestamp when token was burned>`
      *   - `burned = `true`
+     *   - `extraData` = `<Extra data when token was burned>`
      *
      * Otherwise:
      *   - `addr` = `<Address of owner>`
      *   - `startTimestamp` = `<Timestamp of start of ownership>`
      *   - `burned = `false`
+     *   - `extraData` = `<Extra data at start of ownership>`
      */
     function explicitOwnershipOf(uint256 tokenId) public view override returns (TokenOwnership memory) {
         TokenOwnership memory ownership;

contracts/mocks/ERC721ATransferCounterMock.sol

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: MIT
+// ERC721A Contracts v4.0.0
+// Creators: Chiru Labs
+
+pragma solidity ^0.8.4;
+
+import './ERC721AMock.sol';
+
+contract ERC721ATransferCounterMock is ERC721AMock {
+    constructor(string memory name_, string memory symbol_) ERC721AMock(name_, symbol_) {}
+
+    function _extraData(
+        address from,
+        address to,
+        uint24 previousExtraData
+    ) internal view virtual override returns (uint24) {
+        if (from == address(0)) {
+            return 42;
+        }
+        if (to == address(0)) {
+            return 1337;
+        }
+        return previousExtraData + 1;
+    }
+}

test/ERC721A.test.js

@@ -259,15 +259,14 @@ const createTestSuite = ({ contract, constructorArgs }) =>
           });
 
           it('approval allows token transfer', async function () {
-            await expect(this.erc721a.connect(this.addr3)
-              .transferFrom(this.addr1.address, this.addr3.address, this.tokenId))
-            .to.be.revertedWith('TransferCallerNotOwnerNorApproved');
+            await expect(
+              this.erc721a.connect(this.addr3).transferFrom(this.addr1.address, this.addr3.address, this.tokenId)
+            ).to.be.revertedWith('TransferCallerNotOwnerNorApproved');
             await this.erc721a.connect(this.addr1).approve(this.addr3.address, this.tokenId);
-            await this.erc721a.connect(this.addr3)
-              .transferFrom(this.addr1.address, this.addr3.address, this.tokenId);
-            await expect(this.erc721a.connect(this.addr1)
-              .transferFrom(this.addr3.address, this.addr1.address, this.tokenId))
-            .to.be.revertedWith('TransferCallerNotOwnerNorApproved');
+            await this.erc721a.connect(this.addr3).transferFrom(this.addr1.address, this.addr3.address, this.tokenId);
+            await expect(
+              this.erc721a.connect(this.addr1).transferFrom(this.addr3.address, this.addr1.address, this.tokenId)
+            ).to.be.revertedWith('TransferCallerNotOwnerNorApproved');
           });
         });
 
@@ -680,3 +679,8 @@ describe(
   'ERC721A override _startTokenId()',
   createTestSuite({ contract: 'ERC721AStartTokenIdMock', constructorArgs: ['Azuki', 'AZUKI', 1] })
 );
+
+describe(
+  'ERC721A override _extraData()',
+  createTestSuite({ contract: 'ERC721ATransferCounterMock', constructorArgs: ['Azuki', 'AZUKI'] })
+);

test/extensions/ERC721AQueryable.test.js

@@ -23,18 +23,21 @@ const createTestSuite = ({ contract, constructorArgs }) =>
         expect(explicitOwnership.burned).to.eql(true);
         expect(explicitOwnership.addr).to.eql(address);
         expect(explicitOwnership.startTimestamp).to.not.eql(BigNumber.from(0));
+        expect(explicitOwnership.extraData).to.equal(BigNumber.from(0));
       };
 
       const expectExplicitOwnershipNotExists = function (explicitOwnership) {
         expect(explicitOwnership.burned).to.eql(false);
         expect(explicitOwnership.addr).to.eql(ZERO_ADDRESS);
         expect(explicitOwnership.startTimestamp).to.eql(BigNumber.from(0));
+        expect(explicitOwnership.extraData).to.equal(BigNumber.from(0));
       };
 
       const expectExplicitOwnershipExists = function (explicitOwnership, address) {
         expect(explicitOwnership.burned).to.eql(false);
         expect(explicitOwnership.addr).to.eql(address);
         expect(explicitOwnership.startTimestamp).to.not.eql(BigNumber.from(0));
+        expect(explicitOwnership.extraData).to.equal(BigNumber.from(0));
       };
 
       context('with no minted tokens', async function () {

test/extensions/ERC721ATransferCounter.test.js

@@ -0,0 +1,97 @@
+const { deployContract, offsettedIndex } = require('../helpers.js');
+const { expect } = require('chai');
+
+const createTestSuite = ({ contract, constructorArgs }) =>
+  function () {
+    let offsetted;
+
+    context(`${contract}`, function () {
+      beforeEach(async function () {
+        this.erc721aCounter = await deployContract(contract, constructorArgs);
+
+        this.startTokenId = this.erc721aCounter.startTokenId
+          ? (await this.erc721aCounter.startTokenId()).toNumber()
+          : 0;
+
+        offsetted = (...arr) => offsettedIndex(this.startTokenId, arr);
+      });
+
+      context('with minted tokens', async function () {
+        beforeEach(async function () {
+          const [owner, addr1] = await ethers.getSigners();
+          this.owner = owner;
+          this.addr1 = addr1;
+
+          this.addr1.expected = {
+            balance: 1,
+            tokens: [offsetted(0)],
+          };
+
+          this.owner.expected = {
+            balance: 2,
+            tokens: offsetted(1, 2),
+          };
+
+          this.mintOrder = [this.addr1, this.owner];
+
+          for (const minter of this.mintOrder) {
+            const balance = minter.expected.balance;
+            if (balance > 0) {
+              await this.erc721aCounter['safeMint(address,uint256)'](minter.address, balance);
+            }
+            // sanity check
+            expect(await this.erc721aCounter.balanceOf(minter.address)).to.equal(minter.expected.balance);
+          }
+        });
+
+        describe('_ownershipOf', function () {
+          it('initial', async function () {
+            for (const minter of this.mintOrder) {
+              for (const tokenId in minter.expected.tokens) {
+                const ownership = await this.erc721aCounter.getOwnershipOf(tokenId);
+                expect(ownership.extraData).to.equal(42);
+              }
+            }
+          });
+
+          it('after a transfer', async function () {
+            await this.erc721aCounter.transferFrom(this.owner.address, this.addr1.address, 1);
+
+            const tests = [
+              { tokenId: 0, expectedData: 42 },
+              { tokenId: 1, expectedData: 43 },
+              { tokenId: 2, expectedData: 42 },
+            ];
+
+            for (const test of tests) {
+              const ownership = await this.erc721aCounter.getOwnershipOf(test.tokenId);
+              expect(ownership.extraData).to.equal(test.expectedData);
+            }
+          });
+
+          it('after a burn', async function () {
+            await this.erc721aCounter['burn(uint256)'](2);
+
+            const tests = [
+              { tokenId: 0, expectedData: 42 },
+              { tokenId: 1, expectedData: 42 },
+              { tokenId: 2, expectedData: 1337 },
+            ];
+
+            for (const test of tests) {
+              const ownership = await this.erc721aCounter.getOwnershipAt(test.tokenId);
+              expect(ownership.extraData).to.equal(test.expectedData);
+            }
+          });
+        });
+      });
+    });
+  };
+
+describe(
+  'ERC721A override _extraData()',
+  createTestSuite({
+    contract: 'ERC721ATransferCounterMock',
+    constructorArgs: ['Azuki', 'AZUKI'],
+  })
+);