Sov 3680 rm optimization

contracts/helpers/EnumerableAddressSet.sol

@@ -0,0 +1,139 @@
+pragma solidity ^0.5.0;
+
+/**
+ * @dev Based on Library for managing
+ * https://en.wikipedia.org/wiki/Set_(abstract_data_type)[sets] of primitive
+ * types.
+ *
+ * Sets have the following properties:
+ *
+ * - Elements are added, removed, and checked for existence in constant time
+ * (O(1)).
+ * - Elements are enumerated in O(n). No guarantees are made on the ordering.
+ *
+ * As of v2.5.0, only `address` sets are supported.
+ *
+ * Include with `using EnumerableSet for EnumerableSet.AddressSet;`.
+ *
+ * _Available since v2.5.0._
+ */
+library EnumerableAddressSet {
+    struct AddressSet {
+        // Position of the value in the `values` array, plus 1 because index 0
+        // means a value is not in the set.
+        mapping(address => uint256) index;
+        address[] values;
+    }
+
+    /**
+     * @dev Add a value to a set. O(1).
+     * Returns false if the value was already in the set.
+     */
+    function add(AddressSet storage set, address value) internal returns (bool) {
+        if (!contains(set, value)) {
+            set.index[value] = set.values.push(value);
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    /**
+     * @dev Removes a value from a set. O(1).
+     * Returns false if the value was not present in the set.
+     */
+    function remove(AddressSet storage set, address value) internal returns (bool) {
+        if (contains(set, value)) {
+            uint256 toDeleteIndex = set.index[value] - 1;
+            uint256 lastIndex = set.values.length - 1;
+
+            // If the element we're deleting is the last one, we can just remove it without doing a swap
+            if (lastIndex != toDeleteIndex) {
+                address lastValue = set.values[lastIndex];
+
+                // Move the last value to the index where the deleted value is
+                set.values[toDeleteIndex] = lastValue;
+                // Update the index for the moved value
+                set.index[lastValue] = toDeleteIndex + 1; // All indexes are 1-based
+            }
+
+            // Delete the index entry for the deleted value
+            delete set.index[value];
+
+            // Delete the old entry for the moved value
+            set.values.pop();
+
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    /**
+     * @dev Returns true if the value is in the set. O(1).
+     */
+    function contains(AddressSet storage set, address value) internal view returns (bool) {
+        return set.index[value] != 0;
+    }
+
+    /**
+     * @dev Returns an array with all values in the set. O(N).
+     * 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.
+
+     * WARNING: This function may run out of gas on large sets: use {length} and
+     * {get} instead in these cases.
+     */
+    function enumerate(AddressSet storage set) internal view returns (address[] memory) {
+        return set.values;
+    }
+
+    /**
+     * @dev Returns a chunk of array as recommended in enumerate() to avoid running of gas.
+     * 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.
+
+     * WARNING: This function may run out of gas on large sets: use {length} and
+     * {get} instead in these cases.
+
+     * @param start start index of chunk
+     * @param count num of element to return; if count == 0 then returns all the elements from the @param start
+     */
+    function enumerateChunk(
+        AddressSet storage set,
+        uint256 start,
+        uint256 count
+    ) internal view returns (address[] memory output) {
+        uint256 end = start + count;
+        require(end >= start, "addition overflow");
+        end = (set.values.length < end || count == 0) ? set.values.length : end;
+        if (end == 0 || start >= end) {
+            return output;
+        }
+
+        output = new address[](end - start);
+        for (uint256 i; i < end - start; i++) {
+            output[i] = set.values[i + start];
+        }
+        return output;
+    }
+
+    /**
+     * @dev Returns the number of elements on the set. O(1).
+     */
+    function length(AddressSet storage set) internal view returns (uint256) {
+        return set.values.length;
+    }
+
+    /** @dev Returns the element stored at position `index` in the set. 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.
+     *
+     * Requirements:
+     *
+     * - `index` must be strictly less than {length}.
+     */
+    function get(AddressSet storage set, uint256 index) internal view returns (address) {
+        return set.values[index];
+    }
+}

contracts/masset/RewardManager.sol

@@ -7,11 +7,12 @@ import { IERC20 } from "../openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { Address } from "../openzeppelin/contracts/utils/Address.sol";
 import { IRewardManager } from "./IRewardManager.sol";
 import { IMasset } from "./IMasset.sol";
-import { MappingAddressToUint256 } from "../helpers/MappingAddressToUint256.sol";
+import { EnumerableAddressSet } from "../helpers/EnumerableAddressSet.sol";
 
 /// @title A contract to manage rewards and penalties
 /// @author Derek Matter
 contract RewardManager is IRewardManager, Ownable {
+    using EnumerableAddressSet for EnumerableAddressSet.AddressSet; // enumerable map of addresses
 
     using SafeMath for uint256;
     uint256 constant ONE = 1000000000000000000;
@@ -21,7 +22,8 @@ contract RewardManager is IRewardManager, Ownable {
 
     IMasset private masset;
     uint256 private factor = 0;
-    MappingAddressToUint256 private targetWeights = new MappingAddressToUint256();
+    EnumerableAddressSet.AddressSet private tokensSet;
+    mapping(address => uint256) private targetWeights;
     uint256 globalMaxPenaltyPerc = 0;
     uint256 globalMaxRewardPerc = 0;
 
@@ -73,19 +75,24 @@ contract RewardManager is IRewardManager, Ownable {
     /// @notice Get the list of tokens
     /// @return tokens the tokens
     function getTokens() public view returns (address[] memory) {
-        return targetWeights.getKeys();
+        return tokensSet.enumerate();
     }
 
     /// @notice Get the target weight for a token
     /// @return targetWeight the target weight
     function getTargetWeight(address _tokenAddress) public view returns (uint256) {
-        return targetWeights.get(_tokenAddress);
+        return targetWeights[_tokenAddress];
     }
 
     /// @notice Get the target weights for all tokens
     /// @return targetWeights the target weights
     function getTargetWeights() public view returns (uint256[] memory) {
-        return targetWeights.getValues();
+        address[] memory keys = getTokens();
+        uint256[] memory r = new uint256[](keys.length);
+        for(uint i=0; i<r.length; i++) {
+            r[i] = getTargetWeight(keys[i]);
+        }
+        return r;
     }
 
     /// @notice Get the masset contract
@@ -139,13 +146,15 @@ contract RewardManager is IRewardManager, Ownable {
         uint256[] memory _targetWeights) public onlyOwner {
 
         require(_tokenAddresses.length == _targetWeights.length, "arrays not same length");
-        targetWeights = new MappingAddressToUint256();
+
+        delete tokensSet;
         uint256 total = 0;
 
         for(uint i=0; i<_tokenAddresses.length; i++) {
             require(Address.isContract(_tokenAddresses[i]), "token address not a contract");
             total = total.add(_targetWeights[i]);
-            targetWeights.set(_tokenAddresses[i], _targetWeights[i]);
+            tokensSet.add(_tokenAddresses[i]);
+            targetWeights[_tokenAddresses[i]] = _targetWeights[i];
             emit onTargetWeightChanged(msg.sender, _tokenAddresses[i], _targetWeights[i]);
         }
         require(total == ONE, "total not one");
@@ -198,9 +207,9 @@ contract RewardManager is IRewardManager, Ownable {
     }
 
     function getTotalBalanceInMasset() public view returns(uint256 total) {
-        address[] memory bassets = targetWeights.getKeys();
-        for(uint i = 0; i < bassets.length; i++) {
-            total = total.add(getBalanceInMasset(bassets[i]));
+        address[] memory tokens = getTokens();
+        for(uint i = 0; i < tokens.length; i++) {
+            total = total.add(getBalanceInMasset(tokens[i]));
         }
     } 
 
@@ -216,12 +225,12 @@ contract RewardManager is IRewardManager, Ownable {
 
         uint256 totalBefore = getTotalBalanceInMasset();
         uint256 totalAfter = _addOrSub(totalBefore, _sum);
-        address[] memory tokens = targetWeights.getKeys();
+        address[] memory tokens = getTokens();
 
         for(uint i=0; i<tokens.length; i++) {
             address tokenAddress = tokens[i];
             uint256 balanceBefore = getBalanceInMasset(tokenAddress);
-            uint256 targetWeight = targetWeights.get(tokenAddress);
+            uint256 targetWeight = targetWeights[tokenAddress];
             uint256 weightBefore = getWeight(balanceBefore, totalBefore);
             uint256 balanceAfter = tokenAddress == _basset ? _addOrSub(balanceBefore, _sum) : balanceBefore;
             uint256 weightAfter = getWeight(balanceAfter, totalAfter);
@@ -245,7 +254,7 @@ contract RewardManager is IRewardManager, Ownable {
 
         uint256 totalAfter = getTotalBalanceInMasset();
         uint256 totalBefore = totalAfter.sub(_sum);
-        address[] memory tokens = targetWeights.getKeys();
+        address[] memory tokens = getTokens();
 
         for(uint i=0; i<tokens.length; i++) {
             address tokenAddress = tokens[i];
@@ -254,7 +263,7 @@ contract RewardManager is IRewardManager, Ownable {
             if(tokenAddress == _basset) {
                 balanceBefore = balanceBefore.sub(_sum);
             }
-            uint256 targetWeight = targetWeights.get(tokenAddress);
+            uint256 targetWeight = targetWeights[tokenAddress];
             uint256 weightBefore = getWeight(balanceBefore, totalBefore);
             uint256 weightAfter = getWeight(balanceAfter, totalAfter);
             dsqrBefore = dsqrBefore.add(_dsqr(weightBefore, targetWeight));
@@ -274,7 +283,7 @@ contract RewardManager is IRewardManager, Ownable {
             uint256 _sum,
             bool _bridgeMode
         ) public view returns (bool) {
-        uint256 targetWeight = targetWeights.get(_bassetAddress);
+        uint256 targetWeight = targetWeights[_bassetAddress];
         uint256 balance = getBalanceInMasset(_bassetAddress);
         uint256 totalBalance = getTotalBalanceInMasset();
         if(!_bridgeMode) {
@@ -299,7 +308,7 @@ contract RewardManager is IRewardManager, Ownable {
         require(_sum < 1000000000000 * ONE, "_sum is too big"); // prevent overflow when casting to signed
 
         // uninitialized! avoid division by zero
-        require(targetWeights.getKeys().length > 0, "no target weights");
+        require(getTokens().length > 0, "no target weights");
 
         if(!isDepositDeservesReward(_bassetAddress, _sum, _bridgeMode)) {
             return 0;
@@ -340,7 +349,7 @@ contract RewardManager is IRewardManager, Ownable {
         require(_sum < 1000000000000 * ONE, "_sum is too big"); // prevent overflow when casting to signed
 
         // uninitialized! avoid division by zero
-        require(targetWeights.getKeys().length > 0, "no target weights");
+        require(getTokens().length > 0, "no target weights");
 
         (uint256 dsqrBefore, uint256 dsqrAfter) = getAverageDsqrs(_bassetAddress, 0 - int256(_sum));