EnumerableMap.sol 7.5 KB

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