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

@ -0,0 +1,211 @@
pragma solidity ^0.6.0;
library EnumerableMap {
// To implement this library for multiple types with as little code
// repetition as possible, we write it in terms of a generic Map type with
// bytes32 keys and values.
// The Map implementation uses private functions, and user-facing
// implementations (such as Uint256ToAddressMap) are just wrappers around
// the underlying Map.
// This means that we can only create new EnumerableMaps for types that fit
// in bytes32.
struct MapEntry {
bytes32 _key;
bytes32 _value;
}
struct Map {
// Storage of map keys and values
MapEntry[] _entries;
// Position of the entry defined by a key in the `entries` array, plus 1
// because index 0 means a key is not in the map.
mapping (bytes32 => uint256) _indexes;
}
/**
* @dev Adds a key-value pair to a map, or updates the value for an existing
* key. O(1).
*
* Returns true if the key was added to the map, that is if it was not
* already present.
*/
function _set(Map storage map, bytes32 key, bytes32 value) private returns (bool) {
// We read and store the key's index to prevent multiple reads from the same storage slot
uint256 keyIndex = map._indexes[key];
if (keyIndex == 0) { // Equivalent to !contains(map, key)
map._entries.push(MapEntry({ _key: key, _value: value }));
// The entry is stored at length-1, but we add 1 to all indexes
// and use 0 as a sentinel value
map._indexes[key] = map._entries.length;
return true;
} else {
map._entries[keyIndex - 1]._value = value;
return false;
}
}
/**
* @dev Removes a key-value pair from a map. O(1).
*
* Returns true if the key was removed from the map, that is if it was present.
*/
function _remove(Map storage map, bytes32 key) private returns (bool) {
// We read and store the key's index to prevent multiple reads from the same storage slot
uint256 keyIndex = map._indexes[key];
if (keyIndex != 0) { // Equivalent to contains(map, key)
// To delete a key-value pair from the _entries array in O(1), we swap the entry to delete with the last one
// in the array, and then remove the last entry (sometimes called as 'swap and pop').
// This modifies the order of the array, as noted in {at}.
uint256 toDeleteIndex = keyIndex - 1;
uint256 lastIndex = map._entries.length - 1;
// When the entry to delete is the last one, the swap operation is unnecessary. However, since this occurs
// so rarely, we still do the swap anyway to avoid the gas cost of adding an 'if' statement.
MapEntry storage lastEntry = map._entries[lastIndex];
// Move the last entry to the index where the entry to delete is
map._entries[toDeleteIndex] = lastEntry;
// Update the index for the moved entry
map._indexes[lastEntry._key] = toDeleteIndex + 1; // All indexes are 1-based
// Delete the slot where the moved entry was stored
map._entries.pop();
// Delete the index for the deleted slot
delete map._indexes[key];
return true;
} else {
return false;
}
}
/**
* @dev Returns true if the key is in the map. O(1).
*/
function _contains(Map storage map, bytes32 key) private view returns (bool) {
return map._indexes[key] != 0;
}
/**
* @dev Returns the number of key-value pairs in the map. O(1).
*/
function _length(Map storage map) private view returns (uint256) {
return map._entries.length;
}
/**
* @dev Returns the key-value pair stored at position `index` in the map. O(1).
*
* Note that there are no guarantees on the ordering of entries inside the
* array, and it may change when more entries are added or removed.
*
* Requirements:
*
* - `index` must be strictly less than {length}.
*/
function _at(Map storage map, uint256 index) private view returns (bytes32, bytes32) {
require(map._entries.length > index, "EnumerableMap: index out of bounds");
MapEntry storage entry = map._entries[index];
return (entry._key, entry._value);
}
/**
* @dev Returns the value associated with `key`. O(1).
*
* Requirements:
*
* - `key` must be in the map.
*/
function _get(Map storage map, bytes32 key) private view returns (bytes32) {
return _get(map, key, "EnumerableMap: nonexistent key");
}
/**
* @dev Same as {_get}, with a custom error message when `key` is not in the map.
*/
function _get(Map storage map, bytes32 key, string memory errorMessage) private view returns (bytes32) {
uint256 keyIndex = map._indexes[key];
require(keyIndex != 0, errorMessage); // Equivalent to contains(map, key)
return map._entries[keyIndex - 1]._value; // All indexes are 1-based
}
// UintToAddressMap
struct UintToAddressMap {
Map _inner;
}
/**
* @dev Adds a key-value pair to a map, or updates the value for an existing
* key. O(1).
*
* Returns true if the key was added to the map, that is if it was not
* already present.
*/
function set(UintToAddressMap storage map, uint256 key, address value) internal returns (bool) {
return _set(map._inner, bytes32(key), bytes32(uint256(value)));
}
/**
* @dev Removes a value from a set. O(1).
*
* Returns true if the key was removed from the map, that is if it was present.
*/
function remove(UintToAddressMap storage map, uint256 key) internal returns (bool) {
return _remove(map._inner, bytes32(key));
}
/**
* @dev Returns true if the key is in the map. O(1).
*/
function contains(UintToAddressMap storage map, uint256 key) internal view returns (bool) {
return _contains(map._inner, bytes32(key));
}
/**
* @dev Returns the number of elements in the map. O(1).
*/
function length(UintToAddressMap storage map) internal view returns (uint256) {
return _length(map._inner);
}
/**
* @dev Returns the element stored at position `index` in the set. O(1).
* Note that there are no guarantees on the ordering of values inside the
* array, and it may change when more values are added or removed.
*
* Requirements:
*
* - `index` must be strictly less than {length}.
*/
function at(UintToAddressMap storage map, uint256 index) internal view returns (uint256, address) {
(bytes32 key, bytes32 value) = _at(map._inner, index);
return (uint256(key), address(uint256(value)));
}
/**
* @dev Returns the value associated with `key`. O(1).
*
* Requirements:
*
* - `key` must be in the map.
*/
function get(UintToAddressMap storage map, uint256 key) internal view returns (address) {
return address(uint256(_get(map._inner, bytes32(key))));
}
/**
* @dev Same as {get}, with a custom error message when `key` is not in the map.
*/
function get(UintToAddressMap storage map, uint256 key, string memory errorMessage) internal view returns (address) {
return address(uint256(_get(map._inner, bytes32(key), errorMessage)));
}
}

