Skip to content

Commit

Permalink
Add keys() accessor to EnumerableMaps (#3920)
Browse files Browse the repository at this point in the history
Co-authored-by: Francisco Giordano <frangio.1@gmail.com>
  • Loading branch information
Amxx and frangio authored Jan 3, 2023
1 parent 2fc24fc commit 88754d0
Show file tree
Hide file tree
Showing 5 changed files with 147 additions and 10 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
108 changes: 100 additions & 8 deletions contracts/utils/structs/EnumerableMap.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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.
*/
Expand All @@ -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.
*
Expand Down Expand Up @@ -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 {
Expand All @@ -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.
*/
Expand All @@ -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.
*
Expand Down Expand Up @@ -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 {
Expand All @@ -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.
*/
Expand All @@ -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.
*
Expand Down Expand Up @@ -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 {
Expand All @@ -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.
*/
Expand All @@ -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.
*
Expand Down Expand Up @@ -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;
}
}
36 changes: 34 additions & 2 deletions scripts/generate/templates/EnumerableMap.js
Original file line number Diff line number Diff line change
Expand Up @@ -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 }) => `\
Expand All @@ -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.
*/
Expand All @@ -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.
*
Expand Down Expand Up @@ -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
Expand Down
7 changes: 7 additions & 0 deletions test/utils/structs/EnumerableMap.behavior.js
Original file line number Diff line number Diff line change
Expand Up @@ -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 () {
Expand Down
5 changes: 5 additions & 0 deletions test/utils/structs/EnumerableMap.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand All @@ -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',
Expand All @@ -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',
Expand All @@ -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',
Expand All @@ -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',
Expand Down

0 comments on commit 88754d0

Please sign in to comment.