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:
@ -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