Skip to content

Commit

Permalink
Add bytes32 to bytes32 enumerable map (#3192)
Browse files Browse the repository at this point in the history 
* feat(enumerablemap): add bytes32 to bytes32 map

* chore(changelog): edit CHANGELOG

* feat(enumerable map): edit struct visibility
  • Loading branch information
@wiasliaw
wiasliaw committed Mar 22, 2022
1 parent c028c56 commit b13bdb0
Show file tree
Hide file tree
Showing 8 changed files with 334 additions and 378 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Expand Up @@ -4,6 +4,7 @@

* `AccessControl`: add a virtual `_checkRole(bytes32)` function that can be overridden to alter the `onlyRole` modifier behavior. ([#3137](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3137))
* `EnumerableMap`: add new `AddressToUintMap` map type. ([#3150](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3150))
* `EnumerableMap`: add new `Bytes32ToBytes32Map` map type. ([#3192](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3192))
* `ERC1155`: Add a `_afterTokenTransfer` hook for improved extensibility. ([#3166](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3166))
* `DoubleEndedQueue`: a new data structure that supports efficient push and pop to both front and back, useful for FIFO and LIFO queues. ([#3153](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3153))
* `Governor`: improved security of `onlyGovernance` modifier when using an external executor contract (e.g. a timelock) that can operate without necessarily going through the governance protocol. ([#3147](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3147))
Expand Down
46 changes: 46 additions & 0 deletions contracts/mocks/EnumerableMapMock.sol
View file
Expand Up @@ -84,4 +84,50 @@ contract AddressToUintMapMock {
function get(address key) public view returns (uint256) {
return _map.get(key);
}

function getWithMessage(address key, string calldata errorMessage) public view returns (uint256) {
return _map.get(key, errorMessage);
}
}

contract Bytes32ToBytes32MapMock {
using EnumerableMap for EnumerableMap.Bytes32ToBytes32Map;

event OperationResult(bool result);

EnumerableMap.Bytes32ToBytes32Map private _map;

function contains(bytes32 key) public view returns (bool) {
return _map.contains(key);
}

function set(bytes32 key, bytes32 value) public {
bool result = _map.set(key, value);
emit OperationResult(result);
}

function remove(bytes32 key) public {
bool result = _map.remove(key);
emit OperationResult(result);
}

function length() public view returns (uint256) {
return _map.length();
}

function at(uint256 index) public view returns (bytes32 key, bytes32 value) {
return _map.at(index);
}

function tryGet(bytes32 key) public view returns (bool, bytes32) {
return _map.tryGet(key);
}

function get(bytes32 key) public view returns (bytes32) {
return _map.get(key);
}

function getWithMessage(bytes32 key, string calldata errorMessage) public view returns (bytes32) {
return _map.get(key, errorMessage);
}
}
81 changes: 48 additions & 33 deletions contracts/utils/structs/EnumerableMap.sol
View file
Expand Up @@ -30,6 +30,7 @@ import "./EnumerableSet.sol";
*
* - `uint256 -> address` (`UintToAddressMap`) since v3.0.0
* - `address -> uint256` (`AddressToUintMap`) since v4.6.0
* - `bytes32 -> bytes32` (`Bytes32ToBytes32`) since v4.6.0
*/
library EnumerableMap {
using EnumerableSet for EnumerableSet.Bytes32Set;
Expand All @@ -43,7 +44,7 @@ library EnumerableMap {
// This means that we can only create new EnumerableMaps for types that fit
// in bytes32.

struct Map {
struct Bytes32ToBytes32Map {
// Storage of keys
EnumerableSet.Bytes32Set _keys;
mapping(bytes32 => bytes32) _values;
Expand All @@ -56,11 +57,11 @@ library EnumerableMap {
* Returns true if the key was added to the map, that is if it was not
* already present.
*/
function _set(
Map storage map,
function set(
Bytes32ToBytes32Map storage map,
bytes32 key,
bytes32 value
) private returns (bool) {
) internal returns (bool) {
map._values[key] = value;
return map._keys.add(key);
}
Expand All @@ -70,22 +71,22 @@ library EnumerableMap {
*
* Returns true if the key was removed from the map, that is if it was present.
*/
function _remove(Map storage map, bytes32 key) private returns (bool) {
function remove(Bytes32ToBytes32Map storage map, bytes32 key) internal returns (bool) {
delete map._values[key];
return map._keys.remove(key);
}

/**
* @dev Returns true if the key is in the map. O(1).
*/
function _contains(Map storage map, bytes32 key) private view returns (bool) {
function contains(Bytes32ToBytes32Map storage map, bytes32 key) internal view returns (bool) {
return map._keys.contains(key);
}

/**
* @dev Returns the number of key-value pairs in the map. O(1).
*/
function _length(Map storage map) private view returns (uint256) {
function length(Bytes32ToBytes32Map storage map) internal view returns (uint256) {
return map._keys.length();
}

Expand All @@ -99,7 +100,7 @@ library EnumerableMap {
*
* - `index` must be strictly less than {length}.
*/
function _at(Map storage map, uint256 index) private view returns (bytes32, bytes32) {
function at(Bytes32ToBytes32Map storage map, uint256 index) internal view returns (bytes32, bytes32) {
bytes32 key = map._keys.at(index);
return (key, map._values[key]);
}
Expand All @@ -108,10 +109,10 @@ library EnumerableMap {
* @dev Tries to returns the value associated with `key`. O(1).
* Does not revert if `key` is not in the map.
*/
function _tryGet(Map storage map, bytes32 key) private view returns (bool, bytes32) {
function tryGet(Bytes32ToBytes32Map storage map, bytes32 key) internal view returns (bool, bytes32) {
bytes32 value = map._values[key];
if (value == bytes32(0)) {
return (_contains(map, key), bytes32(0));
return (contains(map, key), bytes32(0));
} else {
return (true, value);
}
Expand All @@ -124,9 +125,9 @@ library EnumerableMap {
*
* - `key` must be in the map.
*/
function _get(Map storage map, bytes32 key) private view returns (bytes32) {
function get(Bytes32ToBytes32Map storage map, bytes32 key) internal view returns (bytes32) {
bytes32 value = map._values[key];
require(value != 0 || _contains(map, key), "EnumerableMap: nonexistent key");
require(value != 0 || contains(map, key), "EnumerableMap: nonexistent key");
return value;
}

Expand All @@ -136,20 +137,20 @@ library EnumerableMap {
* CAUTION: This function is deprecated because it requires allocating memory for the error
* message unnecessarily. For custom revert reasons use {_tryGet}.
*/
function _get(
Map storage map,
function get(
Bytes32ToBytes32Map storage map,
bytes32 key,
string memory errorMessage
) private view returns (bytes32) {
) internal view returns (bytes32) {
bytes32 value = map._values[key];
require(value != 0 || _contains(map, key), errorMessage);
require(value != 0 || contains(map, key), errorMessage);
return value;
}

// UintToAddressMap

struct UintToAddressMap {
Map _inner;
Bytes32ToBytes32Map _inner;
}

/**
Expand All @@ -164,7 +165,7 @@ library EnumerableMap {
uint256 key,
address value
) internal returns (bool) {
return _set(map._inner, bytes32(key), bytes32(uint256(uint160(value))));
return set(map._inner, bytes32(key), bytes32(uint256(uint160(value))));
}

/**
Expand All @@ -173,21 +174,21 @@ library EnumerableMap {
* Returns true if the key was removed from the map, that is if it was present.
*/
function remove(UintToAddressMap storage map, uint256 key) internal returns (bool) {
return _remove(map._inner, bytes32(key));
return remove(map._inner, bytes32(key));
}

/**
* @dev Returns true if the key is in the map. O(1).
*/
function contains(UintToAddressMap storage map, uint256 key) internal view returns (bool) {
return _contains(map._inner, bytes32(key));
return contains(map._inner, bytes32(key));
}

/**
* @dev Returns the number of elements in the map. O(1).
*/
function length(UintToAddressMap storage map) internal view returns (uint256) {
return _length(map._inner);
return length(map._inner);
}

/**
Expand All @@ -200,7 +201,7 @@ library EnumerableMap {
* - `index` must be strictly less than {length}.
*/
function at(UintToAddressMap storage map, uint256 index) internal view returns (uint256, address) {
(bytes32 key, bytes32 value) = _at(map._inner, index);
(bytes32 key, bytes32 value) = at(map._inner, index);
return (uint256(key), address(uint160(uint256(value))));
}

Expand All @@ -211,7 +212,7 @@ library EnumerableMap {
* _Available since v3.4._
*/
function tryGet(UintToAddressMap storage map, uint256 key) internal view returns (bool, address) {
(bool success, bytes32 value) = _tryGet(map._inner, bytes32(key));
(bool success, bytes32 value) = tryGet(map._inner, bytes32(key));
return (success, address(uint160(uint256(value))));
}

Expand All @@ -223,7 +224,7 @@ library EnumerableMap {
* - `key` must be in the map.
*/
function get(UintToAddressMap storage map, uint256 key) internal view returns (address) {
return address(uint160(uint256(_get(map._inner, bytes32(key)))));
return address(uint160(uint256(get(map._inner, bytes32(key)))));
}

/**
Expand All @@ -237,13 +238,13 @@ library EnumerableMap {
uint256 key,
string memory errorMessage
) internal view returns (address) {
return address(uint160(uint256(_get(map._inner, bytes32(key), errorMessage))));
return address(uint160(uint256(get(map._inner, bytes32(key), errorMessage))));
}

// AddressToUintMap

struct AddressToUintMap {
Map _inner;
Bytes32ToBytes32Map _inner;
}

/**
Expand All @@ -258,7 +259,7 @@ library EnumerableMap {
address key,
uint256 value
) internal returns (bool) {
return _set(map._inner, bytes32(uint256(uint160(key))), bytes32(value));
return set(map._inner, bytes32(uint256(uint160(key))), bytes32(value));
}

/**
Expand All @@ -267,21 +268,21 @@ library EnumerableMap {
* Returns true if the key was removed from the map, that is if it was present.
*/
function remove(AddressToUintMap storage map, address key) internal returns (bool) {
return _remove(map._inner, bytes32(uint256(uint160(key))));
return remove(map._inner, bytes32(uint256(uint160(key))));
}

/**
* @dev Returns true if the key is in the map. O(1).
*/
function contains(AddressToUintMap storage map, address key) internal view returns (bool) {
return _contains(map._inner, bytes32(uint256(uint160(key))));
return contains(map._inner, bytes32(uint256(uint160(key))));
}

/**
* @dev Returns the number of elements in the map. O(1).
*/
function length(AddressToUintMap storage map) internal view returns (uint256) {
return _length(map._inner);
return length(map._inner);
}

/**
Expand All @@ -294,7 +295,7 @@ library EnumerableMap {
* - `index` must be strictly less than {length}.
*/
function at(AddressToUintMap storage map, uint256 index) internal view returns (address, uint256) {
(bytes32 key, bytes32 value) = _at(map._inner, index);
(bytes32 key, bytes32 value) = at(map._inner, index);
return (address(uint160(uint256(key))), uint256(value));
}

Expand All @@ -305,7 +306,7 @@ library EnumerableMap {
* _Available since v3.4._
*/
function tryGet(AddressToUintMap storage map, address key) internal view returns (bool, uint256) {
(bool success, bytes32 value) = _tryGet(map._inner, bytes32(uint256(uint160(key))));
(bool success, bytes32 value) = tryGet(map._inner, bytes32(uint256(uint160(key))));
return (success, uint256(value));
}

Expand All @@ -317,6 +318,20 @@ library EnumerableMap {
* - `key` must be in the map.
*/
function get(AddressToUintMap storage map, address key) internal view returns (uint256) {
return uint256(_get(map._inner, bytes32(uint256(uint160(key)))));
return uint256(get(map._inner, bytes32(uint256(uint160(key)))));
}

/**
* @dev Same as {get}, with a custom error message when `key` is not in the map.
*
* CAUTION: This function is deprecated because it requires allocating memory for the error
* message unnecessarily. For custom revert reasons use {tryGet}.
*/
function get(
AddressToUintMap storage map,
address key,
string memory errorMessage
) internal view returns (uint256) {
return uint256(get(map._inner, bytes32(uint256(uint160(key))), errorMessage));
}
}

0 comments on commit b13bdb0

Please sign in to comment.