From 88754d0b36cde4e6b1daab591058eaef449f5adb Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Tue, 3 Jan 2023 22:25:37 +0100 Subject: [PATCH] Add keys() accessor to EnumerableMaps (#3920) Co-authored-by: Francisco Giordano --- CHANGELOG.md | 1 + contracts/utils/structs/EnumerableMap.sol | 108 +++++++++++++++++-- scripts/generate/templates/EnumerableMap.js | 36 ++++++- test/utils/structs/EnumerableMap.behavior.js | 7 ++ test/utils/structs/EnumerableMap.test.js | 5 + 5 files changed, 147 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 58f21d9990a..7f6a9be4002 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ * `Strings`: add `equal` method. ([#3774](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3774)) * `Strings`: add `toString` method for signed integers. ([#3773](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3773)) * `MerkleProof`: optimize by using unchecked arithmetic. ([#3745](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3745)) + * `EnumerableMap`: add a `keys()` function that returns an array containing all the keys. ([#3920](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3920)) ### Deprecations diff --git a/contracts/utils/structs/EnumerableMap.sol b/contracts/utils/structs/EnumerableMap.sol index d359671f461..8b188c7344f 100644 --- a/contracts/utils/structs/EnumerableMap.sol +++ b/contracts/utils/structs/EnumerableMap.sol @@ -156,6 +156,18 @@ library EnumerableMap { return value; } + /** + * @dev Return the an array containing all the keys + * + * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed + * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that + * this function has an unbounded cost, and using it as part of a state-changing function may render the function + * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block. + */ + function keys(Bytes32ToBytes32Map storage map) internal view returns (bytes32[] memory) { + return map._keys.values(); + } + // UintToUintMap struct UintToUintMap { @@ -174,7 +186,7 @@ library EnumerableMap { } /** - * @dev Removes a value from a set. O(1). + * @dev Removes a value from a map. O(1). * * Returns true if the key was removed from the map, that is if it was present. */ @@ -197,7 +209,7 @@ library EnumerableMap { } /** - * @dev Returns the element stored at position `index` in the set. O(1). + * @dev Returns the element stored at position `index` in the map. O(1). * Note that there are no guarantees on the ordering of values inside the * array, and it may change when more values are added or removed. * @@ -240,6 +252,26 @@ library EnumerableMap { return uint256(get(map._inner, bytes32(key), errorMessage)); } + /** + * @dev Return the an array containing all the keys + * + * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed + * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that + * this function has an unbounded cost, and using it as part of a state-changing function may render the function + * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block. + */ + function keys(UintToUintMap storage map) internal view returns (uint256[] memory) { + bytes32[] memory store = keys(map._inner); + uint256[] memory result; + + /// @solidity memory-safe-assembly + assembly { + result := store + } + + return result; + } + // UintToAddressMap struct UintToAddressMap { @@ -258,7 +290,7 @@ library EnumerableMap { } /** - * @dev Removes a value from a set. O(1). + * @dev Removes a value from a map. O(1). * * Returns true if the key was removed from the map, that is if it was present. */ @@ -281,7 +313,7 @@ library EnumerableMap { } /** - * @dev Returns the element stored at position `index` in the set. O(1). + * @dev Returns the element stored at position `index` in the map. O(1). * Note that there are no guarantees on the ordering of values inside the * array, and it may change when more values are added or removed. * @@ -328,6 +360,26 @@ library EnumerableMap { return address(uint160(uint256(get(map._inner, bytes32(key), errorMessage)))); } + /** + * @dev Return the an array containing all the keys + * + * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed + * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that + * this function has an unbounded cost, and using it as part of a state-changing function may render the function + * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block. + */ + function keys(UintToAddressMap storage map) internal view returns (uint256[] memory) { + bytes32[] memory store = keys(map._inner); + uint256[] memory result; + + /// @solidity memory-safe-assembly + assembly { + result := store + } + + return result; + } + // AddressToUintMap struct AddressToUintMap { @@ -346,7 +398,7 @@ library EnumerableMap { } /** - * @dev Removes a value from a set. O(1). + * @dev Removes a value from a map. O(1). * * Returns true if the key was removed from the map, that is if it was present. */ @@ -369,7 +421,7 @@ library EnumerableMap { } /** - * @dev Returns the element stored at position `index` in the set. O(1). + * @dev Returns the element stored at position `index` in the map. O(1). * Note that there are no guarantees on the ordering of values inside the * array, and it may change when more values are added or removed. * @@ -416,6 +468,26 @@ library EnumerableMap { return uint256(get(map._inner, bytes32(uint256(uint160(key))), errorMessage)); } + /** + * @dev Return the an array containing all the keys + * + * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed + * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that + * this function has an unbounded cost, and using it as part of a state-changing function may render the function + * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block. + */ + function keys(AddressToUintMap storage map) internal view returns (address[] memory) { + bytes32[] memory store = keys(map._inner); + address[] memory result; + + /// @solidity memory-safe-assembly + assembly { + result := store + } + + return result; + } + // Bytes32ToUintMap struct Bytes32ToUintMap { @@ -434,7 +506,7 @@ library EnumerableMap { } /** - * @dev Removes a value from a set. O(1). + * @dev Removes a value from a map. O(1). * * Returns true if the key was removed from the map, that is if it was present. */ @@ -457,7 +529,7 @@ library EnumerableMap { } /** - * @dev Returns the element stored at position `index` in the set. O(1). + * @dev Returns the element stored at position `index` in the map. O(1). * Note that there are no guarantees on the ordering of values inside the * array, and it may change when more values are added or removed. * @@ -503,4 +575,24 @@ library EnumerableMap { ) internal view returns (uint256) { return uint256(get(map._inner, key, errorMessage)); } + + /** + * @dev Return the an array containing all the keys + * + * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed + * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that + * this function has an unbounded cost, and using it as part of a state-changing function may render the function + * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block. + */ + function keys(Bytes32ToUintMap storage map) internal view returns (bytes32[] memory) { + bytes32[] memory store = keys(map._inner); + bytes32[] memory result; + + /// @solidity memory-safe-assembly + assembly { + result := store + } + + return result; + } } diff --git a/scripts/generate/templates/EnumerableMap.js b/scripts/generate/templates/EnumerableMap.js index ca8e0e77de4..dbd502a0b58 100644 --- a/scripts/generate/templates/EnumerableMap.js +++ b/scripts/generate/templates/EnumerableMap.js @@ -168,6 +168,18 @@ function get( require(value != 0 || contains(map, key), errorMessage); return value; } + +/** + * @dev Return the an array containing all the keys + * + * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed + * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that + * this function has an unbounded cost, and using it as part of a state-changing function may render the function + * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block. + */ +function keys(Bytes32ToBytes32Map storage map) internal view returns (bytes32[] memory) { + return map._keys.values(); +} `; const customMap = ({ name, keyType, valueType }) => `\ @@ -193,7 +205,7 @@ function set( } /** - * @dev Removes a value from a set. O(1). + * @dev Removes a value from a map. O(1). * * Returns true if the key was removed from the map, that is if it was present. */ @@ -216,7 +228,7 @@ function length(${name} storage map) internal view returns (uint256) { } /** - * @dev Returns the element stored at position \`index\` in the set. O(1). + * @dev Returns the element stored at position \`index\` in the map. O(1). * Note that there are no guarantees on the ordering of values inside the * array, and it may change when more values are added or removed. * @@ -262,6 +274,26 @@ function get( ) internal view returns (${valueType}) { return ${fromBytes32(valueType, `get(map._inner, ${toBytes32(keyType, 'key')}, errorMessage)`)}; } + +/** + * @dev Return the an array containing all the keys + * + * WARNING: This operation will copy the entire storage to memory, which can be quite expensive. This is designed + * to mostly be used by view accessors that are queried without any gas fees. Developers should keep in mind that + * this function has an unbounded cost, and using it as part of a state-changing function may render the function + * uncallable if the map grows to a point where copying to memory consumes too much gas to fit in a block. + */ +function keys(${name} storage map) internal view returns (${keyType}[] memory) { + bytes32[] memory store = keys(map._inner); + ${keyType}[] memory result; + + /// @solidity memory-safe-assembly + assembly { + result := store + } + + return result; +} `; // GENERATE diff --git a/test/utils/structs/EnumerableMap.behavior.js b/test/utils/structs/EnumerableMap.behavior.js index 4d698f3341c..25ae8421580 100644 --- a/test/utils/structs/EnumerableMap.behavior.js +++ b/test/utils/structs/EnumerableMap.behavior.js @@ -36,6 +36,13 @@ function shouldBehaveLikeMap ( }))).to.have.same.deep.members( zip(keys.map(k => k.toString()), values.map(v => v.toString())), ); + + // This also checks that both arrays have the same length + expect( + (await methods.keys(map)).map(k => k.toString()), + ).to.have.same.members( + keys.map(key => key.toString()), + ); } it('starts empty', async function () { diff --git a/test/utils/structs/EnumerableMap.test.js b/test/utils/structs/EnumerableMap.test.js index 2587f75955f..14f31b4294b 100644 --- a/test/utils/structs/EnumerableMap.test.js +++ b/test/utils/structs/EnumerableMap.test.js @@ -39,6 +39,7 @@ contract('EnumerableMap', function (accounts) { length: '$length_EnumerableMap_AddressToUintMap(uint256)', at: '$at_EnumerableMap_AddressToUintMap(uint256,uint256)', contains: '$contains(uint256,address)', + keys: '$keys_EnumerableMap_AddressToUintMap(uint256)', }), { setReturn: 'return$set_EnumerableMap_AddressToUintMap_address_uint256', @@ -62,6 +63,7 @@ contract('EnumerableMap', function (accounts) { length: '$length_EnumerableMap_UintToAddressMap(uint256)', at: '$at_EnumerableMap_UintToAddressMap(uint256,uint256)', contains: '$contains_EnumerableMap_UintToAddressMap(uint256,uint256)', + keys: '$keys_EnumerableMap_UintToAddressMap(uint256)', }), { setReturn: 'return$set_EnumerableMap_UintToAddressMap_uint256_address', @@ -85,6 +87,7 @@ contract('EnumerableMap', function (accounts) { length: '$length_EnumerableMap_Bytes32ToBytes32Map(uint256)', at: '$at_EnumerableMap_Bytes32ToBytes32Map(uint256,uint256)', contains: '$contains_EnumerableMap_Bytes32ToBytes32Map(uint256,bytes32)', + keys: '$keys_EnumerableMap_Bytes32ToBytes32Map(uint256)', }), { setReturn: 'return$set_EnumerableMap_Bytes32ToBytes32Map_bytes32_bytes32', @@ -108,6 +111,7 @@ contract('EnumerableMap', function (accounts) { length: '$length_EnumerableMap_UintToUintMap(uint256)', at: '$at_EnumerableMap_UintToUintMap(uint256,uint256)', contains: '$contains_EnumerableMap_UintToUintMap(uint256,uint256)', + keys: '$keys_EnumerableMap_UintToUintMap(uint256)', }), { setReturn: 'return$set_EnumerableMap_UintToUintMap_uint256_uint256', @@ -131,6 +135,7 @@ contract('EnumerableMap', function (accounts) { length: '$length_EnumerableMap_Bytes32ToUintMap(uint256)', at: '$at_EnumerableMap_Bytes32ToUintMap(uint256,uint256)', contains: '$contains_EnumerableMap_Bytes32ToUintMap(uint256,bytes32)', + keys: '$keys_EnumerableMap_Bytes32ToUintMap(uint256)', }), { setReturn: 'return$set_EnumerableMap_Bytes32ToUintMap_bytes32_uint256',