ERC721 bugfix + gas optimizations (#1549)

* Now only swapping when needed.

* Removed _addTokenTo and _removeTokenFrom

* Removed removeTokenFrom test.

* Added tests for ERC721 _mint and _burn

* _burn now uses the same swap and pop mechanism as _removeFromOwner

* Gas optimization on burn

contracts/mocks/ERC721FullMock.sol

@@ -8,7 +8,7 @@ import "../token/ERC721/ERC721Burnable.sol";
 /**
  * @title ERC721FullMock
  * This mock just provides public functions for setting metadata URI, getting all tokens of an owner,
- * checking token existence, removal of a token from an address 
+ * checking token existence, removal of a token from an address
  */
 contract ERC721FullMock is ERC721Full, ERC721Mintable, ERC721MetadataMintable, ERC721Burnable {
     constructor (string name, string symbol) public ERC721Mintable() ERC721Full(name, symbol) {}
@@ -24,8 +24,4 @@ contract ERC721FullMock is ERC721Full, ERC721Mintable, ERC721MetadataMintable, E
     function setTokenURI(uint256 tokenId, string uri) public {
         _setTokenURI(tokenId, uri);
     }
-
-    function removeTokenFrom(address from, uint256 tokenId) public {
-        _removeTokenFrom(from, tokenId);
-    }
 }

contracts/mocks/ERC721Mock.sol

@@ -11,6 +11,10 @@ contract ERC721Mock is ERC721 {
         _mint(to, tokenId);
     }
 
+    function burn(address owner, uint256 tokenId) public {
+        _burn(owner, tokenId);
+    }
+
     function burn(uint256 tokenId) public {
         _burn(tokenId);
     }

contracts/token/ERC721/ERC721.sol

@@ -202,7 +202,11 @@ contract ERC721 is ERC165, IERC721 {
      */
     function _mint(address to, uint256 tokenId) internal {
         require(to != address(0));
-        _addTokenTo(to, tokenId);
+        require(!_exists(tokenId));
+
+        _tokenOwner[tokenId] = to;
+        _ownedTokensCount[to] = _ownedTokensCount[to].add(1);
+
         emit Transfer(address(0), to, tokenId);
     }
 
@@ -214,11 +218,16 @@ contract ERC721 is ERC165, IERC721 {
      * @param tokenId uint256 ID of the token being burned
      */
     function _burn(address owner, uint256 tokenId) internal {
+        require(ownerOf(tokenId) == owner);
+
         _clearApproval(tokenId);
-        _removeTokenFrom(owner, tokenId);
+
+        _ownedTokensCount[owner] = _ownedTokensCount[owner].sub(1);
+        _tokenOwner[tokenId] = address(0);
+
         emit Transfer(owner, address(0), tokenId);
     }
-
+    
     /**
      * @dev Internal function to burn a specific token
      * Reverts if the token does not exist
@@ -228,33 +237,6 @@ contract ERC721 is ERC165, IERC721 {
         _burn(ownerOf(tokenId), tokenId);
     }
 
-    /**
-     * @dev Internal function to add a token ID to the list of a given address
-     * Note that this function is left internal to make ERC721Enumerable possible, but is not
-     * intended to be called by custom derived contracts: in particular, it emits no Transfer event.
-     * @param to address representing the new owner of the given token ID
-     * @param tokenId uint256 ID of the token to be added to the tokens list of the given address
-     */
-    function _addTokenTo(address to, uint256 tokenId) internal {
-        require(_tokenOwner[tokenId] == address(0));
-        _tokenOwner[tokenId] = to;
-        _ownedTokensCount[to] = _ownedTokensCount[to].add(1);
-    }
-
-    /**
-     * @dev Internal function to remove a token ID from the list of a given address
-     * Note that this function is left internal to make ERC721Enumerable possible, but is not
-     * intended to be called by custom derived contracts: in particular, it emits no Transfer event,
-     * and doesn't clear approvals.
-     * @param from address representing the previous owner of the given token ID
-     * @param tokenId uint256 ID of the token to be removed from the tokens list of the given address
-     */
-    function _removeTokenFrom(address from, uint256 tokenId) internal {
-        require(ownerOf(tokenId) == from);
-        _ownedTokensCount[from] = _ownedTokensCount[from].sub(1);
-        _tokenOwner[tokenId] = address(0);
-    }
-
     /**
      * @dev Internal function to transfer ownership of a given token ID to another address.
      * As opposed to transferFrom, this imposes no restrictions on msg.sender.

contracts/token/ERC721/ERC721Enumerable.sol

@@ -67,37 +67,6 @@ contract ERC721Enumerable is ERC165, ERC721, IERC721Enumerable {
         return _allTokens[index];
     }
 
-    /**
-     * @dev Internal function to add a token ID to the list of a given address
-     * This function is internal due to language limitations, see the note in ERC721.sol.
-     * It is not intended to be called by custom derived contracts: in particular, it emits no Transfer event.
-     * @param to address representing the new owner of the given token ID
-     * @param tokenId uint256 ID of the token to be added to the tokens list of the given address
-     */
-    function _addTokenTo(address to, uint256 tokenId) internal {
-        super._addTokenTo(to, tokenId);
-
-        _addTokenToOwnerEnumeration(to, tokenId);
-    }
-
-    /**
-     * @dev Internal function to remove a token ID from the list of a given address
-     * This function is internal due to language limitations, see the note in ERC721.sol.
-     * It is not intended to be called by custom derived contracts: in particular, it emits no Transfer event,
-     * and doesn't clear approvals.
-     * @param from address representing the previous owner of the given token ID
-     * @param tokenId uint256 ID of the token to be removed from the tokens list of the given address
-     */
-    function _removeTokenFrom(address from, uint256 tokenId) internal {
-        super._removeTokenFrom(from, tokenId);
-
-        _removeTokenFromOwnerEnumeration(from, tokenId);
-
-        // Since the token is being destroyed, we also clear its index
-        // TODO(nventuro): 0 is still a valid index, so arguably this isnt really helpful, remove?
-        _ownedTokensIndex[tokenId] = 0;
-    }
-
     /**
      * @dev Internal function to transfer ownership of a given token ID to another address.
      * As opposed to transferFrom, this imposes no restrictions on msg.sender.
@@ -122,8 +91,9 @@ contract ERC721Enumerable is ERC165, ERC721, IERC721Enumerable {
     function _mint(address to, uint256 tokenId) internal {
         super._mint(to, tokenId);
 
-        _allTokensIndex[tokenId] = _allTokens.length;
-        _allTokens.push(tokenId);
+        _addTokenToOwnerEnumeration(to, tokenId);
+
+        _addTokenToAllTokensEnumeration(tokenId);
     }
 
     /**
@@ -136,17 +106,11 @@ contract ERC721Enumerable is ERC165, ERC721, IERC721Enumerable {
     function _burn(address owner, uint256 tokenId) internal {
         super._burn(owner, tokenId);
 
-        // Reorg all tokens array
-        uint256 tokenIndex = _allTokensIndex[tokenId];
-        uint256 lastTokenIndex = _allTokens.length.sub(1);
-        uint256 lastToken = _allTokens[lastTokenIndex];
-
-        _allTokens[tokenIndex] = lastToken;
-        _allTokens[lastTokenIndex] = 0;
+        _removeTokenFromOwnerEnumeration(owner, tokenId);
+        // Since tokenId will be deleted, we can clear its slot in _ownedTokensIndex to trigger a gas refund
+        _ownedTokensIndex[tokenId] = 0;
 
-        _allTokens.length--;
-        _allTokensIndex[tokenId] = 0;
-        _allTokensIndex[lastToken] = tokenIndex;
+        _removeTokenFromAllTokensEnumeration(tokenId);
     }
 
     /**
@@ -164,9 +128,17 @@ contract ERC721Enumerable is ERC165, ERC721, IERC721Enumerable {
      * @param tokenId uint256 ID of the token to be added to the tokens list of the given address
      */
     function _addTokenToOwnerEnumeration(address to, uint256 tokenId) private {
-        uint256 newOwnedTokensLength = _ownedTokens[to].push(tokenId);
-        // No need to use SafeMath since the length after a push cannot be zero
-        _ownedTokensIndex[tokenId] = newOwnedTokensLength - 1;
+        _ownedTokensIndex[tokenId] = _ownedTokens[to].length;
+        _ownedTokens[to].push(tokenId);
+    }
+
+    /**
+     * @dev Private function to add a token to this extension's token tracking data structures.
+     * @param tokenId uint256 ID of the token to be added to the tokens list
+     */
+    function _addTokenToAllTokensEnumeration(uint256 tokenId) private {
+        _allTokensIndex[tokenId] = _allTokens.length;
+        _allTokens.push(tokenId);
     }
 
     /**
@@ -182,21 +154,45 @@ contract ERC721Enumerable is ERC165, ERC721, IERC721Enumerable {
         // then delete the last slot (swap and pop).
 
         uint256 lastTokenIndex = _ownedTokens[from].length.sub(1);
-        uint256 lastTokenId = _ownedTokens[from][lastTokenIndex];
-
         uint256 tokenIndex = _ownedTokensIndex[tokenId];
 
-        _ownedTokens[from][tokenIndex] = lastTokenId; // Move the last token to the slot of the to-delete token
-        _ownedTokensIndex[lastTokenId] = tokenIndex; // Update the moved token's index
+        // When the token to delete is the last token, the swap operation is unnecessary
+        if (tokenIndex != lastTokenIndex) {
+            uint256 lastTokenId = _ownedTokens[from][lastTokenIndex];
 
-        // Note that this will handle single-element arrays. In that case, both tokenIndex and lastTokenIndex are going
-        // to be zero. The swap operation will therefore have no effect, but the token _will_ be deleted during the
-        // 'pop' operation.
+            _ownedTokens[from][tokenIndex] = lastTokenId; // Move the last token to the slot of the to-delete token
+            _ownedTokensIndex[lastTokenId] = tokenIndex; // Update the moved token's index
+        }
 
         // This also deletes the contents at the last position of the array
         _ownedTokens[from].length--;
 
         // Note that _ownedTokensIndex[tokenId] hasn't been cleared: it still points to the old slot (now occcupied by
-        // lasTokenId).
+        // lasTokenId, or just over the end of the array if the token was the last one).
+    }
+
+    /**
+     * @dev Private function to remove a token from this extension's token tracking data structures.
+     * This has O(1) time complexity, but alters the order of the _allTokens array.
+     * @param tokenId uint256 ID of the token to be removed from the tokens list
+     */
+    function _removeTokenFromAllTokensEnumeration(uint256 tokenId) private {
+        // To prevent a gap in the tokens array, we store the last token in the index of the token to delete, and
+        // then delete the last slot (swap and pop).
+
+        uint256 lastTokenIndex = _allTokens.length.sub(1);
+        uint256 tokenIndex = _allTokensIndex[tokenId];
+
+        // When the token to delete is the last token, the swap operation is unnecessary. However, since this occurs so
+        // rarely (when the last minted token is burnt) that we still do the swap here to avoid the gas cost of adding
+        // an 'if' statement (like in _removeTokenFromOwnerEnumeration)
+        uint256 lastTokenId = _allTokens[lastTokenIndex];
+
+        _allTokens[tokenIndex] = lastTokenId; // Move the last token to the slot of the to-delete token
+        _allTokensIndex[lastTokenId] = tokenIndex; // Update the moved token's index
+
+        // This also deletes the contents at the last position of the array
+        _allTokens.length--;
+        _allTokensIndex[tokenId] = 0;
     }
 }

test/token/ERC721/ERC721.test.js

@@ -1,13 +1,112 @@
-const { shouldBehaveLikeERC721 } = require('./ERC721.behavior');
+require('../../helpers/setup');
+const { ZERO_ADDRESS } = require('../../helpers/constants');
+const expectEvent = require('../../helpers/expectEvent');
+const send = require('../../helpers/send');
+const shouldFail = require('../../helpers/shouldFail');
 
+const { shouldBehaveLikeERC721 } = require('./ERC721.behavior');
 const ERC721Mock = artifacts.require('ERC721Mock.sol');
 
-require('../../helpers/setup');
-
-contract('ERC721', function ([_, creator, ...accounts]) {
+contract('ERC721', function ([_, creator, tokenOwner, anyone, ...accounts]) {
   beforeEach(async function () {
     this.token = await ERC721Mock.new({ from: creator });
   });
 
   shouldBehaveLikeERC721(creator, creator, accounts);
+
+  describe('internal functions', function () {
+    const tokenId = 5042;
+
+    describe('_mint(address, uint256)', function () {
+      it('reverts with a null destination address', async function () {
+        await shouldFail.reverting(this.token.mint(ZERO_ADDRESS, tokenId));
+      });
+
+      context('with minted token', async function () {
+        beforeEach(async function () {
+          ({ logs: this.logs } = await this.token.mint(tokenOwner, tokenId));
+        });
+
+        it('emits a Transfer event', function () {
+          expectEvent.inLogs(this.logs, 'Transfer', { from: ZERO_ADDRESS, to: tokenOwner, tokenId });
+        });
+
+        it('creates the token', async function () {
+          (await this.token.balanceOf(tokenOwner)).should.be.bignumber.equal(1);
+          (await this.token.ownerOf(tokenId)).should.equal(tokenOwner);
+        });
+
+        it('reverts when adding a token id that already exists', async function () {
+          await shouldFail.reverting(this.token.mint(tokenOwner, tokenId));
+        });
+      });
+    });
+
+    describe('_burn(address, uint256)', function () {
+      it('reverts when burning a non-existent token id', async function () {
+        await shouldFail.reverting(send.transaction(this.token, 'burn', 'address,uint256', [tokenOwner, tokenId]));
+      });
+
+      context('with minted token', function () {
+        beforeEach(async function () {
+          await this.token.mint(tokenOwner, tokenId);
+        });
+
+        it('reverts when the account is not the owner', async function () {
+          await shouldFail.reverting(send.transaction(this.token, 'burn', 'address,uint256', [anyone, tokenId]));
+        });
+
+        context('with burnt token', function () {
+          beforeEach(async function () {
+            ({ logs: this.logs } =
+              await send.transaction(this.token, 'burn', 'address,uint256', [tokenOwner, tokenId]));
+          });
+
+          it('emits a Transfer event', function () {
+            expectEvent.inLogs(this.logs, 'Transfer', { from: tokenOwner, to: ZERO_ADDRESS, tokenId });
+          });
+
+          it('deletes the token', async function () {
+            (await this.token.balanceOf(tokenOwner)).should.be.bignumber.equal(0);
+            await shouldFail.reverting(this.token.ownerOf(tokenId));
+          });
+
+          it('reverts when burning a token id that has been deleted', async function () {
+            await shouldFail.reverting(send.transaction(this.token, 'burn', 'address,uint256', [tokenOwner, tokenId]));
+          });
+        });
+      });
+    });
+
+    describe('_burn(uint256)', function () {
+      it('reverts when burning a non-existent token id', async function () {
+        await shouldFail.reverting(send.transaction(this.token, 'burn', 'uint256', [tokenId]));
+      });
+
+      context('with minted token', function () {
+        beforeEach(async function () {
+          await this.token.mint(tokenOwner, tokenId);
+        });
+
+        context('with burnt token', function () {
+          beforeEach(async function () {
+            ({ logs: this.logs } = await send.transaction(this.token, 'burn', 'uint256', [tokenId]));
+          });
+
+          it('emits a Transfer event', function () {
+            expectEvent.inLogs(this.logs, 'Transfer', { from: tokenOwner, to: ZERO_ADDRESS, tokenId });
+          });
+
+          it('deletes the token', async function () {
+            (await this.token.balanceOf(tokenOwner)).should.be.bignumber.equal(0);
+            await shouldFail.reverting(this.token.ownerOf(tokenId));
+          });
+
+          it('reverts when burning a token id that has been deleted', async function () {
+            await shouldFail.reverting(send.transaction(this.token, 'burn', 'uint256', [tokenId]));
+          });
+        });
+      });
+    });
+  });
 });