Add function to directly set extra data

contracts/ERC721A.sol

@@ -59,9 +59,12 @@ contract ERC721A is IERC721A {
     // The bit position of `extraData` in packed ownership.
     uint256 private constant BITPOS_EXTRA_DATA = 232;
 
+    // Mask of all 256 bits in a packed ownership except the 24 bits for `extraData`.
+    uint256 private constant BITMASK_EXTRA_DATA_COMPLEMENT = (1 << 232) - 1;
+
     // The mask of the lower 160 bits for addresses.
     uint256 private constant BITMASK_ADDRESS = (1 << 160) - 1;
-    
+
     // The maximum `quantity` that can be minted with `_mintERC2309`.
     // This limit is to prevent overflows on the address data entries.
     // For a limit of 5000, a total of 3.689e15 calls to `_mintERC2309`
@@ -208,8 +211,8 @@ contract ERC721A is IERC721A {
     function _setAux(address owner, uint64 aux) internal {
         uint256 packed = _packedAddressData[owner];
         uint256 auxCasted;
+        // Cast `aux` with assembly to avoid redundant masking.
         assembly {
-            // Cast aux without masking.
             auxCasted := aux
         }
         packed = (packed & BITMASK_AUX_COMPLEMENT) | (auxCasted << BITPOS_AUX);
@@ -763,6 +766,21 @@ contract ERC721A is IERC721A {
         }
     }
 
+    /**
+     * @dev Directly sets the extra data for the ownership data `index`.
+     */
+    function _setExtraDataAt(uint256 index, uint24 extraData) internal {
+        uint256 packed = _packedOwnerships[index];
+        if (packed == 0) revert OwnershipNotInitializedForExtraData();
+        uint256 extraDataCasted;
+        // Cast `extraData` with assembly to avoid redundant masking.
+        assembly {
+            extraDataCasted := extraData
+        }
+        packed = (packed & BITMASK_EXTRA_DATA_COMPLEMENT) | (extraDataCasted << BITPOS_EXTRA_DATA);
+        _packedOwnerships[index] = packed;
+    }
+
     /**
      * @dev Returns the next extra data for the packed ownership data.
      * The returned result is shifted into position.
@@ -771,12 +789,9 @@ contract ERC721A is IERC721A {
         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;
+    ) private view returns (uint256) {
+        uint24 extraData = uint24(prevOwnershipPacked >> BITPOS_EXTRA_DATA);
+        return uint256(_extraData(from, to, extraData)) << BITPOS_EXTRA_DATA;
     }
 
     /**

contracts/IERC721A.sol

@@ -73,6 +73,11 @@ interface IERC721A {
      */
     error MintERC2309QuantityExceedsLimit();
 
+    /**
+     * The `extraData` cannot be set on an unintialized ownership slot.
+     */
+    error OwnershipNotInitializedForExtraData();
+
     struct TokenOwnership {
         // The address of the owner.
         address addr;

contracts/mocks/ERC721ATransferCounterMock.sol

@@ -22,4 +22,8 @@ contract ERC721ATransferCounterMock is ERC721AMock {
         }
         return previousExtraData + 1;
     }
+
+    function setExtraDataAt(uint256 index, uint24 extraData) public {
+        _setExtraDataAt(index, extraData);
+    }
 }

test/ERC721A.test.js

@@ -724,8 +724,3 @@ describe('ERC721A with ERC2309', async function () {
     await expect(deployContract('ERC721AWithERC2309Mock', args)).to.be.revertedWith('MintZeroQuantity');
   });
 });
-
-describe(
-  'ERC721A override _extraData()',
-  createTestSuite({ contract: 'ERC721ATransferCounterMock', constructorArgs: ['Azuki', 'AZUKI'] })
-);

test/extensions/ERC721ATransferCounter.test.js

@@ -1,19 +1,11 @@
-const { deployContract, offsettedIndex } = require('../helpers.js');
+const { deployContract } = 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 () {
@@ -24,12 +16,12 @@ const createTestSuite = ({ contract, constructorArgs }) =>
 
           this.addr1.expected = {
             balance: 1,
-            tokens: [offsetted(0)],
+            tokens: [0],
           };
 
           this.owner.expected = {
             balance: 2,
-            tokens: offsetted(1, 2),
+            tokens: [1, 2],
           };
 
           this.mintOrder = [this.addr1, this.owner];
@@ -84,6 +76,25 @@ const createTestSuite = ({ contract, constructorArgs }) =>
             }
           });
         });
+
+        describe('setExtraData', function () {
+          it('can set and get the extraData directly', async function () {
+            const extraData = 12345;
+            await this.erc721aCounter.setExtraDataAt(0, extraData);
+            const ownership = await this.erc721aCounter.getOwnershipAt(0);
+            expect(ownership.extraData).to.equal(extraData);
+          });
+
+          it('setting the extraData for uninitialized slot reverts', async function () {
+            const extraData = 12345;
+            await expect(this.erc721aCounter.setExtraDataAt(2, extraData))
+              .to.be.revertedWith('OwnershipNotInitializedForExtraData');
+            await this.erc721aCounter.transferFrom(this.owner.address, this.addr1.address, 2);
+            await this.erc721aCounter.setExtraDataAt(2, extraData);
+            const ownership = await this.erc721aCounter.getOwnershipAt(2);
+            expect(ownership.extraData).to.equal(extraData);
+          });
+        });
       });
     });
   };