From f88e5552342c9f9afc8ec76f833e281ca748b960 Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Fri, 16 Jul 2021 17:06:47 +0200 Subject: [PATCH] Add values() functions to EnumerableSets (#2768) Co-authored-by: Francisco Giordano --- CHANGELOG.md | 1 + contracts/mocks/EnumerableSetMock.sol | 12 ++++ contracts/utils/structs/EnumerableSet.sol | 62 ++++++++++++++++++++ test/utils/structs/EnumerableSet.behavior.js | 25 +++++--- 4 files changed, 92 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 67ae86dcc3b..2104d110ffc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ ## Unreleased * `ERC2771Context`: use private variable from storage to store the forwarder address. Fixes issues where `_msgSender()` was not callable from constructors. ([#2754](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2754)) + * `EnumerableSet`: add `values()` functions that returns an array containing all values in a single call. ([#2768](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2768)) ## 4.2.0 (2021-06-30) diff --git a/contracts/mocks/EnumerableSetMock.sol b/contracts/mocks/EnumerableSetMock.sol index 26f6f60c21f..922ce46d250 100644 --- a/contracts/mocks/EnumerableSetMock.sol +++ b/contracts/mocks/EnumerableSetMock.sol @@ -33,6 +33,10 @@ contract EnumerableBytes32SetMock { function at(uint256 index) public view returns (bytes32) { return _set.at(index); } + + function values() public view returns (bytes32[] memory) { + return _set.values(); + } } // AddressSet @@ -64,6 +68,10 @@ contract EnumerableAddressSetMock { function at(uint256 index) public view returns (address) { return _set.at(index); } + + function values() public view returns (address[] memory) { + return _set.values(); + } } // UintSet @@ -95,4 +103,8 @@ contract EnumerableUintSetMock { function at(uint256 index) public view returns (uint256) { return _set.at(index); } + + function values() public view returns (uint256[] memory) { + return _set.values(); + } } diff --git a/contracts/utils/structs/EnumerableSet.sol b/contracts/utils/structs/EnumerableSet.sol index e90d55dedaf..e3cf13f12e3 100644 --- a/contracts/utils/structs/EnumerableSet.sol +++ b/contracts/utils/structs/EnumerableSet.sol @@ -130,6 +130,18 @@ library EnumerableSet { return set._values[index]; } + /** + * @dev Return the entire set in an array + * + * 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 set grows to a point where copying to memory consumes too much gas to fit in a block. + */ + function _values(Set storage set) private view returns (bytes32[] memory) { + return set._values; + } + // Bytes32Set struct Bytes32Set { @@ -184,6 +196,18 @@ library EnumerableSet { return _at(set._inner, index); } + /** + * @dev Return the entire set in an array + * + * 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 set grows to a point where copying to memory consumes too much gas to fit in a block. + */ + function values(Bytes32Set storage set) internal view returns (bytes32[] memory) { + return _values(set._inner); + } + // AddressSet struct AddressSet { @@ -238,6 +262,25 @@ library EnumerableSet { return address(uint160(uint256(_at(set._inner, index)))); } + /** + * @dev Return the entire set in an array + * + * 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 set grows to a point where copying to memory consumes too much gas to fit in a block. + */ + function values(AddressSet storage set) internal view returns (address[] memory) { + bytes32[] memory store = _values(set._inner); + address[] memory result; + + assembly { + result := store + } + + return result; + } + // UintSet struct UintSet { @@ -291,4 +334,23 @@ library EnumerableSet { function at(UintSet storage set, uint256 index) internal view returns (uint256) { return uint256(_at(set._inner, index)); } + + /** + * @dev Return the entire set in an array + * + * 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 set grows to a point where copying to memory consumes too much gas to fit in a block. + */ + function values(UintSet storage set) internal view returns (uint256[] memory) { + bytes32[] memory store = _values(set._inner); + uint256[] memory result; + + assembly { + result := store + } + + return result; + } } diff --git a/test/utils/structs/EnumerableSet.behavior.js b/test/utils/structs/EnumerableSet.behavior.js index be3f52a5e5a..17e08667129 100644 --- a/test/utils/structs/EnumerableSet.behavior.js +++ b/test/utils/structs/EnumerableSet.behavior.js @@ -3,18 +3,27 @@ const { expect } = require('chai'); function shouldBehaveLikeSet (valueA, valueB, valueC) { async function expectMembersMatch (set, values) { - await Promise.all(values.map(async value => - expect(await set.contains(value)).to.equal(true), - )); + const contains = await Promise.all(values.map(value => set.contains(value))); + expect(contains.every(Boolean)).to.be.equal(true); - expect(await set.length()).to.bignumber.equal(values.length.toString()); + const length = await set.length(); + expect(length).to.bignumber.equal(values.length.toString()); // To compare values we convert to strings to workaround Chai // limitations when dealing with nested arrays (required for BNs) - expect(await Promise.all([...Array(values.length).keys()].map(async (index) => { - const entry = await set.at(index); - return entry.toString(); - }))).to.have.same.members(values.map(v => v.toString())); + const indexedValues = await Promise.all(Array(values.length).fill().map((_, index) => set.at(index))); + expect( + indexedValues.map(v => v.toString()), + ).to.have.same.members( + values.map(v => v.toString()), + ); + + const returnedValues = await set.values(); + expect( + returnedValues.map(v => v.toString()), + ).to.have.same.members( + values.map(v => v.toString()), + ); } it('starts empty', async function () {