ERC721 bugfix + gas optimizations

contracts/token/ERC721/ERC721Enumerable.sol

@@ -93,8 +93,7 @@ contract ERC721Enumerable is ERC165, ERC721, IERC721Enumerable {
 
         _addTokenToOwnerEnumeration(to, tokenId);
 
-        _allTokensIndex[tokenId] = _allTokens.length;
-        _allTokens.push(tokenId);
+        _addTokenToAllTokensEnumeration(tokenId);
     }
 
     /**
@@ -111,17 +110,7 @@ contract ERC721Enumerable is ERC165, ERC721, IERC721Enumerable {
         // Since tokenId will be deleted, we can clear its slot in _ownedTokensIndex to trigger a gas refund
         _ownedTokensIndex[tokenId] = 0;
 
-        // Reorg all tokens array
-        uint256 tokenIndex = _allTokensIndex[tokenId];
-        uint256 lastTokenIndex = _allTokens.length.sub(1);
-        uint256 lastToken = _allTokens[lastTokenIndex];
-
-        _allTokens[tokenIndex] = lastToken;
-        _allTokens[lastTokenIndex] = 0;
-
-        _allTokens.length--;
-        _allTokensIndex[tokenId] = 0;
-        _allTokensIndex[lastToken] = tokenIndex;
+        _removeTokenFromAllTokensEnumeration(tokenId);
     }
 
     /**
@@ -139,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);
     }
 
     /**
@@ -173,4 +170,29 @@ contract ERC721Enumerable is ERC165, ERC721, IERC721Enumerable {
         // Note that _ownedTokensIndex[tokenId] hasn't been cleared: it still points to the old slot (now occcupied by
         // 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
+        if (tokenIndex != lastTokenIndex) {
+            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;
+    }
 }