Move _setOwnersExplicit to separate extension

contracts/ERC721A.sol

@@ -34,7 +34,7 @@ contract ERC721A is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
         uint128 numberMinted;
     }
 
-    uint256 private currentIndex = 0;
+    uint256 internal currentIndex = 0;
 
     uint256 internal immutable maxBatchSize;
 
@@ -46,7 +46,7 @@ contract ERC721A is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
 
     // Mapping from token ID to ownership details
     // An empty struct value does not necessarily mean the token is unowned. See ownershipOf implementation for details.
-    mapping(uint256 => TokenOwnership) private _ownerships;
+    mapping(uint256 => TokenOwnership) internal _ownerships;
 
     // Mapping owner address to address data
     mapping(address => AddressData) private _addressData;
@@ -395,29 +395,6 @@ contract ERC721A is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
         emit Approval(owner, to, tokenId);
     }
 
-    uint256 public nextOwnerToExplicitlySet = 0;
-
-    /**
-     * @dev Explicitly set `owners` to eliminate loops in future calls of ownerOf().
-     */
-    function _setOwnersExplicit(uint256 quantity) internal {
-        uint256 oldNextOwnerToSet = nextOwnerToExplicitlySet;
-        require(quantity > 0, 'quantity must be nonzero');
-        uint256 endIndex = oldNextOwnerToSet + quantity - 1;
-        if (endIndex > currentIndex - 1) {
-            endIndex = currentIndex - 1;
-        }
-        // We know if the last one in the group exists, all in the group exist, due to serial ordering.
-        require(_exists(endIndex), 'not enough minted yet for this cleanup');
-        for (uint256 i = oldNextOwnerToSet; i <= endIndex; i++) {
-            if (_ownerships[i].addr == address(0)) {
-                TokenOwnership memory ownership = ownershipOf(i);
-                _ownerships[i] = TokenOwnership(ownership.addr, ownership.startTimestamp);
-            }
-        }
-        nextOwnerToExplicitlySet = endIndex + 1;
-    }
-
     /**
      * @dev Internal function to invoke {IERC721Receiver-onERC721Received} on a target address.
      * The call is not executed if the target address is not a contract.

contracts/extensions/ERC721AOwnersExplicit.sol

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: MIT
+// Creator: Chiru Labs
+
+pragma solidity ^0.8.0;
+
+import '../ERC721A.sol';
+
+abstract contract ERC721AOwnersExplicit is ERC721A {
+    uint256 public nextOwnerToExplicitlySet = 0;
+
+    /**
+     * @dev Explicitly set `owners` to eliminate loops in future calls of ownerOf().
+     */
+    function _setOwnersExplicit(uint256 quantity) internal {
+        require(quantity > 0, 'quantity must be nonzero');
+        require(currentIndex > 0, 'no tokens minted yet');
+        require(nextOwnerToExplicitlySet < currentIndex, 'all ownerships have been set');
+
+        uint256 oldNextOwnerToSet = nextOwnerToExplicitlySet;
+        uint256 endIndex = oldNextOwnerToSet + quantity - 1;
+        // Set the end index to be the last token index
+        if (endIndex > currentIndex - 1) {
+            endIndex = currentIndex - 1;
+        }
+
+        for (uint256 i = oldNextOwnerToSet; i <= endIndex; i++) {
+            if (_ownerships[i].addr == address(0)) {
+                TokenOwnership memory ownership = ownershipOf(i);
+                _ownerships[i] = TokenOwnership(ownership.addr, ownership.startTimestamp);
+            }
+        }
+        nextOwnerToExplicitlySet = endIndex + 1;
+    }
+}

contracts/mocks/ERC721AExplicitOwnershipMock.sol

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: MIT
+// Creators: Chiru Labs
+
+pragma solidity ^0.8.0;
+
+import '../extensions/ERC721AOwnersExplicit.sol';
+
+contract ERC721AOwnersExplicitMock is ERC721AOwnersExplicit {
+    constructor(
+        string memory name_,
+        string memory symbol_,
+        uint256 maxBatchSize_
+    ) ERC721A(name_, symbol_, maxBatchSize_) {}
+
+    function safeMint(address to, uint256 quantity) public {
+        _safeMint(to, quantity);
+    }
+
+    function setOwnersExplicit(uint256 quantity) public {
+        _setOwnersExplicit(quantity);
+    }
+
+    function getOwnershipAt(uint256 index) public view returns (TokenOwnership memory) {
+        return _ownerships[index];
+    }
+}

