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:
211
contracts/utils/EnumerableMap.sol
Normal file
211
contracts/utils/EnumerableMap.sol
Normal 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)));
|
||||
}
|
||||
}
|
||||
@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user