Skip to content

Commit

Permalink
Add values() functions to EnumerableSets (#2768)
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 Jul 16, 2021
1 parent 87826f8 commit f88e555
Show file tree
Hide file tree
Showing 4 changed files with 92 additions and 8 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
12 changes: 12 additions & 0 deletions contracts/mocks/EnumerableSetMock.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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();
}
}
62 changes: 62 additions & 0 deletions contracts/utils/structs/EnumerableSet.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
}
}
25 changes: 17 additions & 8 deletions test/utils/structs/EnumerableSet.behavior.js
Original file line number Diff line number Diff line change
Expand Up @@ -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 () {
Expand Down

0 comments on commit f88e555

Please sign in to comment.