Add a generic interface for ownership storage hitchhiking

contracts/ERC721A.sol

@@ -47,14 +47,17 @@ contract ERC721A is IERC721A {
     // The bit position of `startTimestamp` in packed ownership.
     uint256 private constant BITPOS_START_TIMESTAMP = 160;
 
+    // The bit position of `extraData` in packed ownership.
+    uint256 private constant BITPOS_EXTRA_DATA = 224;
+
     // The bit mask of the `burned` bit in packed ownership.
-    uint256 private constant BITMASK_BURNED = 1 << 224;
+    uint256 private constant BITMASK_BURNED = 1 << 248;
 
     // The bit position of the `nextInitialized` bit in packed ownership.
-    uint256 private constant BITPOS_NEXT_INITIALIZED = 225;
+    uint256 private constant BITPOS_NEXT_INITIALIZED = 249;
 
     // The bit mask of the `nextInitialized` bit in packed ownership.
-    uint256 private constant BITMASK_NEXT_INITIALIZED = 1 << 225;
+    uint256 private constant BITMASK_NEXT_INITIALIZED = 1 << 249;
 
     // The tokenId of the next token to be minted.
     uint256 private _currentIndex;
@@ -75,8 +78,9 @@ contract ERC721A is IERC721A {
     // Bits Layout:
     // - [0..159]   `addr`
     // - [160..223] `startTimestamp`
-    // - [224]      `burned`
-    // - [225]      `nextInitialized`
+    // - [224..247] `extraData`
+    // - [248]      `burned`
+    // - [249]      `nextInitialized`
     mapping(uint256 => uint256) private _packedOwnerships;
 
     // Mapping owner address to address data.
@@ -238,9 +242,17 @@ contract ERC721A is IERC721A {
     function _unpackedOwnership(uint256 packed) private pure returns (TokenOwnership memory ownership) {
         ownership.addr = address(uint160(packed));
         ownership.startTimestamp = uint64(packed >> BITPOS_START_TIMESTAMP);
+        ownership.extraData = uint24(packed >> BITPOS_EXTRA_DATA);
         ownership.burned = packed & BITMASK_BURNED != 0;
     }
 
+    /**
+     * Returns the unpacked `extraData` field from `packed`.
+     */
+    function _unpackExtraData(uint256 packed) private pure returns (uint24) {
+        return uint24(packed >> BITPOS_EXTRA_DATA);
+    }
+
     /**
      * Returns the unpacked `TokenOwnership` struct at `index`.
      */
@@ -468,6 +480,7 @@ contract ERC721A is IERC721A {
      */
     function _mint(address to, uint256 quantity) internal {
         uint256 startTokenId = _currentIndex;
+
         if (_addressToUint256(to) == 0) revert MintToZeroAddress();
         if (quantity == 0) revert MintZeroQuantity();
 
@@ -492,6 +505,7 @@ contract ERC721A is IERC721A {
             _packedOwnerships[startTokenId] =
                 _addressToUint256(to) |
                 (block.timestamp << BITPOS_START_TIMESTAMP) |
+                (_extraData(address(0), to, 0) << BITPOS_EXTRA_DATA) |
                 (_boolToUint256(quantity == 1) << BITPOS_NEXT_INITIALIZED);
 
             uint256 offset;
@@ -555,6 +569,7 @@ contract ERC721A is IERC721A {
             _packedOwnerships[tokenId] =
                 _addressToUint256(to) |
                 (block.timestamp << BITPOS_START_TIMESTAMP) |
+                (_extraData(from, to, _unpackExtraData(prevOwnershipPacked)) << BITPOS_EXTRA_DATA) |
                 BITMASK_NEXT_INITIALIZED;
 
             // If the next slot may not have been initialized (i.e. `nextInitialized == false`) .
@@ -633,6 +648,7 @@ contract ERC721A is IERC721A {
             _packedOwnerships[tokenId] =
                 _addressToUint256(from) |
                 (block.timestamp << BITPOS_START_TIMESTAMP) |
+                (_extraData(from, address(0), _unpackExtraData(prevOwnershipPacked)) << BITPOS_EXTRA_DATA) |
                 BITMASK_BURNED |
                 BITMASK_NEXT_INITIALIZED;
 
@@ -689,6 +705,26 @@ contract ERC721A is IERC721A {
         }
     }
 
+    /**
+     * @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 (uint256) {}
+
     /**
      * @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

@@ -73,6 +73,8 @@ interface IERC721A {
         address addr;
         // Keeps track of the start time of ownership with minimal overhead for tokenomics.
         uint64 startTimestamp;
+        // Arbitrary data similar to `startTimestamp` that can be set through `_extraData`.
+        uint24 extraData;
         // Whether the token has been burned.
         bool burned;
     }

contracts/extensions/IERC721AQueryable.sol

@@ -21,16 +21,19 @@ interface IERC721AQueryable is IERC721A {
      * If the `tokenId` is out of bounds:
      *   - `addr` = `address(0)`
      *   - `startTimestamp` = `0`
+     *   - `extraData` = `0`
      *   - `burned` = `false`
      *
      * If the `tokenId` is burned:
      *   - `addr` = `<Address of owner before token was burned>`
      *   - `startTimestamp` = `<Timestamp when token was burned>`
+     *   - `extraData` = `<_extraData when token was burned>`
      *   - `burned = `true`
      *
      * Otherwise:
      *   - `addr` = `<Address of owner>`
      *   - `startTimestamp` = `<Timestamp of start of ownership>`
+     *   - `extraData` = `<_extraData at start of ownership>`
      *   - `burned = `false`
      */
     function explicitOwnershipOf(uint256 tokenId) external view returns (TokenOwnership memory);

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 (uint256) {
+        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'],
+  })
+);