fix IterableMapping lib > contract & include it; Function to cast key…

…s uint256[] >> uint32[] for endowment IDs;
AngelProtocolFinance · Jul 28, 2023 · ae0dcc7 · ae0dcc7
1 parent e1f281a
commit ae0dcc7
Show file tree

Hide file tree

Showing 4 changed files with 82 additions and 123 deletions.
diff --git a/contracts/core/index-fund/IndexFund.sol b/contracts/core/index-fund/IndexFund.sol
@@ -11,7 +11,7 @@ import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import {IIndexFund} from "./IIndexFund.sol";
 import {Array, Array32} from "../../lib/array.sol";
-import {IterableMappingEndow, IterableMappingFund} from "../../lib/IterableMappings.sol";
+import {IterableMapping} from "../../lib/IterableMapping.sol";
 import {Utils} from "../../lib/utils.sol";
 import {IRegistrar} from "../registrar/interfaces/IRegistrar.sol";
 import {RegistrarStorage} from "../registrar/storage.sol";
@@ -29,11 +29,9 @@ uint256 constant MIN_AMOUNT_PER_ENDOWMENT = 100;
  * It is responsible for creating new funds, adding endowments to funds, and
  * distributing funds to the endowment members
  */
-contract IndexFund is IIndexFund, Storage, OwnableUpgradeable, ReentrancyGuard {
+contract IndexFund is IIndexFund, Storage, OwnableUpgradeable, ReentrancyGuard, IterableMapping {
   using SafeERC20 for IERC20;
   using SafeMath for uint256;
-  using IterableMappingEndow for IterableMappingEndow.Map;
-  using IterableMappingFund for IterableMappingFund.Map;
 
   /**
    * @notice Initializer function for index fund contract, to be called when proxy is deployed
@@ -137,8 +135,8 @@ contract IndexFund is IIndexFund, Storage, OwnableUpgradeable, ReentrancyGuard {
 
     // set all Fund <> Endowment mappings
     for (uint8 i = 0; i < endowments.length; i++) {
-      state.FundsByEndowment[endowments[i]].set(state.nextFundId, true);
-      state.EndowmentsByFund[state.nextFundId].set(endowments[i], true);
+      IterableMapping.set(state.FundsByEndowment[endowments[i]], state.nextFundId, true);
+      IterableMapping.set(state.EndowmentsByFund[state.nextFundId], endowments[i], true);
     }
 
     if (rotatingFund) {
@@ -180,8 +178,8 @@ contract IndexFund is IIndexFund, Storage, OwnableUpgradeable, ReentrancyGuard {
 
     // remove endowment from all involved funds if in their endowments array
     for (uint i = 0; i < state.FundsByEndowment[endowment].keys.length; i++) {
-      uint256 fundId = state.FundsByEndowment[endowment].getKeyAtIndex(i);
-      state.EndowmentsByFund[fundId].remove(endowment);
+      uint256 fundId = IterableMapping.getKeyAtIndex(state.FundsByEndowment[endowment], i);
+      IterableMapping.remove(state.EndowmentsByFund[fundId], endowment);
       emit MemberRemoved(fundId, endowment);
       // if endowment removal results in a fund having zero endowment members left, close out the fund
       if (state.EndowmentsByFund[fundId].keys.length == 0) {
@@ -212,14 +210,14 @@ contract IndexFund is IIndexFund, Storage, OwnableUpgradeable, ReentrancyGuard {
 
     // add Endowments passed to a Fund members and FundsByEndowment mappings
     for (uint32 i = 0; i < endowmentsAdd.length; i++) {
-      state.EndowmentsByFund[fundId].set(endowmentsAdd[i], true);
-      state.FundsByEndowment[endowmentsAdd[i]].set(fundId, true);
+      IterableMapping.set(state.FundsByEndowment[endowmentsAdd[i]], fundId, true);
+      IterableMapping.set(state.EndowmentsByFund[fundId], endowmentsAdd[i], true);
     }
 
     // Endowments to be removed from a Fund
     for (uint32 i = 0; i < endowmentsRemove.length; i++) {
-      state.EndowmentsByFund[fundId].remove(endowmentsRemove[i]);
-      state.FundsByEndowment[endowmentsRemove[i]].remove(fundId);
+      IterableMapping.remove(state.EndowmentsByFund[fundId], endowmentsRemove[i]);
+      IterableMapping.remove(state.FundsByEndowment[endowmentsRemove[i]], fundId);
     }
 
     // resulting fund has no members, remove it
@@ -232,7 +230,7 @@ contract IndexFund is IIndexFund, Storage, OwnableUpgradeable, ReentrancyGuard {
       state.EndowmentsByFund[fundId].keys.length <= MAX_ENDOWMENT_MEMBERS,
       "Fund endowment members exceeds upper limit"
     );
-    emit MembersUpdated(fundId, state.EndowmentsByFund[fundId].keys);
+    emit MembersUpdated(fundId, IterableMapping.keysAsUint32(state.EndowmentsByFund[fundId]));
   }
 
   /**
@@ -409,7 +407,7 @@ contract IndexFund is IIndexFund, Storage, OwnableUpgradeable, ReentrancyGuard {
         id: state.Funds[fundId].id,
         name: state.Funds[fundId].name,
         description: state.Funds[fundId].description,
-        endowments: state.EndowmentsByFund[fundId].keys,
+        endowments: IterableMapping.keysAsUint32(state.EndowmentsByFund[fundId]),
         splitToLiquid: state.Funds[fundId].splitToLiquid,
         expiryTime: state.Funds[fundId].expiryTime
       });
@@ -436,7 +434,7 @@ contract IndexFund is IIndexFund, Storage, OwnableUpgradeable, ReentrancyGuard {
         id: state.Funds[state.activeFund].id,
         name: state.Funds[state.activeFund].name,
         description: state.Funds[state.activeFund].description,
-        endowments: state.EndowmentsByFund[state.activeFund].keys,
+        endowments: IterableMapping.keysAsUint32(state.EndowmentsByFund[state.activeFund]),
         splitToLiquid: state.Funds[state.activeFund].splitToLiquid,
         expiryTime: state.Funds[state.activeFund].expiryTime
       });
@@ -492,7 +490,7 @@ contract IndexFund is IIndexFund, Storage, OwnableUpgradeable, ReentrancyGuard {
     for (uint256 i = 0; i < state.EndowmentsByFund[fundId].keys.length; i++) {
       IAccounts(accountsContract).depositERC20(
         AccountMessages.DepositRequest({
-          id: state.EndowmentsByFund[fundId].getKeyAtIndex(i),
+          id: uint32(IterableMapping.getKeyAtIndex(state.EndowmentsByFund[fundId], i)),
           lockedPercentage: 100 - liquidSplit,
           liquidPercentage: liquidSplit
         }),

diff --git a/contracts/core/index-fund/storage.sol b/contracts/core/index-fund/storage.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.16;
 
-import {IterableMappingEndow, IterableMappingFund} from "../../lib/IterableMappings.sol";
+import {IterableMapping} from "../../lib/IterableMapping.sol";
 import {IIndexFund} from "./IIndexFund.sol";
 
 library IndexFundStorage {
@@ -32,9 +32,9 @@ library IndexFundStorage {
     // Fund ID >> Fund
     mapping(uint256 => Fund) Funds;
     // Fund ID >> Mapping (Endow ID >> bool)
-    mapping(uint256 => IterableMappingEndow.Map) EndowmentsByFund;
+    mapping(uint256 => IterableMapping.Map) EndowmentsByFund;
     // Endow ID >> Mapping (Fund ID >> bool)
-    mapping(uint32 => IterableMappingFund.Map) FundsByEndowment;
+    mapping(uint32 => IterableMapping.Map) FundsByEndowment;
   }
 }
 

diff --git a/contracts/lib/IterableMapping.sol b/contracts/lib/IterableMapping.sol
@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.16;
+
+contract IterableMapping {
+  struct Map {
+    uint256[] keys;
+    mapping(uint256 => bool) values;
+    mapping(uint256 => uint) indexOf;
+    mapping(uint256 => bool) inserted;
+  }
+
+  function get(Map storage map, uint256 key) internal view returns (bool) {
+    return map.values[key];
+  }
+
+  function getKeyAtIndex(Map storage map, uint index) internal view returns (uint256) {
+    return map.keys[index];
+  }
+
+  function size(Map storage map) internal view returns (uint) {
+    return map.keys.length;
+  }
+
+  function set(Map storage map, uint256 key, bool val) internal {
+    if (map.inserted[key]) {
+      map.values[key] = val;
+    } else {
+      map.inserted[key] = true;
+      map.values[key] = val;
+      map.indexOf[key] = map.keys.length;
+      map.keys.push(key);
+    }
+  }
+
+  function remove(Map storage map, uint256 key) internal {
+    if (!map.inserted[key]) {
+      return;
+    }
+
+    delete map.inserted[key];
+    delete map.values[key];
+
+    uint256 index = map.indexOf[key];
+    uint256 lastKey = map.keys[map.keys.length - 1];
+
+    map.indexOf[lastKey] = index;
+    delete map.indexOf[key];
+
+    map.keys[index] = lastKey;
+    map.keys.pop();
+  }
+
+  /**
+   * @dev Converts a Map's keys from a Uint256 Array to Uint32 Array
+   * @param map Map
+   * @return keys32 Map's keys as a Uint32 Array
+   */
+  function keysAsUint32(Map storage map) internal view returns (uint32[] memory) {
+    uint32[] memory keys32 = new uint32[](map.keys.length);
+    for (uint256 i = 0; i < map.keys.length; i++) {
+      keys32[i] = uint32(map.keys[i]);
+    }
+    return keys32;
+  }
+}
diff --git a/contracts/lib/IterableMappings.sol b/contracts/lib/IterableMappings.sol