EnumerableMap.js 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297
  1. const format = require('../format-lines');
  2. const { fromBytes32, toBytes32 } = require('./conversion');
  3. const { TYPES } = require('./EnumerableMap.opts');
  4. const header = `\
  5. pragma solidity ^0.8.20;
  6. import {EnumerableSet} from "./EnumerableSet.sol";
  7. /**
  8. * @dev Library for managing an enumerable variant of Solidity's
  9. * https://solidity.readthedocs.io/en/latest/types.html#mapping-types[\`mapping\`]
  10. * type.
  11. *
  12. * Maps have the following properties:
  13. *
  14. * - Entries are added, removed, and checked for existence in constant time
  15. * (O(1)).
  16. * - Entries are enumerated in O(n). No guarantees are made on the ordering.
  17. * - Map can be cleared (all entries removed) in O(n).
  18. *
  19. * \`\`\`solidity
  20. * contract Example {
  21. * // Add the library methods
  22. * using EnumerableMap for EnumerableMap.UintToAddressMap;
  23. *
  24. * // Declare a set state variable
  25. * EnumerableMap.UintToAddressMap private myMap;
  26. * }
  27. * \`\`\`
  28. *
  29. * The following map types are supported:
  30. *
  31. * - \`uint256 -> address\` (\`UintToAddressMap\`) since v3.0.0
  32. * - \`address -> uint256\` (\`AddressToUintMap\`) since v4.6.0
  33. * - \`bytes32 -> bytes32\` (\`Bytes32ToBytes32Map\`) since v4.6.0
  34. * - \`uint256 -> uint256\` (\`UintToUintMap\`) since v4.7.0
  35. * - \`bytes32 -> uint256\` (\`Bytes32ToUintMap\`) since v4.7.0
  36. * - \`uint256 -> bytes32\` (\`UintToBytes32Map\`) since v5.1.0
  37. * - \`address -> address\` (\`AddressToAddressMap\`) since v5.1.0
  38. * - \`address -> bytes32\` (\`AddressToBytes32Map\`) since v5.1.0
  39. * - \`bytes32 -> address\` (\`Bytes32ToAddressMap\`) since v5.1.0
  40. *
  41. * [WARNING]
  42. * ====
  43. * Trying to delete such a structure from storage will likely result in data corruption, rendering the structure
  44. * unusable.
  45. * See https://github.com/ethereum/solidity/pull/11843[ethereum/solidity#11843] for more info.
  46. *
  47. * In order to clean an EnumerableMap, you can either remove all elements one by one or create a fresh instance using an
  48. * array of EnumerableMap.
  49. * ====
  50. */
  51. `;
  52. const defaultMap = `\
  53. // To implement this library for multiple types with as little code repetition as possible, we write it in
  54. // terms of a generic Map type with bytes32 keys and values. The Map implementation uses private functions,
  55. // and user-facing implementations such as \`UintToAddressMap\` are just wrappers around the underlying Map.
  56. // This means that we can only create new EnumerableMaps for types that fit in bytes32.
  57. /**
  58. * @dev Query for a nonexistent map key.
  59. */
  60. error EnumerableMapNonexistentKey(bytes32 key);
  61. struct Bytes32ToBytes32Map {
  62. // Storage of keys
  63. EnumerableSet.Bytes32Set _keys;
  64. mapping(bytes32 key => bytes32) _values;
  65. }
  66. /**
  67. * @dev Adds a key-value pair to a map, or updates the value for an existing
  68. * key. O(1).
  69. *
  70. * Returns true if the key was added to the map, that is if it was not
  71. * already present.
  72. */
  73. function set(Bytes32ToBytes32Map storage map, bytes32 key, bytes32 value) internal returns (bool) {
  74. map._values[key] = value;
  75. return map._keys.add(key);
  76. }
  77. /**
  78. * @dev Removes a key-value pair from a map. O(1).
  79. *
  80. * Returns true if the key was removed from the map, that is if it was present.
  81. */
  82. function remove(Bytes32ToBytes32Map storage map, bytes32 key) internal returns (bool) {
  83. delete map._values[key];
  84. return map._keys.remove(key);
  85. }
  86. /**
  87. * @dev Removes all the entries from a map. O(n).
  88. *
  89. * WARNING: Developers should keep in mind that this function has an unbounded cost and using it may render the
  90. * function uncallable if the map grows to the point where clearing it consumes too much gas to fit in a block.
  91. */
  92. function clear(Bytes32ToBytes32Map storage map) internal {
  93. uint256 len = length(map);
  94. for (uint256 i = 0; i < len; ++i) {
  95. delete map._values[map._keys.at(i)];
  96. }
  97. map._keys.clear();
  98. }
  99. /**
  100. * @dev Returns true if the key is in the map. O(1).
  101. */
  102. function contains(Bytes32ToBytes32Map storage map, bytes32 key) internal view returns (bool) {
  103. return map._keys.contains(key);
  104. }
  105. /**
  106. * @dev Returns the number of key-value pairs in the map. O(1).
  107. */
  108. function length(Bytes32ToBytes32Map storage map) internal view returns (uint256) {
  109. return map._keys.length();
  110. }
  111. /**
  112. * @dev Returns the key-value pair stored at position \`index\` in the map. O(1).
  113. *
  114. * Note that there are no guarantees on the ordering of entries inside the
  115. * array, and it may change when more entries are added or removed.
  116. *
  117. * Requirements:
  118. *
  119. * - \`index\` must be strictly less than {length}.
  120. */
  121. function at(Bytes32ToBytes32Map storage map, uint256 index) internal view returns (bytes32 key, bytes32 value) {
  122. bytes32 atKey = map._keys.at(index);
  123. return (atKey, map._values[atKey]);
  124. }
  125. /**
  126. * @dev Tries to returns the value associated with \`key\`. O(1).
  127. * Does not revert if \`key\` is not in the map.
  128. */
  129. function tryGet(Bytes32ToBytes32Map storage map, bytes32 key) internal view returns (bool exists, bytes32 value) {
  130. bytes32 val = map._values[key];
  131. if (val == bytes32(0)) {
  132. return (contains(map, key), bytes32(0));
  133. } else {
  134. return (true, val);
  135. }
  136. }
  137. /**
  138. * @dev Returns the value associated with \`key\`. O(1).
  139. *
  140. * Requirements:
  141. *
  142. * - \`key\` must be in the map.
  143. */
  144. function get(Bytes32ToBytes32Map storage map, bytes32 key) internal view returns (bytes32) {
  145. bytes32 value = map._values[key];
  146. if (value == 0 && !contains(map, key)) {
  147. revert EnumerableMapNonexistentKey(key);
  148. }
  149. return value;
  150. }
  151. /**
  152. * @dev Return the an array containing all the keys
  153. *
  154. * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed
  155. * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that
  156. * this function has an unbounded cost, and using it as part of a state-changing function may render the function
  157. * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block.
  158. */
  159. function keys(Bytes32ToBytes32Map storage map) internal view returns (bytes32[] memory) {
  160. return map._keys.values();
  161. }
  162. `;
  163. const customMap = ({ name, keyType, valueType }) => `\
  164. // ${name}
  165. struct ${name} {
  166. Bytes32ToBytes32Map _inner;
  167. }
  168. /**
  169. * @dev Adds a key-value pair to a map, or updates the value for an existing
  170. * key. O(1).
  171. *
  172. * Returns true if the key was added to the map, that is if it was not
  173. * already present.
  174. */
  175. function set(${name} storage map, ${keyType} key, ${valueType} value) internal returns (bool) {
  176. return set(map._inner, ${toBytes32(keyType, 'key')}, ${toBytes32(valueType, 'value')});
  177. }
  178. /**
  179. * @dev Removes a value from a map. O(1).
  180. *
  181. * Returns true if the key was removed from the map, that is if it was present.
  182. */
  183. function remove(${name} storage map, ${keyType} key) internal returns (bool) {
  184. return remove(map._inner, ${toBytes32(keyType, 'key')});
  185. }
  186. /**
  187. * @dev Removes all the entries from a map. O(n).
  188. *
  189. * WARNING: Developers should keep in mind that this function has an unbounded cost and using it may render the
  190. * function uncallable if the map grows to the point where clearing it consumes too much gas to fit in a block.
  191. */
  192. function clear(${name} storage map) internal {
  193. clear(map._inner);
  194. }
  195. /**
  196. * @dev Returns true if the key is in the map. O(1).
  197. */
  198. function contains(${name} storage map, ${keyType} key) internal view returns (bool) {
  199. return contains(map._inner, ${toBytes32(keyType, 'key')});
  200. }
  201. /**
  202. * @dev Returns the number of elements in the map. O(1).
  203. */
  204. function length(${name} storage map) internal view returns (uint256) {
  205. return length(map._inner);
  206. }
  207. /**
  208. * @dev Returns the element stored at position \`index\` in the map. O(1).
  209. * Note that there are no guarantees on the ordering of values inside the
  210. * array, and it may change when more values are added or removed.
  211. *
  212. * Requirements:
  213. *
  214. * - \`index\` must be strictly less than {length}.
  215. */
  216. function at(${name} storage map, uint256 index) internal view returns (${keyType} key, ${valueType} value) {
  217. (bytes32 atKey, bytes32 val) = at(map._inner, index);
  218. return (${fromBytes32(keyType, 'atKey')}, ${fromBytes32(valueType, 'val')});
  219. }
  220. /**
  221. * @dev Tries to returns the value associated with \`key\`. O(1).
  222. * Does not revert if \`key\` is not in the map.
  223. */
  224. function tryGet(${name} storage map, ${keyType} key) internal view returns (bool exists, ${valueType} value) {
  225. (bool success, bytes32 val) = tryGet(map._inner, ${toBytes32(keyType, 'key')});
  226. return (success, ${fromBytes32(valueType, 'val')});
  227. }
  228. /**
  229. * @dev Returns the value associated with \`key\`. O(1).
  230. *
  231. * Requirements:
  232. *
  233. * - \`key\` must be in the map.
  234. */
  235. function get(${name} storage map, ${keyType} key) internal view returns (${valueType}) {
  236. return ${fromBytes32(valueType, `get(map._inner, ${toBytes32(keyType, 'key')})`)};
  237. }
  238. /**
  239. * @dev Return the an array containing all the keys
  240. *
  241. * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed
  242. * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that
  243. * this function has an unbounded cost, and using it as part of a state-changing function may render the function
  244. * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block.
  245. */
  246. function keys(${name} storage map) internal view returns (${keyType}[] memory) {
  247. bytes32[] memory store = keys(map._inner);
  248. ${keyType}[] memory result;
  249. assembly ("memory-safe") {
  250. result := store
  251. }
  252. return result;
  253. }
  254. `;
  255. // GENERATE
  256. module.exports = format(
  257. header.trimEnd(),
  258. 'library EnumerableMap {',
  259. format(
  260. [].concat(
  261. 'using EnumerableSet for EnumerableSet.Bytes32Set;',
  262. '',
  263. defaultMap,
  264. TYPES.map(details => customMap(details)),
  265. ),
  266. ).trimEnd(),
  267. '}',
  268. );