EnumerableMap.js 9.7 KB

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