EnumerableMap.sol 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224
  1. // SPDX-License-Identifier: MIT
  2. pragma solidity ^0.8.0;
  3. import "./EnumerableSet.sol";
  4. /**
  5. * @dev Library for managing an enumerable variant of Solidity's
  6. * https://solidity.readthedocs.io/en/latest/types.html#mapping-types[`mapping`]
  7. * type.
  8. *
  9. * Maps have the following properties:
  10. *
  11. * - Entries are added, removed, and checked for existence in constant time
  12. * (O(1)).
  13. * - Entries are enumerated in O(n). No guarantees are made on the ordering.
  14. *
  15. * ```
  16. * contract Example {
  17. * // Add the library methods
  18. * using EnumerableMap for EnumerableMap.UintToAddressMap;
  19. *
  20. * // Declare a set state variable
  21. * EnumerableMap.UintToAddressMap private myMap;
  22. * }
  23. * ```
  24. *
  25. * As of v3.0.0, only maps of type `uint256 -> address` (`UintToAddressMap`) are
  26. * supported.
  27. */
  28. library EnumerableMap {
  29. using EnumerableSet for EnumerableSet.Bytes32Set;
  30. // To implement this library for multiple types with as little code
  31. // repetition as possible, we write it in terms of a generic Map type with
  32. // bytes32 keys and values.
  33. // The Map implementation uses private functions, and user-facing
  34. // implementations (such as Uint256ToAddressMap) are just wrappers around
  35. // the underlying Map.
  36. // This means that we can only create new EnumerableMaps for types that fit
  37. // in bytes32.
  38. struct Map {
  39. // Storage of keys
  40. EnumerableSet.Bytes32Set _keys;
  41. mapping (bytes32 => bytes32) _values;
  42. }
  43. /**
  44. * @dev Adds a key-value pair to a map, or updates the value for an existing
  45. * key. O(1).
  46. *
  47. * Returns true if the key was added to the map, that is if it was not
  48. * already present.
  49. */
  50. function _set(Map storage map, bytes32 key, bytes32 value) private returns (bool) {
  51. map._values[key] = value;
  52. return map._keys.add(key);
  53. }
  54. /**
  55. * @dev Removes a key-value pair from a map. O(1).
  56. *
  57. * Returns true if the key was removed from the map, that is if it was present.
  58. */
  59. function _remove(Map storage map, bytes32 key) private returns (bool) {
  60. delete map._values[key];
  61. return map._keys.remove(key);
  62. }
  63. /**
  64. * @dev Returns true if the key is in the map. O(1).
  65. */
  66. function _contains(Map storage map, bytes32 key) private view returns (bool) {
  67. return map._keys.contains(key);
  68. }
  69. /**
  70. * @dev Returns the number of key-value pairs in the map. O(1).
  71. */
  72. function _length(Map storage map) private view returns (uint256) {
  73. return map._keys.length();
  74. }
  75. /**
  76. * @dev Returns the key-value pair stored at position `index` in the map. O(1).
  77. *
  78. * Note that there are no guarantees on the ordering of entries inside the
  79. * array, and it may change when more entries are added or removed.
  80. *
  81. * Requirements:
  82. *
  83. * - `index` must be strictly less than {length}.
  84. */
  85. function _at(Map storage map, uint256 index) private view returns (bytes32, bytes32) {
  86. bytes32 key = map._keys.at(index);
  87. return (key, map._values[key]);
  88. }
  89. /**
  90. * @dev Tries to returns the value associated with `key`. O(1).
  91. * Does not revert if `key` is not in the map.
  92. */
  93. function _tryGet(Map storage map, bytes32 key) private view returns (bool, bytes32) {
  94. bytes32 value = map._values[key];
  95. if (value == bytes32(0)) {
  96. return (_contains(map, key), bytes32(0));
  97. } else {
  98. return (true, value);
  99. }
  100. }
  101. /**
  102. * @dev Returns the value associated with `key`. O(1).
  103. *
  104. * Requirements:
  105. *
  106. * - `key` must be in the map.
  107. */
  108. function _get(Map storage map, bytes32 key) private view returns (bytes32) {
  109. bytes32 value = map._values[key];
  110. require(value != 0 || _contains(map, key), "EnumerableMap: nonexistent key");
  111. return value;
  112. }
  113. /**
  114. * @dev Same as {_get}, with a custom error message when `key` is not in the map.
  115. *
  116. * CAUTION: This function is deprecated because it requires allocating memory for the error
  117. * message unnecessarily. For custom revert reasons use {_tryGet}.
  118. */
  119. function _get(Map storage map, bytes32 key, string memory errorMessage) private view returns (bytes32) {
  120. bytes32 value = map._values[key];
  121. require(value != 0 || _contains(map, key), errorMessage);
  122. return value;
  123. }
  124. // UintToAddressMap
  125. struct UintToAddressMap {
  126. Map _inner;
  127. }
  128. /**
  129. * @dev Adds a key-value pair to a map, or updates the value for an existing
  130. * key. O(1).
  131. *
  132. * Returns true if the key was added to the map, that is if it was not
  133. * already present.
  134. */
  135. function set(UintToAddressMap storage map, uint256 key, address value) internal returns (bool) {
  136. return _set(map._inner, bytes32(key), bytes32(uint256(uint160(value))));
  137. }
  138. /**
  139. * @dev Removes a value from a set. O(1).
  140. *
  141. * Returns true if the key was removed from the map, that is if it was present.
  142. */
  143. function remove(UintToAddressMap storage map, uint256 key) internal returns (bool) {
  144. return _remove(map._inner, bytes32(key));
  145. }
  146. /**
  147. * @dev Returns true if the key is in the map. O(1).
  148. */
  149. function contains(UintToAddressMap storage map, uint256 key) internal view returns (bool) {
  150. return _contains(map._inner, bytes32(key));
  151. }
  152. /**
  153. * @dev Returns the number of elements in the map. O(1).
  154. */
  155. function length(UintToAddressMap storage map) internal view returns (uint256) {
  156. return _length(map._inner);
  157. }
  158. /**
  159. * @dev Returns the element stored at position `index` in the set. O(1).
  160. * Note that there are no guarantees on the ordering of values inside the
  161. * array, and it may change when more values are added or removed.
  162. *
  163. * Requirements:
  164. *
  165. * - `index` must be strictly less than {length}.
  166. */
  167. function at(UintToAddressMap storage map, uint256 index) internal view returns (uint256, address) {
  168. (bytes32 key, bytes32 value) = _at(map._inner, index);
  169. return (uint256(key), address(uint160(uint256(value))));
  170. }
  171. /**
  172. * @dev Tries to returns the value associated with `key`. O(1).
  173. * Does not revert if `key` is not in the map.
  174. *
  175. * _Available since v3.4._
  176. */
  177. function tryGet(UintToAddressMap storage map, uint256 key) internal view returns (bool, address) {
  178. (bool success, bytes32 value) = _tryGet(map._inner, bytes32(key));
  179. return (success, address(uint160(uint256(value))));
  180. }
  181. /**
  182. * @dev Returns the value associated with `key`. O(1).
  183. *
  184. * Requirements:
  185. *
  186. * - `key` must be in the map.
  187. */
  188. function get(UintToAddressMap storage map, uint256 key) internal view returns (address) {
  189. return address(uint160(uint256(_get(map._inner, bytes32(key)))));
  190. }
  191. /**
  192. * @dev Same as {get}, with a custom error message when `key` is not in the map.
  193. *
  194. * CAUTION: This function is deprecated because it requires allocating memory for the error
  195. * message unnecessarily. For custom revert reasons use {tryGet}.
  196. */
  197. function get(UintToAddressMap storage map, uint256 key, string memory errorMessage) internal view returns (address) {
  198. return address(uint160(uint256(_get(map._inner, bytes32(key), errorMessage))));
  199. }
  200. }