Add EnumerableMap, refactor ERC721 (#2160)

* Implement AddressSet in terms of a generic Set

* Add Uint256Set

* Add EnumerableMap

* Fix wording on EnumerableSet docs and tests

* Refactor ERC721 using EnumerableSet and EnumerableMap

* Fix tests

* Fix linter error

* Gas optimization for EnumerableMap

* Gas optimization for EnumerableSet

* Remove often not-taken if from Enumerable data structures

* Fix failing test

* Gas optimization for EnumerableMap

* Fix linter errors

* Add comment for clarification

* Improve test naming

* Rename EnumerableMap.add to set

* Add overload for EnumerableMap.get with custom error message

* Improve Enumerable docs

* Rename Uint256Set to UintSet

* Add changelog entry
This commit is contained in:
Nicolás Venturo
2020-04-02 15:43:06 -03:00
committed by GitHub
parent 0408e51ae6
commit bd0778461d
12 changed files with 606 additions and 252 deletions

View File

@ -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.UintSet;
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.UintSet) 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,8 +167,7 @@ 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);
}
/**
@ -189,7 +175,8 @@ contract ERC721 is Context, ERC165, IERC721, IERC721Metadata, IERC721Enumerable
* @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;
_holderTokens[owner].remove(tokenId);
_removeTokenFromAllTokensEnumeration(tokenId);
// Clear approvals
_approve(address(0), 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.