EnumerableSet.sol 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234
  1. pragma solidity ^0.6.0;
  2. /**
  3. * @dev Library for managing
  4. * https://en.wikipedia.org/wiki/Set_(abstract_data_type)[sets] of primitive
  5. * types.
  6. *
  7. * Sets have the following properties:
  8. *
  9. * - Elements are added, removed, and checked for existence in constant time
  10. * (O(1)).
  11. * - Elements are enumerated in O(n). No guarantees are made on the ordering.
  12. *
  13. * As of v2.5.0, only `address` sets are supported.
  14. *
  15. * Include with `using EnumerableSet for EnumerableSet.AddressSet;`.
  16. *
  17. * @author Alberto Cuesta Cañada
  18. */
  19. library EnumerableSet {
  20. // To implement this library for multiple types with as little code
  21. // repetition as possible, we write it in terms of a generic Set type with
  22. // bytes32 values.
  23. // The Set implementation uses private functions, and user-facing
  24. // implementations (such as AddressSet) are just wrappers around the
  25. // underlying Set.
  26. // This means that we can only create new EnumerableSets for types that fit
  27. // in bytes32.
  28. struct Set {
  29. // Storage of set values
  30. bytes32[] _values;
  31. // Position of the value in the `values` array, plus 1 because index 0
  32. // means a value is not in the set.
  33. mapping (bytes32 => uint256) _indexes;
  34. }
  35. /**
  36. * @dev Add a value to a set. O(1).
  37. *
  38. * Returns true if the value was added to the set, that is if it was not
  39. * already present.
  40. */
  41. function _add(Set storage set, bytes32 value) private returns (bool) {
  42. if (!_contains(set, value)) {
  43. set._values.push(value);
  44. // The value is stored at length-1, but we add 1 to all indexes
  45. // and use 0 as a sentinel value
  46. set._indexes[value] = set._values.length;
  47. return true;
  48. } else {
  49. return false;
  50. }
  51. }
  52. /**
  53. * @dev Removes a value from a set. O(1).
  54. *
  55. * Returns true if the value was removed from the set, that is if it was
  56. * present.
  57. */
  58. function _remove(Set storage set, bytes32 value) private returns (bool) {
  59. // We read and store the value's index to prevent multiple reads from the same storage slot
  60. uint256 valueIndex = set._indexes[value];
  61. if (valueIndex != 0) { // Equivalent to contains(set, value)
  62. // To delete an element from the _values array in O(1), we swap the element to delete with the last one in
  63. // the array, and then remove the last element (sometimes called as 'swap and pop').
  64. // This modifies the order of the array, as noted in {at}.
  65. uint256 toDeleteIndex = valueIndex - 1;
  66. uint256 lastIndex = set._values.length - 1;
  67. // When the value to delete is the last one, the swap operation is unnecessary. However, since this occurs
  68. // so rarely, we still do the swap anyway to avoid the gas cost of adding an 'if' statement.
  69. bytes32 lastvalue = set._values[lastIndex];
  70. // Move the last value to the index where the value to delete is
  71. set._values[toDeleteIndex] = lastvalue;
  72. // Update the index for the moved value
  73. set._indexes[lastvalue] = toDeleteIndex + 1; // All indexes are 1-based
  74. // Delete the slot where the moved value was stored
  75. set._values.pop();
  76. // Delete the index for the deleted slot
  77. delete set._indexes[value];
  78. return true;
  79. } else {
  80. return false;
  81. }
  82. }
  83. /**
  84. * @dev Returns true if the value is in the set. O(1).
  85. */
  86. function _contains(Set storage set, bytes32 value) private view returns (bool) {
  87. return set._indexes[value] != 0;
  88. }
  89. /**
  90. * @dev Returns the number of values on the set. O(1).
  91. */
  92. function _length(Set storage set) private view returns (uint256) {
  93. return set._values.length;
  94. }
  95. /**
  96. * @dev Returns the value stored at position `index` in the set. O(1).
  97. *
  98. * Note that there are no guarantees on the ordering of values inside the
  99. * array, and it may change when more values are added or removed.
  100. *
  101. * Requirements:
  102. *
  103. * - `index` must be strictly less than {length}.
  104. */
  105. function _at(Set storage set, uint256 index) private view returns (bytes32) {
  106. require(set._values.length > index, "EnumerableSet: index out of bounds");
  107. return set._values[index];
  108. }
  109. // AddressSet
  110. struct AddressSet {
  111. Set _inner;
  112. }
  113. /**
  114. * @dev Add a value to a set. O(1).
  115. *
  116. * Returns true if the value was added to the set, that is if it was not
  117. * already present.
  118. */
  119. function add(AddressSet storage set, address value) internal returns (bool) {
  120. return _add(set._inner, bytes32(uint256(value)));
  121. }
  122. /**
  123. * @dev Removes a value from a set. O(1).
  124. *
  125. * Returns true if the value was removed from the set, that is if it was
  126. * present.
  127. */
  128. function remove(AddressSet storage set, address value) internal returns (bool) {
  129. return _remove(set._inner, bytes32(uint256(value)));
  130. }
  131. /**
  132. * @dev Returns true if the value is in the set. O(1).
  133. */
  134. function contains(AddressSet storage set, address value) internal view returns (bool) {
  135. return _contains(set._inner, bytes32(uint256(value)));
  136. }
  137. /**
  138. * @dev Returns the number of values in the set. O(1).
  139. */
  140. function length(AddressSet storage set) internal view returns (uint256) {
  141. return _length(set._inner);
  142. }
  143. /**
  144. * @dev Returns the value stored at position `index` in the set. O(1).
  145. *
  146. * Note that there are no guarantees on the ordering of values inside the
  147. * array, and it may change when more values are added or removed.
  148. *
  149. * Requirements:
  150. *
  151. * - `index` must be strictly less than {length}.
  152. */
  153. function at(AddressSet storage set, uint256 index) internal view returns (address) {
  154. return address(uint256(_at(set._inner, index)));
  155. }
  156. // UintSet
  157. struct UintSet {
  158. Set _inner;
  159. }
  160. /**
  161. * @dev Add a value to a set. O(1).
  162. *
  163. * Returns true if the value was added to the set, that is if it was not
  164. * already present.
  165. */
  166. function add(UintSet storage set, uint256 value) internal returns (bool) {
  167. return _add(set._inner, bytes32(value));
  168. }
  169. /**
  170. * @dev Removes a value from a set. O(1).
  171. *
  172. * Returns true if the value was removed from the set, that is if it was
  173. * present.
  174. */
  175. function remove(UintSet storage set, uint256 value) internal returns (bool) {
  176. return _remove(set._inner, bytes32(value));
  177. }
  178. /**
  179. * @dev Returns true if the value is in the set. O(1).
  180. */
  181. function contains(UintSet storage set, uint256 value) internal view returns (bool) {
  182. return _contains(set._inner, bytes32(value));
  183. }
  184. /**
  185. * @dev Returns the number of values on the set. O(1).
  186. */
  187. function length(UintSet storage set) internal view returns (uint256) {
  188. return _length(set._inner);
  189. }
  190. /**
  191. * @dev Returns the value stored at position `index` in the set. O(1).
  192. *
  193. * Note that there are no guarantees on the ordering of values inside the
  194. * array, and it may change when more values are added or removed.
  195. *
  196. * Requirements:
  197. *
  198. * - `index` must be strictly less than {length}.
  199. */
  200. function at(UintSet storage set, uint256 index) internal view returns (uint256) {
  201. return uint256(_at(set._inner, index));
  202. }
  203. }