EnumerableSet.sol 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  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. struct AddressSet {
  21. // Position of the value in the `values` array, plus 1 because index 0
  22. // means a value is not in the set.
  23. mapping (address => uint256) index;
  24. address[] values;
  25. }
  26. /**
  27. * @dev Add a value to a set. O(1).
  28. * Returns false if the value was already in the set.
  29. */
  30. function add(AddressSet storage set, address value)
  31. internal
  32. returns (bool)
  33. {
  34. if (!contains(set, value)) {
  35. set.values.push(value);
  36. // The element is stored at length-1, but we add 1 to all indexes
  37. // and use 0 as a sentinel value
  38. set.index[value] = set.values.length;
  39. return true;
  40. } else {
  41. return false;
  42. }
  43. }
  44. /**
  45. * @dev Removes a value from a set. O(1).
  46. * Returns false if the value was not present in the set.
  47. */
  48. function remove(AddressSet storage set, address value)
  49. internal
  50. returns (bool)
  51. {
  52. if (contains(set, value)){
  53. uint256 toDeleteIndex = set.index[value] - 1;
  54. uint256 lastIndex = set.values.length - 1;
  55. // If the element we're deleting is the last one, we can just remove it without doing a swap
  56. if (lastIndex != toDeleteIndex) {
  57. address lastValue = set.values[lastIndex];
  58. // Move the last value to the index where the deleted value is
  59. set.values[toDeleteIndex] = lastValue;
  60. // Update the index for the moved value
  61. set.index[lastValue] = toDeleteIndex + 1; // All indexes are 1-based
  62. }
  63. // Delete the index entry for the deleted value
  64. delete set.index[value];
  65. // Delete the old entry for the moved value
  66. set.values.pop();
  67. return true;
  68. } else {
  69. return false;
  70. }
  71. }
  72. /**
  73. * @dev Returns true if the value is in the set. O(1).
  74. */
  75. function contains(AddressSet storage set, address value)
  76. internal
  77. view
  78. returns (bool)
  79. {
  80. return set.index[value] != 0;
  81. }
  82. /**
  83. * @dev Returns an array with all values in the set. O(N).
  84. * Note that there are no guarantees on the ordering of values inside the
  85. * array, and it may change when more values are added or removed.
  86. * WARNING: This function may run out of gas on large sets: use {length} and
  87. * {get} instead in these cases.
  88. */
  89. function enumerate(AddressSet storage set)
  90. internal
  91. view
  92. returns (address[] memory)
  93. {
  94. address[] memory output = new address[](set.values.length);
  95. for (uint256 i; i < set.values.length; i++){
  96. output[i] = set.values[i];
  97. }
  98. return output;
  99. }
  100. /**
  101. * @dev Returns the number of elements on the set. O(1).
  102. */
  103. function length(AddressSet storage set)
  104. internal
  105. view
  106. returns (uint256)
  107. {
  108. return set.values.length;
  109. }
  110. /** @dev Returns the element stored at position `index` in the set. O(1).
  111. * Note that there are no guarantees on the ordering of values inside the
  112. * array, and it may change when more values are added or removed.
  113. *
  114. * Requirements:
  115. *
  116. * - `index` must be strictly less than {length}.
  117. */
  118. function get(AddressSet storage set, uint256 index)
  119. internal
  120. view
  121. returns (address)
  122. {
  123. return set.values[index];
  124. }
  125. }