From 090b17347dbadc884b9d73372150b2fa78f3fe50 Mon Sep 17 00:00:00 2001 From: Vectorized Date: Sun, 8 Dec 2024 06:16:56 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20EnumerableMapLib=20(#1215)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- prep/all.js | 1 + prep/gen-enumerable-map-lib.js | 88 +++++ src/utils/EnumerableMapLib.sol | 622 ++++++++++++++++++++++++++++++ src/utils/EnumerableSetLib.sol | 10 +- src/utils/g/EnumerableMapLib.sol | 634 +++++++++++++++++++++++++++++++ src/utils/g/EnumerableSetLib.sol | 10 +- test/EnumerableMapLib.t.sol | 81 ++++ 7 files changed, 1436 insertions(+), 10 deletions(-) create mode 100644 prep/gen-enumerable-map-lib.js create mode 100644 src/utils/EnumerableMapLib.sol create mode 100644 src/utils/g/EnumerableMapLib.sol create mode 100644 test/EnumerableMapLib.t.sol diff --git a/prep/all.js b/prep/all.js index 1e533d9c2..52ef3f925 100644 --- a/prep/all.js +++ b/prep/all.js @@ -5,6 +5,7 @@ const { runCommandSync } = require('./common.js'); async function main() { const scripts = [ 'gen-efficient-hash-lib.js', + 'gen-enumerable-map-lib.js', 'gen-safe-cast-lib.js', 'gen-globalized-libs.js', 'remove-trailing-whitespace.js' diff --git a/prep/gen-enumerable-map-lib.js b/prep/gen-enumerable-map-lib.js new file mode 100644 index 000000000..4ebdc16c9 --- /dev/null +++ b/prep/gen-enumerable-map-lib.js @@ -0,0 +1,88 @@ +#!/usr/bin/env node +const { genSectionRegex, readSync, writeAndFmtSync, normalizeNewlines, hexNoPrefix } = require('./common.js'); + +async function main() { + const srcPath = 'src/utils/EnumerableMapLib.sol'; + const maxDepth = 14; + let src = readSync(srcPath); + + const capitalize = s => s.charAt(0).toUpperCase() + s.slice(1); + + const mapType = (f, t) => capitalize(f) + 'To' + capitalize(t) + 'Map'; + + const crossForEach = (a, fn) => a.forEach(x => a.forEach(y => fn(x, y))); + + const genStructDef = (f, t) => { + return '/// @dev A enumerable map of `' + f + '` to `' + t + '`.\n' + + 'struct ' + mapType(f, t) + '{\n' + + 'EnumerableSetLib.' + capitalize(f) + 'Set _keys;\n' + + 'mapping(' + f + ' => ' + t + ') _values;\n}\n\n'; + }; + + const genGettersAndSettersDef = (f, t) => { + const mt = mapType(f, t); + let s = '/// @dev Adds a key-value pair to the map, or updates the value for an existing key.\n'; + s += '/// Returns true if `key` was added to the map, that is if it was not already present.\n'; + s += 'function set(' + mt + ' storage map, ' + f + ' key, ' + t + ' value) internal returns (bool) {\n'; + s += 'map._values[key] = value;\nreturn EnumerableSetLib.add(map._keys, key);\n}\n\n'; + + s += '/// @dev Removes a key-value pair from the map.\n'; + s += '/// Returns true if `key` was removed from the map, that is if it was present.\n'; + s += 'function remove(' + mt + ' storage map, ' + f + ' key) internal returns (bool) {\n'; + s += 'delete map._values[key];\nreturn EnumerableSetLib.remove(map._keys, key);\n}\n\n'; + + s += '/// @dev Returns true if the key is in the map.\n'; + s += 'function contains(' + mt + ' storage map, ' + f + ' key) internal view returns (bool) {\n'; + s += 'return EnumerableSetLib.contains(map._keys, key);\n}\n\n'; + + s += '/// @dev Returns the number of key-value pairs in the map.\n'; + s += 'function length(' + mt + ' storage map) internal view returns (uint256) {\n'; + s += 'return EnumerableSetLib.length(map._keys);\n}\n\n'; + + s += '/// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds.\n'; + s += 'function at(' + mt + ' storage map, uint256 i)' + s += 'internal view returns (' + f + ' key, ' + t + ' value) {\n'; + s += 'value = map._values[key = EnumerableSetLib.at(map._keys, i)];\n}\n\n'; + + s += '/// @dev Tries to return the value associated with the key.\n'; + s += 'function tryGet(' + mt + ' storage map, ' + f + ' key)' + s += 'internal view returns (bool exists, ' + t + ' value) {\n'; + s += 'exists = (value = map._values[key]) != ' + t + '(0) || contains(map, key);\n}\n\n'; + + s += '/// @dev Returns the value for the key. Reverts if the key is not found.\n'; + s += 'function get(' + mt + ' storage map, ' + f + ' key)' + s += 'internal view returns (' + t + ' value)\n{\n'; + s += 'if ((value = map._values[key]) == ' + t + '(0)) if (!contains(map, key)) _revertNotFound();\n}\n\n'; + + s += '/// @dev Returns the keys. May run out-of-gas if the map is too big.\n'; + s += 'function keys(' + mt + ' storage map) internal view returns (' + f + '[] memory) {\n'; + s += 'return EnumerableSetLib.values(map._keys);\n}\n\n'; + return s; + } + + const types = ['bytes32', 'uint256', 'address']; + + src = src.replace( + genSectionRegex('STRUCTS'), + (m0, m1, m2) => { + let chunks = [m1]; + crossForEach(types, (f, t) => chunks.push(genStructDef(f, t))); + chunks.push(m2); + return normalizeNewlines(chunks.join('\n\n\n')); + } + ).replace( + genSectionRegex('GETTERS / SETTERS'), + (m0, m1, m2) => { + let chunks = [m1]; + crossForEach(types, (f, t) => chunks.push(genGettersAndSettersDef(f, t))); + chunks.push(m2); + return normalizeNewlines(chunks.join('\n\n\n')); + } + ); + writeAndFmtSync(srcPath, src); +}; + +main().catch(e => { + console.error(e); + process.exit(1); +}); diff --git a/src/utils/EnumerableMapLib.sol b/src/utils/EnumerableMapLib.sol new file mode 100644 index 000000000..9296e42e8 --- /dev/null +++ b/src/utils/EnumerableMapLib.sol @@ -0,0 +1,622 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +import {EnumerableSetLib} from "./EnumerableSetLib.sol"; + +/// @notice Library for managing enumerable maps in storage. +/// @author Solady (https://github.com/vectorized/solady/blob/main/src/utils/EnumerableMapLib.sol) +/// @author Modified from OpenZeppelin (https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/utils/structs/EnumerableMap.sol) +library EnumerableMapLib { + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* CUSTOM ERRORS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev The key does not exist in the enumerable map. + error EnumerableMapKeyNotFound(); + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* STRUCTS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev A enumerable map of `bytes32` to `bytes32`. + struct Bytes32ToBytes32Map { + EnumerableSetLib.Bytes32Set _keys; + mapping(bytes32 => bytes32) _values; + } + + /// @dev A enumerable map of `bytes32` to `uint256`. + struct Bytes32ToUint256Map { + EnumerableSetLib.Bytes32Set _keys; + mapping(bytes32 => uint256) _values; + } + + /// @dev A enumerable map of `bytes32` to `address`. + struct Bytes32ToAddressMap { + EnumerableSetLib.Bytes32Set _keys; + mapping(bytes32 => address) _values; + } + + /// @dev A enumerable map of `uint256` to `bytes32`. + struct Uint256ToBytes32Map { + EnumerableSetLib.Uint256Set _keys; + mapping(uint256 => bytes32) _values; + } + + /// @dev A enumerable map of `uint256` to `uint256`. + struct Uint256ToUint256Map { + EnumerableSetLib.Uint256Set _keys; + mapping(uint256 => uint256) _values; + } + + /// @dev A enumerable map of `uint256` to `address`. + struct Uint256ToAddressMap { + EnumerableSetLib.Uint256Set _keys; + mapping(uint256 => address) _values; + } + + /// @dev A enumerable map of `address` to `bytes32`. + struct AddressToBytes32Map { + EnumerableSetLib.AddressSet _keys; + mapping(address => bytes32) _values; + } + + /// @dev A enumerable map of `address` to `uint256`. + struct AddressToUint256Map { + EnumerableSetLib.AddressSet _keys; + mapping(address => uint256) _values; + } + + /// @dev A enumerable map of `address` to `address`. + struct AddressToAddressMap { + EnumerableSetLib.AddressSet _keys; + mapping(address => address) _values; + } + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* GETTERS / SETTERS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Bytes32ToBytes32Map storage map, bytes32 key, bytes32 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Bytes32ToBytes32Map storage map, bytes32 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Bytes32ToBytes32Map storage map, bytes32 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Bytes32ToBytes32Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Bytes32ToBytes32Map storage map, uint256 i) + internal + view + returns (bytes32 key, bytes32 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Bytes32ToBytes32Map storage map, bytes32 key) + internal + view + returns (bool exists, bytes32 value) + { + exists = (value = map._values[key]) != bytes32(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Bytes32ToBytes32Map storage map, bytes32 key) + internal + view + returns (bytes32 value) + { + if ((value = map._values[key]) == bytes32(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Bytes32ToBytes32Map storage map) internal view returns (bytes32[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Bytes32ToUint256Map storage map, bytes32 key, uint256 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Bytes32ToUint256Map storage map, bytes32 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Bytes32ToUint256Map storage map, bytes32 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Bytes32ToUint256Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Bytes32ToUint256Map storage map, uint256 i) + internal + view + returns (bytes32 key, uint256 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Bytes32ToUint256Map storage map, bytes32 key) + internal + view + returns (bool exists, uint256 value) + { + exists = (value = map._values[key]) != uint256(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Bytes32ToUint256Map storage map, bytes32 key) + internal + view + returns (uint256 value) + { + if ((value = map._values[key]) == uint256(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Bytes32ToUint256Map storage map) internal view returns (bytes32[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Bytes32ToAddressMap storage map, bytes32 key, address value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Bytes32ToAddressMap storage map, bytes32 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Bytes32ToAddressMap storage map, bytes32 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Bytes32ToAddressMap storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Bytes32ToAddressMap storage map, uint256 i) + internal + view + returns (bytes32 key, address value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Bytes32ToAddressMap storage map, bytes32 key) + internal + view + returns (bool exists, address value) + { + exists = (value = map._values[key]) != address(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Bytes32ToAddressMap storage map, bytes32 key) + internal + view + returns (address value) + { + if ((value = map._values[key]) == address(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Bytes32ToAddressMap storage map) internal view returns (bytes32[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Uint256ToBytes32Map storage map, uint256 key, bytes32 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Uint256ToBytes32Map storage map, uint256 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Uint256ToBytes32Map storage map, uint256 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Uint256ToBytes32Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Uint256ToBytes32Map storage map, uint256 i) + internal + view + returns (uint256 key, bytes32 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Uint256ToBytes32Map storage map, uint256 key) + internal + view + returns (bool exists, bytes32 value) + { + exists = (value = map._values[key]) != bytes32(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Uint256ToBytes32Map storage map, uint256 key) + internal + view + returns (bytes32 value) + { + if ((value = map._values[key]) == bytes32(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Uint256ToBytes32Map storage map) internal view returns (uint256[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Uint256ToUint256Map storage map, uint256 key, uint256 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Uint256ToUint256Map storage map, uint256 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Uint256ToUint256Map storage map, uint256 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Uint256ToUint256Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Uint256ToUint256Map storage map, uint256 i) + internal + view + returns (uint256 key, uint256 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Uint256ToUint256Map storage map, uint256 key) + internal + view + returns (bool exists, uint256 value) + { + exists = (value = map._values[key]) != uint256(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Uint256ToUint256Map storage map, uint256 key) + internal + view + returns (uint256 value) + { + if ((value = map._values[key]) == uint256(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Uint256ToUint256Map storage map) internal view returns (uint256[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Uint256ToAddressMap storage map, uint256 key, address value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Uint256ToAddressMap storage map, uint256 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Uint256ToAddressMap storage map, uint256 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Uint256ToAddressMap storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Uint256ToAddressMap storage map, uint256 i) + internal + view + returns (uint256 key, address value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Uint256ToAddressMap storage map, uint256 key) + internal + view + returns (bool exists, address value) + { + exists = (value = map._values[key]) != address(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Uint256ToAddressMap storage map, uint256 key) + internal + view + returns (address value) + { + if ((value = map._values[key]) == address(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Uint256ToAddressMap storage map) internal view returns (uint256[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(AddressToBytes32Map storage map, address key, bytes32 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(AddressToBytes32Map storage map, address key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(AddressToBytes32Map storage map, address key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(AddressToBytes32Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(AddressToBytes32Map storage map, uint256 i) + internal + view + returns (address key, bytes32 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(AddressToBytes32Map storage map, address key) + internal + view + returns (bool exists, bytes32 value) + { + exists = (value = map._values[key]) != bytes32(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(AddressToBytes32Map storage map, address key) + internal + view + returns (bytes32 value) + { + if ((value = map._values[key]) == bytes32(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(AddressToBytes32Map storage map) internal view returns (address[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(AddressToUint256Map storage map, address key, uint256 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(AddressToUint256Map storage map, address key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(AddressToUint256Map storage map, address key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(AddressToUint256Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(AddressToUint256Map storage map, uint256 i) + internal + view + returns (address key, uint256 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(AddressToUint256Map storage map, address key) + internal + view + returns (bool exists, uint256 value) + { + exists = (value = map._values[key]) != uint256(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(AddressToUint256Map storage map, address key) + internal + view + returns (uint256 value) + { + if ((value = map._values[key]) == uint256(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(AddressToUint256Map storage map) internal view returns (address[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(AddressToAddressMap storage map, address key, address value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(AddressToAddressMap storage map, address key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(AddressToAddressMap storage map, address key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(AddressToAddressMap storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(AddressToAddressMap storage map, uint256 i) + internal + view + returns (address key, address value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(AddressToAddressMap storage map, address key) + internal + view + returns (bool exists, address value) + { + exists = (value = map._values[key]) != address(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(AddressToAddressMap storage map, address key) + internal + view + returns (address value) + { + if ((value = map._values[key]) == address(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(AddressToAddressMap storage map) internal view returns (address[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* PRIVATE HELPERS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev Reverts with `EnumerableMapKeyNotFound()`. + function _revertNotFound() private pure { + /// @solidity memory-safe-assembly + assembly { + mstore(0x00, 0x88682bf3) // `EnumerableMapKeyNotFound()`. + revert(0x1c, 0x04) + } + } +} diff --git a/src/utils/EnumerableSetLib.sol b/src/utils/EnumerableSetLib.sol index c0216b682..cdebf5507 100644 --- a/src/utils/EnumerableSetLib.sol +++ b/src/utils/EnumerableSetLib.sol @@ -612,7 +612,7 @@ library EnumerableSetLib { } } - /// @dev Returns the element at index `i` in the set. + /// @dev Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds. function at(AddressSet storage set, uint256 i) internal view returns (address result) { bytes32 rootSlot = _rootSlot(set); /// @solidity memory-safe-assembly @@ -623,7 +623,7 @@ library EnumerableSetLib { if (i >= length(set)) revert IndexOutOfBounds(); } - /// @dev Returns the element at index `i` in the set. + /// @dev Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds. function at(Bytes32Set storage set, uint256 i) internal view returns (bytes32 result) { result = _rootSlot(set); /// @solidity memory-safe-assembly @@ -634,17 +634,17 @@ library EnumerableSetLib { if (i >= length(set)) revert IndexOutOfBounds(); } - /// @dev Returns the element at index `i` in the set. + /// @dev Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds. function at(Uint256Set storage set, uint256 i) internal view returns (uint256 result) { result = uint256(at(_toBytes32Set(set), i)); } - /// @dev Returns the element at index `i` in the set. + /// @dev Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds. function at(Int256Set storage set, uint256 i) internal view returns (int256 result) { result = int256(uint256(at(_toBytes32Set(set), i))); } - /// @dev Returns the element at index `i` in the set. + /// @dev Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds. function at(Uint8Set storage set, uint256 i) internal view returns (uint8 result) { /// @solidity memory-safe-assembly assembly { diff --git a/src/utils/g/EnumerableMapLib.sol b/src/utils/g/EnumerableMapLib.sol new file mode 100644 index 000000000..91e3e9c7b --- /dev/null +++ b/src/utils/g/EnumerableMapLib.sol @@ -0,0 +1,634 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.13; + +// This file is auto-generated. + +/*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ +/* STRUCTS */ +/*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + +/// @dev A enumerable map of `bytes32` to `bytes32`. +struct Bytes32ToBytes32Map { + EnumerableSetLib.Bytes32Set _keys; + mapping(bytes32 => bytes32) _values; +} + +/// @dev A enumerable map of `bytes32` to `uint256`. +struct Bytes32ToUint256Map { + EnumerableSetLib.Bytes32Set _keys; + mapping(bytes32 => uint256) _values; +} + +/// @dev A enumerable map of `bytes32` to `address`. +struct Bytes32ToAddressMap { + EnumerableSetLib.Bytes32Set _keys; + mapping(bytes32 => address) _values; +} + +/// @dev A enumerable map of `uint256` to `bytes32`. +struct Uint256ToBytes32Map { + EnumerableSetLib.Uint256Set _keys; + mapping(uint256 => bytes32) _values; +} + +/// @dev A enumerable map of `uint256` to `uint256`. +struct Uint256ToUint256Map { + EnumerableSetLib.Uint256Set _keys; + mapping(uint256 => uint256) _values; +} + +/// @dev A enumerable map of `uint256` to `address`. +struct Uint256ToAddressMap { + EnumerableSetLib.Uint256Set _keys; + mapping(uint256 => address) _values; +} + +/// @dev A enumerable map of `address` to `bytes32`. +struct AddressToBytes32Map { + EnumerableSetLib.AddressSet _keys; + mapping(address => bytes32) _values; +} + +/// @dev A enumerable map of `address` to `uint256`. +struct AddressToUint256Map { + EnumerableSetLib.AddressSet _keys; + mapping(address => uint256) _values; +} + +/// @dev A enumerable map of `address` to `address`. +struct AddressToAddressMap { + EnumerableSetLib.AddressSet _keys; + mapping(address => address) _values; +} + +using EnumerableMapLib for Bytes32ToBytes32Map global; +using EnumerableMapLib for Bytes32ToUint256Map global; +using EnumerableMapLib for Bytes32ToAddressMap global; +using EnumerableMapLib for Uint256ToBytes32Map global; +using EnumerableMapLib for Uint256ToUint256Map global; +using EnumerableMapLib for Uint256ToAddressMap global; +using EnumerableMapLib for AddressToBytes32Map global; +using EnumerableMapLib for AddressToUint256Map global; +using EnumerableMapLib for AddressToAddressMap global; + +import {EnumerableSetLib} from "../EnumerableSetLib.sol"; + +/// @notice Library for managing enumerable maps in storage. +/// @author Solady (https://github.com/vectorized/solady/blob/main/src/utils/g/EnumerableMapLib.sol) +/// @author Modified from OpenZeppelin (https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/utils/structs/EnumerableMap.sol) +library EnumerableMapLib { + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* CUSTOM ERRORS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev The key does not exist in the enumerable map. + error EnumerableMapKeyNotFound(); + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* GETTERS / SETTERS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Bytes32ToBytes32Map storage map, bytes32 key, bytes32 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Bytes32ToBytes32Map storage map, bytes32 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Bytes32ToBytes32Map storage map, bytes32 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Bytes32ToBytes32Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Bytes32ToBytes32Map storage map, uint256 i) + internal + view + returns (bytes32 key, bytes32 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Bytes32ToBytes32Map storage map, bytes32 key) + internal + view + returns (bool exists, bytes32 value) + { + exists = (value = map._values[key]) != bytes32(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Bytes32ToBytes32Map storage map, bytes32 key) + internal + view + returns (bytes32 value) + { + if ((value = map._values[key]) == bytes32(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Bytes32ToBytes32Map storage map) internal view returns (bytes32[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Bytes32ToUint256Map storage map, bytes32 key, uint256 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Bytes32ToUint256Map storage map, bytes32 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Bytes32ToUint256Map storage map, bytes32 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Bytes32ToUint256Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Bytes32ToUint256Map storage map, uint256 i) + internal + view + returns (bytes32 key, uint256 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Bytes32ToUint256Map storage map, bytes32 key) + internal + view + returns (bool exists, uint256 value) + { + exists = (value = map._values[key]) != uint256(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Bytes32ToUint256Map storage map, bytes32 key) + internal + view + returns (uint256 value) + { + if ((value = map._values[key]) == uint256(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Bytes32ToUint256Map storage map) internal view returns (bytes32[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Bytes32ToAddressMap storage map, bytes32 key, address value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Bytes32ToAddressMap storage map, bytes32 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Bytes32ToAddressMap storage map, bytes32 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Bytes32ToAddressMap storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Bytes32ToAddressMap storage map, uint256 i) + internal + view + returns (bytes32 key, address value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Bytes32ToAddressMap storage map, bytes32 key) + internal + view + returns (bool exists, address value) + { + exists = (value = map._values[key]) != address(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Bytes32ToAddressMap storage map, bytes32 key) + internal + view + returns (address value) + { + if ((value = map._values[key]) == address(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Bytes32ToAddressMap storage map) internal view returns (bytes32[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Uint256ToBytes32Map storage map, uint256 key, bytes32 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Uint256ToBytes32Map storage map, uint256 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Uint256ToBytes32Map storage map, uint256 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Uint256ToBytes32Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Uint256ToBytes32Map storage map, uint256 i) + internal + view + returns (uint256 key, bytes32 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Uint256ToBytes32Map storage map, uint256 key) + internal + view + returns (bool exists, bytes32 value) + { + exists = (value = map._values[key]) != bytes32(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Uint256ToBytes32Map storage map, uint256 key) + internal + view + returns (bytes32 value) + { + if ((value = map._values[key]) == bytes32(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Uint256ToBytes32Map storage map) internal view returns (uint256[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Uint256ToUint256Map storage map, uint256 key, uint256 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Uint256ToUint256Map storage map, uint256 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Uint256ToUint256Map storage map, uint256 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Uint256ToUint256Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Uint256ToUint256Map storage map, uint256 i) + internal + view + returns (uint256 key, uint256 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Uint256ToUint256Map storage map, uint256 key) + internal + view + returns (bool exists, uint256 value) + { + exists = (value = map._values[key]) != uint256(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Uint256ToUint256Map storage map, uint256 key) + internal + view + returns (uint256 value) + { + if ((value = map._values[key]) == uint256(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Uint256ToUint256Map storage map) internal view returns (uint256[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(Uint256ToAddressMap storage map, uint256 key, address value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(Uint256ToAddressMap storage map, uint256 key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(Uint256ToAddressMap storage map, uint256 key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(Uint256ToAddressMap storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(Uint256ToAddressMap storage map, uint256 i) + internal + view + returns (uint256 key, address value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(Uint256ToAddressMap storage map, uint256 key) + internal + view + returns (bool exists, address value) + { + exists = (value = map._values[key]) != address(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(Uint256ToAddressMap storage map, uint256 key) + internal + view + returns (address value) + { + if ((value = map._values[key]) == address(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(Uint256ToAddressMap storage map) internal view returns (uint256[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(AddressToBytes32Map storage map, address key, bytes32 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(AddressToBytes32Map storage map, address key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(AddressToBytes32Map storage map, address key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(AddressToBytes32Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(AddressToBytes32Map storage map, uint256 i) + internal + view + returns (address key, bytes32 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(AddressToBytes32Map storage map, address key) + internal + view + returns (bool exists, bytes32 value) + { + exists = (value = map._values[key]) != bytes32(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(AddressToBytes32Map storage map, address key) + internal + view + returns (bytes32 value) + { + if ((value = map._values[key]) == bytes32(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(AddressToBytes32Map storage map) internal view returns (address[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(AddressToUint256Map storage map, address key, uint256 value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(AddressToUint256Map storage map, address key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(AddressToUint256Map storage map, address key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(AddressToUint256Map storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(AddressToUint256Map storage map, uint256 i) + internal + view + returns (address key, uint256 value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(AddressToUint256Map storage map, address key) + internal + view + returns (bool exists, uint256 value) + { + exists = (value = map._values[key]) != uint256(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(AddressToUint256Map storage map, address key) + internal + view + returns (uint256 value) + { + if ((value = map._values[key]) == uint256(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(AddressToUint256Map storage map) internal view returns (address[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /// @dev Adds a key-value pair to the map, or updates the value for an existing key. + /// Returns true if `key` was added to the map, that is if it was not already present. + function set(AddressToAddressMap storage map, address key, address value) + internal + returns (bool) + { + map._values[key] = value; + return EnumerableSetLib.add(map._keys, key); + } + + /// @dev Removes a key-value pair from the map. + /// Returns true if `key` was removed from the map, that is if it was present. + function remove(AddressToAddressMap storage map, address key) internal returns (bool) { + delete map._values[key]; + return EnumerableSetLib.remove(map._keys, key); + } + + /// @dev Returns true if the key is in the map. + function contains(AddressToAddressMap storage map, address key) internal view returns (bool) { + return EnumerableSetLib.contains(map._keys, key); + } + + /// @dev Returns the number of key-value pairs in the map. + function length(AddressToAddressMap storage map) internal view returns (uint256) { + return EnumerableSetLib.length(map._keys); + } + + /// @dev Returns the key-value pair at index `i`. Reverts if `i` is out-of-bounds. + function at(AddressToAddressMap storage map, uint256 i) + internal + view + returns (address key, address value) + { + value = map._values[key = EnumerableSetLib.at(map._keys, i)]; + } + + /// @dev Tries to return the value associated with the key. + function tryGet(AddressToAddressMap storage map, address key) + internal + view + returns (bool exists, address value) + { + exists = (value = map._values[key]) != address(0) || contains(map, key); + } + + /// @dev Returns the value for the key. Reverts if the key is not found. + function get(AddressToAddressMap storage map, address key) + internal + view + returns (address value) + { + if ((value = map._values[key]) == address(0)) if (!contains(map, key)) _revertNotFound(); + } + + /// @dev Returns the keys. May run out-of-gas if the map is too big. + function keys(AddressToAddressMap storage map) internal view returns (address[] memory) { + return EnumerableSetLib.values(map._keys); + } + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* PRIVATE HELPERS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev Reverts with `EnumerableMapKeyNotFound()`. + function _revertNotFound() private pure { + /// @solidity memory-safe-assembly + assembly { + mstore(0x00, 0x88682bf3) // `EnumerableMapKeyNotFound()`. + revert(0x1c, 0x04) + } + } +} diff --git a/src/utils/g/EnumerableSetLib.sol b/src/utils/g/EnumerableSetLib.sol index ca417a5f2..b9cfea853 100644 --- a/src/utils/g/EnumerableSetLib.sol +++ b/src/utils/g/EnumerableSetLib.sol @@ -620,7 +620,7 @@ library EnumerableSetLib { } } - /// @dev Returns the element at index `i` in the set. + /// @dev Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds. function at(AddressSet storage set, uint256 i) internal view returns (address result) { bytes32 rootSlot = _rootSlot(set); /// @solidity memory-safe-assembly @@ -631,7 +631,7 @@ library EnumerableSetLib { if (i >= length(set)) revert IndexOutOfBounds(); } - /// @dev Returns the element at index `i` in the set. + /// @dev Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds. function at(Bytes32Set storage set, uint256 i) internal view returns (bytes32 result) { result = _rootSlot(set); /// @solidity memory-safe-assembly @@ -642,17 +642,17 @@ library EnumerableSetLib { if (i >= length(set)) revert IndexOutOfBounds(); } - /// @dev Returns the element at index `i` in the set. + /// @dev Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds. function at(Uint256Set storage set, uint256 i) internal view returns (uint256 result) { result = uint256(at(_toBytes32Set(set), i)); } - /// @dev Returns the element at index `i` in the set. + /// @dev Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds. function at(Int256Set storage set, uint256 i) internal view returns (int256 result) { result = int256(uint256(at(_toBytes32Set(set), i))); } - /// @dev Returns the element at index `i` in the set. + /// @dev Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds. function at(Uint8Set storage set, uint256 i) internal view returns (uint8 result) { /// @solidity memory-safe-assembly assembly { diff --git a/test/EnumerableMapLib.t.sol b/test/EnumerableMapLib.t.sol new file mode 100644 index 000000000..c29fe0557 --- /dev/null +++ b/test/EnumerableMapLib.t.sol @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +import "./utils/SoladyTest.sol"; +import {EnumerableMapLib} from "../src/utils/EnumerableMapLib.sol"; + +contract EnumerableMapLibTest is SoladyTest { + using EnumerableMapLib for *; + + EnumerableMapLib.AddressToUint256Map map; + + struct _TestTemps { + bool exists; + address key; + uint256 value; + address[] keys; + } + + function testEnumerableMap(bytes32) public { + address key0 = _randomNonZeroAddress(); + address key1 = _randomNonZeroAddress(); + uint256 value0 = _random(); + uint256 value1 = _random(); + _TestTemps memory t; + + assertFalse(map.contains(key0)); + assertTrue(map.set(key0, value0)); + assertTrue(map.contains(key0)); + + if (key0 != key1) { + (t.exists, t.value) = map.tryGet(key0); + assertTrue(t.exists); + assertEq(t.value, value0); + (t.exists, t.value) = map.tryGet(key1); + assertFalse(t.exists); + assertEq(t.value, 0); + vm.expectRevert(EnumerableMapLib.EnumerableMapKeyNotFound.selector); + this.get(key1); + assertEq(this.get(key0), value0); + } + + assertEq(map.length(), 1); + assertFalse(map.set(key0, value0)); + assertEq(map.length(), 1); + + if (key0 != key1) { + assertTrue(map.set(key1, value1)); + assertEq(map.length(), 2); + + (t.key, t.value) = map.at(0); + assertEq(t.key, key0); + assertEq(t.value, value0); + (t.key, t.value) = map.at(1); + assertEq(t.key, key1); + assertEq(t.value, value1); + + t.keys = map.keys(); + assertEq(t.keys.length, 2); + assertEq(t.keys[0], key0); + assertEq(t.keys[1], key1); + + assertTrue(map.remove(key0)); + assertEq(map.length(), 1); + assertFalse(map.remove(key0)); + assertEq(map.length(), 1); + } else { + t.keys = map.keys(); + assertEq(t.keys.length, 1); + assertEq(t.keys[0], key0); + + assertTrue(map.remove(key0)); + assertEq(map.length(), 0); + assertFalse(map.remove(key0)); + assertEq(map.length(), 0); + } + } + + function get(address key) public view returns (uint256) { + return map.get(key); + } +}