Update EnumerableMap.sol

OpenZeppelin · May 20, 2022 · 25f7c16 · 25f7c16
1 parent 29bb077
commit 25f7c16
Showing 1 changed file with 61 additions and 61 deletions.
diff --git a/contracts/utils/structs/EnumerableMap.sol b/contracts/utils/structs/EnumerableMap.sol
@@ -117,8 +117,6 @@ library EnumerableMap {
     /**
      * @dev Tries to returns the value associated with `key`.  O(1).
      * Does not revert if `key` is not in the map.
-     *
-     * _Available since v3.4._
      */
     function tryGet(Bytes32ToBytes32Map storage map, bytes32 key) internal view returns (bool, bytes32) {
         bytes32 value = map._values[key];
@@ -158,9 +156,9 @@ library EnumerableMap {
         return value;
     }
 
-    // UintToAddressMap
+    // UintToUintMap
 
-    struct UintToAddressMap {
+    struct UintToUintMap {
         Bytes32ToBytes32Map _inner;
     }
 
@@ -172,33 +170,33 @@ library EnumerableMap {
      * already present.
      */
     function set(
-        UintToAddressMap storage map,
+        UintToUintMap storage map,
         uint256 key,
-        address value
+        uint256 value
     ) internal returns (bool) {
-        return set(map._inner, bytes32(key), bytes32(uint256(uint160(value))));
+        return set(map._inner, bytes32(key), bytes32(value));
     }
 
     /**
      * @dev Removes a value from a set. O(1).
      *
      * 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) {
+    function remove(UintToUintMap storage map, uint256 key) internal returns (bool) {
         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) {
+    function contains(UintToUintMap storage map, uint256 key) internal view returns (bool) {
         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) {
+    function length(UintToUintMap storage map) internal view returns (uint256) {
         return length(map._inner);
     }
 
@@ -211,18 +209,18 @@ library EnumerableMap {
      *
      * - `index` must be strictly less than {length}.
      */
-    function at(UintToAddressMap storage map, uint256 index) internal view returns (uint256, address) {
+    function at(UintToUintMap storage map, uint256 index) internal view returns (uint256, uint256) {
         (bytes32 key, bytes32 value) = at(map._inner, index);
-        return (uint256(key), address(uint160(uint256(value))));
+        return (uint256(key), uint256(value));
     }
 
     /**
      * @dev Tries to returns the value associated with `key`.  O(1).
      * Does not revert if `key` is not in the map.
      */
-    function tryGet(UintToAddressMap storage map, uint256 key) internal view returns (bool, address) {
+    function tryGet(UintToUintMap storage map, uint256 key) internal view returns (bool, uint256) {
         (bool success, bytes32 value) = tryGet(map._inner, bytes32(key));
-        return (success, address(uint160(uint256(value))));
+        return (success, uint256(value));
     }
 
     /**
@@ -232,8 +230,8 @@ 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)))));
+    function get(UintToUintMap storage map, uint256 key) internal view returns (uint256) {
+        return uint256(get(map._inner, bytes32(key)));
     }
 
     /**
@@ -243,16 +241,16 @@ library EnumerableMap {
      * message unnecessarily. For custom revert reasons use {tryGet}.
      */
     function get(
-        UintToAddressMap storage map,
+        UintToUintMap storage map,
         uint256 key,
         string memory errorMessage
-    ) internal view returns (address) {
-        return address(uint160(uint256(get(map._inner, bytes32(key), errorMessage))));
+    ) internal view returns (uint256) {
+        return uint256(get(map._inner, bytes32(key), errorMessage));
     }
 
-    // AddressToUintMap
+    // UintToAddressMap
 
-    struct AddressToUintMap {
+    struct UintToAddressMap {
         Bytes32ToBytes32Map _inner;
     }
 
@@ -264,33 +262,33 @@ library EnumerableMap {
      * already present.
      */
     function set(
-        AddressToUintMap storage map,
-        address key,
-        uint256 value
+        UintToAddressMap storage map,
+        uint256 key,
+        address value
     ) internal returns (bool) {
-        return set(map._inner, bytes32(uint256(uint160(key))), bytes32(value));
+        return set(map._inner, bytes32(key), bytes32(uint256(uint160(value))));
     }
 
     /**
      * @dev Removes a value from a set. O(1).
      *
      * 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))));
+    function remove(UintToAddressMap storage map, uint256 key) internal returns (bool) {
+        return remove(map._inner, bytes32(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))));
+    function contains(UintToAddressMap storage map, uint256 key) internal view returns (bool) {
+        return contains(map._inner, bytes32(key));
     }
 
     /**
      * @dev Returns the number of elements in the map. O(1).
      */
-    function length(AddressToUintMap storage map) internal view returns (uint256) {
+    function length(UintToAddressMap storage map) internal view returns (uint256) {
         return length(map._inner);
     }
 
@@ -303,18 +301,20 @@ library EnumerableMap {
      *
      * - `index` must be strictly less than {length}.
      */
-    function at(AddressToUintMap storage map, uint256 index) internal view returns (address, uint256) {
+    function at(UintToAddressMap storage map, uint256 index) internal view returns (uint256, address) {
         (bytes32 key, bytes32 value) = at(map._inner, index);
-        return (address(uint160(uint256(key))), uint256(value));
+        return (uint256(key), address(uint160(uint256(value))));
     }
 
     /**
      * @dev Tries to returns the value associated with `key`.  O(1).
      * Does not revert if `key` is not in the map.
+     *
+     * _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))));
-        return (success, uint256(value));
+    function tryGet(UintToAddressMap storage map, uint256 key) internal view returns (bool, address) {
+        (bool success, bytes32 value) = tryGet(map._inner, bytes32(key));
+        return (success, address(uint160(uint256(value))));
     }
 
     /**
@@ -324,8 +324,8 @@ 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)))));
+    function get(UintToAddressMap storage map, uint256 key) internal view returns (address) {
+        return address(uint160(uint256(get(map._inner, bytes32(key)))));
     }
 
     /**
@@ -335,16 +335,16 @@ library EnumerableMap {
      * message unnecessarily. For custom revert reasons use {tryGet}.
      */
     function get(
-        AddressToUintMap storage map,
-        address key,
+        UintToAddressMap storage map,
+        uint256 key,
         string memory errorMessage
-    ) internal view returns (uint256) {
-        return uint256(get(map._inner, bytes32(uint256(uint160(key))), errorMessage));
+    ) internal view returns (address) {
+        return address(uint160(uint256(get(map._inner, bytes32(key), errorMessage))));
     }
-
-    // UintToUintMap
 
-    struct UintToUintMap {
+    // AddressToUintMap
+
+    struct AddressToUintMap {
         Bytes32ToBytes32Map _inner;
     }
 
@@ -356,33 +356,33 @@ library EnumerableMap {
      * already present.
      */
     function set(
-        UintToUintMap storage map,
-        uint256 key,
+        AddressToUintMap storage map,
+        address key,
         uint256 value
     ) internal returns (bool) {
-        return set(map._inner, bytes32(key), bytes32(value));
+        return set(map._inner, bytes32(uint256(uint160(key))), bytes32(value));
     }
 
     /**
      * @dev Removes a value from a set. O(1).
      *
      * Returns true if the key was removed from the map, that is if it was present.
      */
-    function remove(UintToUintMap storage map, uint256 key) internal returns (bool) {
-        return remove(map._inner, bytes32(key));
+    function remove(AddressToUintMap storage map, address key) internal returns (bool) {
+        return remove(map._inner, bytes32(uint256(uint160(key))));
     }
 
     /**
      * @dev Returns true if the key is in the map. O(1).
      */
-    function contains(UintToUintMap storage map, uint256 key) internal view returns (bool) {
-        return contains(map._inner, bytes32(key));
+    function contains(AddressToUintMap storage map, address key) internal view returns (bool) {
+        return contains(map._inner, bytes32(uint256(uint160(key))));
     }
 
     /**
      * @dev Returns the number of elements in the map. O(1).
      */
-    function length(UintToUintMap storage map) internal view returns (uint256) {
+    function length(AddressToUintMap storage map) internal view returns (uint256) {
         return length(map._inner);
     }
 
@@ -395,17 +395,17 @@ library EnumerableMap {
      *
      * - `index` must be strictly less than {length}.
      */
-    function at(UintToUintMap storage map, uint256 index) internal view returns (uint256, uint256) {
+    function at(AddressToUintMap storage map, uint256 index) internal view returns (address, uint256) {
         (bytes32 key, bytes32 value) = at(map._inner, index);
-        return (uint256(key), uint256(value));
+        return (address(uint160(uint256(key))), uint256(value));
     }
 
     /**
      * @dev Tries to returns the value associated with `key`.  O(1).
      * Does not revert if `key` is not in the map.
      */
-    function tryGet(UintToUintMap storage map, uint256 key) internal view returns (bool, uint256) {
-        (bool success, bytes32 value) = tryGet(map._inner, bytes32(key));
+    function tryGet(AddressToUintMap storage map, address key) internal view returns (bool, uint256) {
+        (bool success, bytes32 value) = tryGet(map._inner, bytes32(uint256(uint160(key))));
         return (success, uint256(value));
     }
 
@@ -416,8 +416,8 @@ library EnumerableMap {
      *
      * - `key` must be in the map.
      */
-    function get(UintToUintMap storage map, uint256 key) internal view returns (uint256) {
-        return uint256(get(map._inner, bytes32(key)));
+    function get(AddressToUintMap storage map, address key) internal view returns (uint256) {
+        return uint256(get(map._inner, bytes32(uint256(uint160(key)))));
     }
 
     /**
@@ -427,10 +427,10 @@ library EnumerableMap {
      * message unnecessarily. For custom revert reasons use {tryGet}.
      */
     function get(
-        UintToUintMap storage map,
-        uint256 key,
+        AddressToUintMap storage map,
+        address key,
         string memory errorMessage
     ) internal view returns (uint256) {
-        return uint256(get(map._inner, bytes32(key), errorMessage));
+        return uint256(get(map._inner, bytes32(uint256(uint160(key))), errorMessage));
     }
 }