Add keys() accessor to EnumerableMaps (#3920)

Co-authored-by: Francisco Giordano <frangio.1@gmail.com>

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
 

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

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

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 () {

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',