From 403795166ada8ab8359dc2f0d6f39050ed87dea2 Mon Sep 17 00:00:00 2001 From: Micah Zoltu Date: Thu, 30 Jan 2020 22:22:57 +0800 Subject: [PATCH 1/3] Adds EnumerableMap. I'm not sure what the best way to recommend a library to OZ contracts is, so I figured I would start with a PR and see what happens. This is code is _strongly_ based on the OZ EnumerableSet. The biggest usability difference is that you have to instantiate manually and then call `.initialize` rather than doing `EnumerableMap.newMap()`. This is because struct arrays aren't supported in the ways necessary to make `EnumerableMap.newMap()` work. The `require` statements are basically invariants to protect against bugs (and provide an early warning during testing) but could be removed as they should never be hit (really they are assertions with useful error messages). I also changed the array to be 1-indexed, which differs from `EnumerableSet` which has a 0-indexed array and then treats the index values as all off-by-one. I personally found the code easier to understand with a placeholder value at the first entry position rather than having to shift the indexes by 1. Since the placeholder is just a 0 value, I suspect there is no meaningful gas difference between the strategies and it is just a matter of style. The biggest problem here, of course (same as `EnumerableSet`) is that Solidity doesn't contain generics, so the user will need to change the key and value types to match their code. The nice thing here though is that the key can be any valid mapping key and the value can be anything including structs. The user only needs to change the type in a few places. I chose to use `address` and `uint8` because you can do a find/replace as neither type is used outside the key/value types and it is as reasonable of a choice as any. --- contracts/utils/EnumerableMap.sol | 71 +++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 contracts/utils/EnumerableMap.sol diff --git a/contracts/utils/EnumerableMap.sol b/contracts/utils/EnumerableMap.sol new file mode 100644 index 00000000000..5a24f5bad65 --- /dev/null +++ b/contracts/utils/EnumerableMap.sol @@ -0,0 +1,71 @@ +library EnumerableMap { + struct Entry { + address key; + uint8 value; + } + + struct Map { + mapping (address => uint256) index; + Entry[] entries; + } + + // we initialize it with a placeholder item in the first position because we treat the array as 1-indexed since 0 is a special index (means no entry in the index) + function initialize(Map storage map) internal { + map.entries.push(); + } + + function contains(Map storage map, address key) internal view returns (bool) { + return map.index[key] != 0; + } + + function add(Map storage map, address key, uint8 value) internal { + uint256 index = map.index[key]; + if (index == 0) { + // create new entry + Entry memory entry = Entry({ key: key, value: value }); + uint256 newEntryIndex = map.entries.length; + map.entries.push(entry); + map.index[key] = newEntryIndex; + } else { + // update existing entry + map.entries[map.index[key]].value = value; + } + + require(map.entries[map.index[key]].key == key, "Key at inserted location does not match inserted key."); + require(map.entries[map.index[key]].value == value, "Value at inserted location does not match inserted value."); + } + + function remove(Map storage map, address key) internal { + // get the index into entries array that this entry lives at + uint256 index = map.index[key]; + + // if this key doesn't exist in the index, then we have nothing to do + if (index == 0) return; + + // if the entry we are removing isn't the last, overwrite it with the last entry + uint256 lastIndex = map.entries.length - 1; + if (index != lastIndex) { + Entry storage lastEntry = map.entries[lastIndex]; + map.entries[index] = lastEntry; + map.index[lastEntry.key] = index; + } + + // delete the last entry (if the item we are removing isn't last, it will have been overwritten inside the conditional above) + map.entries.pop(); + + // delete the index pointer + delete map.index[key]; + + require(map.index[key] == 0, "Removed key still exists in the index."); + require(map.entries[index].key != key, "Removed key still exists in the array at original index."); + } + + function enumerate(Map storage map) internal view returns (Entry[] memory) { + Entry[] memory output = new Entry[](map.entries.length - 1); + + for (uint256 i = 1; i < map.entries.length; ++i) { + output[i - 1] = map.entries[i]; + } + return output; + } +} From 3752ee6d92e6063a887a817f6f4cd376f3401032 Mon Sep 17 00:00:00 2001 From: Micah Zoltu Date: Thu, 30 Jan 2020 23:05:56 +0800 Subject: [PATCH 2/3] Adds getter. --- contracts/utils/EnumerableMap.sol | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/contracts/utils/EnumerableMap.sol b/contracts/utils/EnumerableMap.sol index 5a24f5bad65..91fd3d05bab 100644 --- a/contracts/utils/EnumerableMap.sol +++ b/contracts/utils/EnumerableMap.sol @@ -60,6 +60,10 @@ library EnumerableMap { require(map.entries[index].key != key, "Removed key still exists in the array at original index."); } + function get(Map storage map, address key) internal view returns (uint8) { + return map.entries[map.index[key]].value; + } + function enumerate(Map storage map) internal view returns (Entry[] memory) { Entry[] memory output = new Entry[](map.entries.length - 1); From 17e64fbd7c72395cd59b65f9340a34f8a1bdfe1c Mon Sep 17 00:00:00 2001 From: Micah Zoltu Date: Sat, 1 Feb 2020 09:56:51 +0800 Subject: [PATCH 3/3] Makes get throw instead of returning a sentinel value. Also added a `tryGet` for when you _want_ a sentinel value instead of throwing. Also allows someone who is copying this code to customize the sentinel value to something besides `0` if desired (e.g., INT_MAX). --- contracts/utils/EnumerableMap.sol | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/contracts/utils/EnumerableMap.sol b/contracts/utils/EnumerableMap.sol index 91fd3d05bab..b8214d1882c 100644 --- a/contracts/utils/EnumerableMap.sol +++ b/contracts/utils/EnumerableMap.sol @@ -61,7 +61,15 @@ library EnumerableMap { } function get(Map storage map, address key) internal view returns (uint8) { - return map.entries[map.index[key]].value; + uint256 index = map.index[key]; + require(index != 0, "Provided key was not in the map."); + return map.entries[index].value; + } + + function tryGet(Map storage map, address key) internal view returns (uint256) { + uint256 index = map.index[key]; + if (index == 0) return 0; + return map.entries[index].value; } function enumerate(Map storage map) internal view returns (Entry[] memory) {