Add EnumerableMap, refactor ERC721

contracts/mocks/ERC721Mock.sol

@@ -13,10 +13,6 @@ contract ERC721Mock is ERC721 {
         return _exists(tokenId);
     }
 
-    function tokensOfOwner(address owner) public view returns (uint256[] memory) {
-        return _tokensOfOwner(owner);
-    }
-
     function setTokenURI(uint256 tokenId, string memory uri) public {
         _setTokenURI(tokenId, uri);
     }

contracts/mocks/EnumerableMapMock.sol

@@ -0,0 +1,38 @@
+pragma solidity ^0.6.0;
+
+import "../utils/EnumerableMap.sol";
+
+contract EnumerableMapMock {
+    using EnumerableMap for EnumerableMap.UintToAddressMap;
+
+    event OperationResult(bool result);
+
+    EnumerableMap.UintToAddressMap private _map;
+
+    function contains(uint256 key) public view returns (bool) {
+        return _map.contains(key);
+    }
+
+    function set(uint256 key, address value) public {
+        bool result = _map.set(key, value);
+        emit OperationResult(result);
+    }
+
+    function remove(uint256 key) public {
+        bool result = _map.remove(key);
+        emit OperationResult(result);
+    }
+
+    function length() public view returns (uint256) {
+        return _map.length();
+    }
+
+    function at(uint256 index) public view returns (uint256 key, address value) {
+        return _map.at(index);
+    }
+
+
+    function get(uint256 key) public view returns (address) {
+        return _map.get(key);
+    }
+}

contracts/mocks/EnumerableSetMock.sol

