Skip to content

Commit

Permalink
Add EnumerableMap, refactor ERC721 (#2160)
Browse files Browse the repository at this point in the history 
* 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
  • Loading branch information
@nventuro
nventuro committed Apr 2, 2020
1 parent 0408e51 commit bd07784
Show file tree
Hide file tree
Showing 12 changed files with 603 additions and 249 deletions.
3 changes: 3 additions & 0 deletions CHANGELOG.md
View file
Expand Up @@ -5,6 +5,7 @@
### New features
* `AccessControl`: new contract for managing permissions in a system, replacement for `Ownable` and `Roles`. ([#2112](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2112))
* `SafeCast`: new functions to convert to and from signed and unsigned values: `toUint256` and `toInt256`. ([#2123](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2123))
* `EnumerableMap`: a new data structure for key-value pairs (like `mapping`) that can be iterated over. ([#2160](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2160))

### Breaking changes
* `ERC721`: `burn(owner, tokenId)` was removed, use `burn(tokenId)` instead. ([#2125](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2125))
Expand All @@ -30,6 +31,8 @@
* `ERC777`: removed `_callsTokensToSend` and `_callTokensReceived`. ([#2134](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2134))
* `EnumerableSet`: renamed `get` to `at`. ([#2151](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2151))
* `ERC165Checker`: functions no longer have a leading underscore. ([#2150](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2150))
* `ERC721Metadata`, `ERC721Enumerable`: these contracts were removed, and their functionality merged into `ERC721`. ([#2160](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2160))
* `ERC721`: added a constructor for `name` and `symbol`. ([#2160](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2160))
* `ERC20Detailed`: this contract was removed and its functionality merged into `ERC20`. ([#2161](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2161))
* `ERC20`: added a constructor for `name` and `symbol`. `decimals` now defaults to 18. ([#2161](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2161))

Expand Down
4 changes: 0 additions & 4 deletions contracts/mocks/ERC721Mock.sol
View file
Expand Up @@ -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);
}
Expand Down
38 changes: 38 additions & 0 deletions contracts/mocks/EnumerableMapMock.sol
View file
@@ -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);
}
}
10 changes: 3 additions & 7 deletions contracts/mocks/EnumerableSetMock.sol
View file
Expand Up @@ -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;

Expand All @@ -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) {
Expand Down
164 changes: 28 additions & 136 deletions contracts/token/ERC721/ERC721.sol
View file
Expand Up @@ -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
Expand All @@ -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;

Expand All @@ -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
Expand Down Expand Up @@ -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();
}

/**
Expand All @@ -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");
}

/**
Expand Down Expand Up @@ -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();
}

/**
Expand All @@ -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;
}

/**
Expand Down Expand Up @@ -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);
}

/**
Expand Down Expand Up @@ -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);
}
Expand All @@ -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);
}
Expand All @@ -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);
}
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit bd07784

Please sign in to comment.