test/extensions/ERC721AOwnersExplicit.test.js

@@ -0,0 +1,92 @@
+const { expect } = require('chai');
+const { constants } = require('@openzeppelin/test-helpers');
+const { ZERO_ADDRESS } = constants;
+
+
+describe('ERC721AOwnersExplicit', function () {
+  beforeEach(async function () {
+    this.ERC721AOwnersExplicit = await ethers.getContractFactory('ERC721AOwnersExplicitMock');
+    this.token = await this.ERC721AOwnersExplicit.deploy('Azuki', 'AZUKI', 5);
+    await this.token.deployed();
+  });
+
+  context('with no minted tokens', async function () {
+    it('does not have enough tokens minted', async function () {
+      await expect(
+        this.token.setOwnersExplicit(1)
+      ).to.be.revertedWith('no tokens minted yet');
+    });
+  });
+
+  context('with minted tokens', async function () {
+    beforeEach(async function () {
+      const [owner, addr1, addr2, addr3] = await ethers.getSigners();
+      this.owner = owner;
+      this.addr1 = addr1;
+      this.addr2 = addr2;
+      this.addr3 = addr3;
+      // After the following mints, our ownership array will look like this:
+      // | 1 | 2 | Empty | 3 | Empty | Empty |
+      await this.token['safeMint(address,uint256)'](addr1.address, 1);
+      await this.token['safeMint(address,uint256)'](addr2.address, 2);
+      await this.token['safeMint(address,uint256)'](addr3.address, 3);
+    });
+
+    describe('setOwnersExplicit', async function () {
+      it('rejects 0 quantity', async function () {
+        await expect(
+          this.token.setOwnersExplicit(0)
+        ).to.be.revertedWith('quantity must be nonzero');
+      });
+
+      it('handles single increment properly', async function () {
+        await this.token.setOwnersExplicit(1);
+        expect(await this.token.nextOwnerToExplicitlySet()).to.equal('1');
+      });
+
+      it('properly sets the ownership of index 2', async function () {
+        let ownerAtTwo = await this.token.getOwnershipAt(2);
+        expect(ownerAtTwo[0]).to.equal(ZERO_ADDRESS);
+        await this.token.setOwnersExplicit(3);
+        ownerAtTwo = await this.token.getOwnershipAt(2);
+        expect(ownerAtTwo[0]).to.equal(this.addr2.address);
+        expect(await this.token.nextOwnerToExplicitlySet()).to.equal('3');
+      });
+
+      it('sets all ownerships in one go', async function () {
+        await this.token.setOwnersExplicit(6);
+        for (let tokenId = 0; tokenId < 6; tokenId++) {
+          let owner = await this.token.getOwnershipAt(tokenId);
+          expect(owner[0]).to.not.equal(ZERO_ADDRESS);
+        }
+      });
+
+      it('sets all ownerships with overflowing quantity', async function () {
+        await this.token.setOwnersExplicit(15);
+        for (let tokenId = 0; tokenId < 6; tokenId++) {
+          let owner = await this.token.getOwnershipAt(tokenId);
+          expect(owner[0]).to.not.equal(ZERO_ADDRESS);
+        }
+      });
+
+      it('sets all ownerships in multiple calls', async function () {
+        await this.token.setOwnersExplicit(2);
+        expect(await this.token.nextOwnerToExplicitlySet()).to.equal('2');
+        await this.token.setOwnersExplicit(1);
+        expect(await this.token.nextOwnerToExplicitlySet()).to.equal('3');
+        await this.token.setOwnersExplicit(3);
+        for (let tokenId = 0; tokenId < 6; tokenId++) {
+          let owner = await this.token.getOwnershipAt(tokenId);
+          expect(owner[0]).to.not.equal(ZERO_ADDRESS);
+        }
+      });
+
+      it('rejects after all ownerships have been set', async function () {
+        await this.token.setOwnersExplicit(6);
+        await expect(
+          this.token.setOwnersExplicit(1)
+        ).to.be.revertedWith('all ownerships have been set');
+      });
+    });
+  });
+});