Forráskód Böngészése

Add values() functions to EnumerableSets (#2768)

Co-authored-by: Francisco Giordano <frangio.1@gmail.com>
Hadrien Croubois 4 éve
szülő
commit
f88e555234

+ 1 - 0
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)
 

+ 12 - 0
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();
+    }
 }

+ 62 - 0
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;
+    }
 }

+ 17 - 8
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 () {