View File

@ -18,24 +18,32 @@ pragma solidity ^0.6.0;
* @author Alberto Cuesta Cañada
*/
library EnumerableSet {
// To implement this library for multiple types with as little code
// repetition as possible, we write it in terms of a generic Set type with
// bytes32 values.
// The Set implementation uses private functions, and user-facing
// implementations (such as AddressSet) are just wrappers around the
// underlying Set.
// This means that we can only create new EnumerableSets for types that fit
// in bytes32.
struct Set {
// Storage of set values
bytes32[] _values;
struct AddressSet {
address[] _values;
// Position of the value in the `values` array, plus 1 because index 0
// means a value is not in the set.
mapping (address => uint256) _indexes;
mapping (bytes32 => uint256) _indexes;
}
/**
* @dev Add a value to a set. O(1).
*
* Returns false if the value was already in the set.
* Returns true if the value was added to the set, that is if it was not
* already present.
*/
function add(AddressSet storage set, address value)
internal
returns (bool)
{
if (!contains(set, value)) {
function _add(Set storage set, bytes32 value) private returns (bool) {
if (!_contains(set, value)) {
set._values.push(value);
// The value is stored at length-1, but we add 1 to all indexes
// and use 0 as a sentinel value
@ -49,25 +57,30 @@ library EnumerableSet {
/**
* @dev Removes a value from a set. O(1).
*
* Returns false if the value was not present in the set.
* Returns true if the value was removed from the set, that is if it was
* present.
*/
function remove(AddressSet storage set, address value)
internal
returns (bool)
{
if (contains(set, value)){
uint256 toDeleteIndex = set._indexes[value] - 1;
function _remove(Set storage set, bytes32 value) private returns (bool) {
// We read and store the value's index to prevent multiple reads from the same storage slot
uint256 valueIndex = set._indexes[value];
if (valueIndex != 0) { // Equivalent to contains(set, value)
// To delete an element from the _values array in O(1), we swap the element to delete with the last one in
// the array, and then remove the last element (sometimes called as 'swap and pop').
// This modifies the order of the array, as noted in {at}.
uint256 toDeleteIndex = valueIndex - 1;
uint256 lastIndex = set._values.length - 1;
// If the value we're deleting is the last one, we can just remove it without doing a swap
if (lastIndex != toDeleteIndex) {
address lastvalue = set._values[lastIndex];
// When the value to delete is the last one, the swap operation is unnecessary. However, since this occurs
// so rarely, we still do the swap anyway to avoid the gas cost of adding an 'if' statement.
// Move the last value to the index where the deleted value is
set._values[toDeleteIndex] = lastvalue;
// Update the index for the moved value
set._indexes[lastvalue] = toDeleteIndex + 1; // All indexes are 1-based
}
bytes32 lastvalue = set._values[lastIndex];
// Move the last value to the index where the value to delete is
set._values[toDeleteIndex] = lastvalue;
// Update the index for the moved value
set._indexes[lastvalue] = toDeleteIndex + 1; // All indexes are 1-based
// Delete the slot where the moved value was stored
set._values.pop();
@ -84,43 +97,14 @@ library EnumerableSet {
/**
* @dev Returns true if the value is in the set. O(1).
*/
function contains(AddressSet storage set, address value)
internal
view
returns (bool)
{
function _contains(Set storage set, bytes32 value) private view returns (bool) {
return set._indexes[value] != 0;
}
/**
* @dev Returns an array with all values in the set. O(N).
*
* Note that there are no guarantees on the ordering of values inside the
* array, and it may change when more values are added or removed.
* WARNING: This function may run out of gas on large sets: use {length} and
* {at} instead in these cases.
*/
function enumerate(AddressSet storage set)
internal
view
returns (address[] memory)
{
address[] memory output = new address[](set._values.length);
for (uint256 i; i < set._values.length; i++){
output[i] = set._values[i];
}
return output;
}
/**
* @dev Returns the number of values on the set. O(1).
*/
function length(AddressSet storage set)
internal
view
returns (uint256)
{
function _length(Set storage set) private view returns (uint256) {
return set._values.length;
}
@ -134,12 +118,117 @@ library EnumerableSet {
*
* - `index` must be strictly less than {length}.
*/
function at(AddressSet storage set, uint256 index)
internal
view
returns (address)
{
function _at(Set storage set, uint256 index) private view returns (bytes32) {
require(set._values.length > index, "EnumerableSet: index out of bounds");
return set._values[index];
}
// AddressSet
struct AddressSet {
Set _inner;
}
/**
* @dev Add a value to a set. O(1).
*
* Returns true if the value was added to the set, that is if it was not
* already present.
*/
function add(AddressSet storage set, address value) internal returns (bool) {
return _add(set._inner, bytes32(uint256(value)));
}
/**
* @dev Removes a value from a set. O(1).
*
* Returns true if the value was removed from the set, that is if it was
* present.
*/
function remove(AddressSet storage set, address value) internal returns (bool) {
return _remove(set._inner, bytes32(uint256(value)));
}
/**
* @dev Returns true if the value is in the set. O(1).
*/
function contains(AddressSet storage set, address value) internal view returns (bool) {
return _contains(set._inner, bytes32(uint256(value)));
}
/**
* @dev Returns the number of values in the set. O(1).
*/
function length(AddressSet storage set) internal view returns (uint256) {
return _length(set._inner);
}
/**
* @dev Returns the value stored at position `index` in the set. O(1).
*
* Note that there are no guarantees on the ordering of values inside the
* array, and it may change when more values are added or removed.
*
* Requirements:
*
* - `index` must be strictly less than {length}.
*/
function at(AddressSet storage set, uint256 index) internal view returns (address) {
return address(uint256(_at(set._inner, index)));
}
// UintSet
struct UintSet {
Set _inner;
}
/**
* @dev Add a value to a set. O(1).
*
* Returns true if the value was added to the set, that is if it was not
* already present.
*/
function add(UintSet storage set, uint256 value) internal returns (bool) {
return _add(set._inner, bytes32(value));
}
/**
* @dev Removes a value from a set. O(1).
*
* Returns true if the value was removed from the set, that is if it was
* present.
*/
function remove(UintSet storage set, uint256 value) internal returns (bool) {
return _remove(set._inner, bytes32(value));
}
/**
* @dev Returns true if the value is in the set. O(1).
*/
function contains(UintSet storage set, uint256 value) internal view returns (bool) {
return _contains(set._inner, bytes32(value));
}
/**
* @dev Returns the number of values on the set. O(1).
*/
function length(UintSet storage set) internal view returns (uint256) {
return _length(set._inner);
}
/**
* @dev Returns the value stored at position `index` in the set. O(1).
*
* Note that there are no guarantees on the ordering of values inside the
* array, and it may change when more values are added or removed.
*
* Requirements:
*
* - `index` must be strictly less than {length}.
*/
function at(UintSet storage set, uint256 index) internal view returns (uint256) {
return uint256(_at(set._inner, index));
}
}