@@ -5,7 +5,7 @@ import "../utils/EnumerableSet.sol";
 contract EnumerableSetMock {
     using EnumerableSet for EnumerableSet.AddressSet;
 
-    event TransactionResult(bool result);
+    event OperationResult(bool result);
 
     EnumerableSet.AddressSet private _set;
 
@@ -15,16 +15,12 @@ contract EnumerableSetMock {
 
     function add(address value) public {
         bool result = _set.add(value);
-        emit TransactionResult(result);
+        emit OperationResult(result);
     }
 
     function remove(address value) public {
         bool result = _set.remove(value);
-        emit TransactionResult(result);
-    }
-
-    function enumerate() public view returns (address[] memory) {
-        return _set.enumerate();
+        emit OperationResult(result);
     }
 
     function length() public view returns (uint256) {

contracts/token/ERC721/ERC721.sol

@@ -8,7 +8,8 @@ import "./IERC721Receiver.sol";
 import "../../introspection/ERC165.sol";
 import "../../math/SafeMath.sol";
 import "../../utils/Address.sol";
-import "../../utils/Counters.sol";
+import "../../utils/EnumerableSet.sol";
+import "../../utils/EnumerableMap.sol";
 
 /**
  * @title ERC721 Non-Fungible Token Standard basic implementation
@@ -17,21 +18,22 @@ import "../../utils/Counters.sol";
 contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable {
     using SafeMath for uint256;
     using Address for address;
-    using Counters for Counters.Counter;
+    using EnumerableSet for EnumerableSet.Uint256Set;
+    using EnumerableMap for EnumerableMap.UintToAddressMap;
 
     // Equals to `bytes4(keccak256("onERC721Received(address,address,uint256,bytes)"))`
     // which can be also obtained as `IERC721Receiver(0).onERC721Received.selector`
     bytes4 private constant _ERC721_RECEIVED = 0x150b7a02;
 
-    // Mapping from token ID to owner
-    mapping (uint256 => address) private _tokenOwner;
+    // Mapping from holder address to their (enumerable) set of owned tokens
+    mapping (address => EnumerableSet.Uint256Set) private _holderTokens;
+
+    // Enumerable mapping from token ids to their owners
+    EnumerableMap.UintToAddressMap private _tokenOwners;
 
     // Mapping from token ID to approved address
     mapping (uint256 => address) private _tokenApprovals;
 
-    // Mapping from owner to number of owned token
-    mapping (address => Counters.Counter) private _ownedTokensCount;
-
     // Mapping from owner to operator approvals
     mapping (address => mapping (address => bool)) private _operatorApprovals;
 
@@ -47,18 +49,6 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
     // Base URI
     string private _baseURI;
 
-    // Mapping from owner to list of owned token IDs
-    mapping(address => uint256[]) private _ownedTokens;
-
-    // Mapping from token ID to index of the owner tokens list
-    mapping(uint256 => uint256) private _ownedTokensIndex;
-
-    // Array with all token ids, used for enumeration
-    uint256[] private _allTokens;
-
-    // Mapping from token id to position in the allTokens array
-    mapping(uint256 => uint256) private _allTokensIndex;
-
     /*
      *     bytes4(keccak256('balanceOf(address)')) == 0x70a08231
      *     bytes4(keccak256('ownerOf(uint256)')) == 0x6352211e
@@ -111,7 +101,7 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
     function balanceOf(address owner) public view override returns (uint256) {
         require(owner != address(0), "ERC721: balance query for the zero address");
 
-        return _ownedTokensCount[owner].current();
+        return _holderTokens[owner].length();
     }
 
     /**
@@ -120,10 +110,7 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
      * @return address currently marked as the owner of the given token ID
      */
     function ownerOf(uint256 tokenId) public view override returns (address) {
-        address owner = _tokenOwner[tokenId];
-        require(owner != address(0), "ERC721: owner query for nonexistent token");
-
-        return owner;
+        return _tokenOwners.get(tokenId, "ERC721: owner query for nonexistent token");
     }
 
     /**
@@ -180,16 +167,16 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
      * @return uint256 token ID at the given index of the tokens list owned by the requested address
      */
     function tokenOfOwnerByIndex(address owner, uint256 index) public view override returns (uint256) {
-        require(index < balanceOf(owner), "ERC721Enumerable: owner index out of bounds");
-        return _ownedTokens[owner][index];
+        return _holderTokens[owner].at(index);
     }
 
     /**
      * @dev Gets the total amount of tokens stored by the contract.
      * @return uint256 representing the total amount of tokens
      */
     function totalSupply() public view override returns (uint256) {
-        return _allTokens.length;
+        // _tokenOwners are indexed by tokenIds, so .length() returns the number of tokenIds
+        return _tokenOwners.length();
     }
 
     /**
@@ -199,8 +186,8 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
      * @return uint256 token ID at the given index of the tokens list
      */
     function tokenByIndex(uint256 index) public view override returns (uint256) {
-        require(index < totalSupply(), "ERC721Enumerable: global index out of bounds");
-        return _allTokens[index];
+        (uint256 tokenId, ) = _tokenOwners.at(index);
+        return tokenId;
     }
 
     /**
@@ -327,8 +314,7 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
      * @return bool whether the token exists
      */
     function _exists(uint256 tokenId) internal view returns (bool) {
-        address owner = _tokenOwner[tokenId];
-        return owner != address(0);
+        return _tokenOwners.contains(tokenId);
     }
 
     /**
@@ -386,11 +372,9 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
 
         _beforeTokenTransfer(address(0), to, tokenId);
 
-        _addTokenToOwnerEnumeration(to, tokenId);
-        _addTokenToAllTokensEnumeration(tokenId);
+        _holderTokens[to].add(tokenId);
 
-        _tokenOwner[tokenId] = to;
-        _ownedTokensCount[to].increment();
+        _tokenOwners.set(tokenId, to);
 
         emit Transfer(address(0), to, tokenId);
     }
@@ -405,22 +389,17 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
 
         _beforeTokenTransfer(owner, address(0), tokenId);
 
+        // Clear approvals
+        _approve(address(0), tokenId);
+
         // Clear metadata (if any)
         if (bytes(_tokenURIs[tokenId]).length != 0) {
             delete _tokenURIs[tokenId];
         }
 
-        _removeTokenFromOwnerEnumeration(owner, tokenId);
-        // Since tokenId will be deleted, we can clear its slot in _ownedTokensIndex to trigger a gas refund
-        _ownedTokensIndex[tokenId] = 0;
-
-        _removeTokenFromAllTokensEnumeration(tokenId);
-
-        // Clear approvals
-        _approve(address(0), tokenId);
+        _holderTokens[owner].remove(tokenId);
 
-        _ownedTokensCount[owner].decrement();
-        _tokenOwner[tokenId] = address(0);
+        _tokenOwners.remove(tokenId);
 
         emit Transfer(owner, address(0), tokenId);
     }
@@ -438,16 +417,13 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
 
         _beforeTokenTransfer(from, to, tokenId);
 
-        _removeTokenFromOwnerEnumeration(from, tokenId);
-        _addTokenToOwnerEnumeration(to, tokenId);
-
-        // Clear approvals
+        // Clear approvals from the previous owner
         _approve(address(0), tokenId);
 
-        _ownedTokensCount[from].decrement();
-        _ownedTokensCount[to].increment();
+        _holderTokens[from].remove(tokenId);
+        _holderTokens[to].add(tokenId);
 
-        _tokenOwner[tokenId] = to;
+        _tokenOwners.set(tokenId, to);
 
         emit Transfer(from, to, tokenId);
     }
@@ -474,15 +450,6 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
         _baseURI = baseURI_;
     }
 
-    /**
-     * @dev Gets the list of token IDs of the requested owner.
-     * @param owner address owning the tokens
-     * @return uint256[] List of token IDs owned by the requested address
-     */
-    function _tokensOfOwner(address owner) internal view returns (uint256[] storage) {
-        return _ownedTokens[owner];
-    }
-
     /**
      * @dev Internal function to invoke {IERC721Receiver-onERC721Received} on a target address.
      * The call is not executed if the target address is not a contract.
@@ -528,81 +495,6 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
         emit Approval(ownerOf(tokenId), to, tokenId);
     }
 
-    /**
-     * @dev Private function to add a token to this extension's ownership-tracking data structures.
-     * @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 _addTokenToOwnerEnumeration(address to, uint256 tokenId) private {
-        _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);
-    }
-
-    /**
-     * @dev Private function to remove a token from this extension's ownership-tracking data structures. Note that
-     * while the token is not assigned a new owner, the `_ownedTokensIndex` mapping is _not_ updated: this allows for
-     * gas optimizations e.g. when performing a transfer operation (avoiding double writes).
-     * This has O(1) time complexity, but alters the order of the _ownedTokens array.
-     * @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 _removeTokenFromOwnerEnumeration(address from, uint256 tokenId) private {
-        // To prevent a gap in from's 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 = _ownedTokens[from].length.sub(1);
-        uint256 tokenIndex = _ownedTokensIndex[tokenId];
-
-        // When the token to delete is the last token, the swap operation is unnecessary
-        if (tokenIndex != lastTokenIndex) {
-            uint256 lastTokenId = _ownedTokens[from][lastTokenIndex];
-
-            _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
-        }
-
-        // Deletes the contents at the last position of the array
-        _ownedTokens[from].pop();
-
-        // Note that _ownedTokensIndex[tokenId] hasn't been cleared: it still points to the old slot (now occupied by
-        // lastTokenId, 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
-
-        // Delete the contents at the last position of the array
-        _allTokens.pop();
-
-        _allTokensIndex[tokenId] = 0;
-    }
-
     /**
      * @dev Hook that is called before any token transfer. This includes minting
      * and burning.