refactor allowlist mappings to use AllowlistType;

AngelProtocolFinance · Aug 15, 2023 · 7b3b4f2 · 7b3b4f2
1 parent 5b8c2b0
commit 7b3b4f2
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 45 deletions.
diff --git a/contracts/core/accounts/facets/AccountsAllowance.sol b/contracts/core/accounts/facets/AccountsAllowance.sol
@@ -61,7 +61,10 @@ contract AccountsAllowance is
         ),
         "Unauthorized"
       );
-      inAllowlist = (IterableMapping.get(state.allowlistedBeneficiaries[endowId], spender) == 1);
+      inAllowlist = (IterableMapping.get(
+        state.allowlists[endowId][LibAccounts.AllowlistType.AllowlistedBeneficiaries],
+        spender
+      ) == 1);
     } else {
       // Only the endowment owner or a delegate whom controls allowlist can update allowances
       require(
@@ -73,7 +76,10 @@ contract AccountsAllowance is
         ),
         "Unauthorized"
       );
-      inAllowlist = (IterableMapping.get(state.maturityAllowlist[endowId], spender) == 1);
+      inAllowlist = (IterableMapping.get(
+        state.allowlists[endowId][LibAccounts.AllowlistType.MaturityAllowlist],
+        spender
+      ) == 1);
     }
     require(inAllowlist, "Spender is not in allowlists");
 

diff --git a/contracts/core/accounts/facets/AccountsCreateEndowment.sol b/contracts/core/accounts/facets/AccountsCreateEndowment.sol
@@ -116,20 +116,24 @@ contract AccountsCreateEndowment is
     // process all initial allowlists provided by user into their respective mappings
     for (uint256 i = 0; i < details.allowlistedBeneficiaries.length; i++) {
       IterableMapping.set(
-        state.allowlistedBeneficiaries[newEndowId],
+        state.allowlists[newEndowId][LibAccounts.AllowlistType.AllowlistedBeneficiaries],
         details.allowlistedBeneficiaries[i],
         1
       );
     }
     for (uint256 i = 0; i < details.allowlistedContributors.length; i++) {
       IterableMapping.set(
-        state.allowlistedContributors[newEndowId],
+        state.allowlists[newEndowId][LibAccounts.AllowlistType.AllowlistedContributors],
         details.allowlistedContributors[i],
         1
       );
     }
     for (uint256 i = 0; i < details.maturityAllowlist.length; i++) {
-      IterableMapping.set(state.maturityAllowlist[newEndowId], details.maturityAllowlist[i], 1);
+      IterableMapping.set(
+        state.allowlists[newEndowId][LibAccounts.AllowlistType.MaturityAllowlist],
+        details.maturityAllowlist[i],
+        1
+      );
     }
 
     emit EndowmentCreated(newEndowId, details.endowType);

diff --git a/contracts/core/accounts/facets/AccountsDepositWithdrawEndowments.sol b/contracts/core/accounts/facets/AccountsDepositWithdrawEndowments.sol
@@ -136,10 +136,11 @@ contract AccountsDepositWithdrawEndowments is
     bool contributorAllowed = true;
     if (
       (msg.sender != registrarConfig.indexFundContract && msg.sender != tempEndowment.owner) &&
-      state.allowlistedContributors[details.id].keys.length > 0
+      state.allowlists[details.id][LibAccounts.AllowlistType.AllowlistedContributors].keys.length >
+      0
     ) {
       contributorAllowed = (IterableMapping.get(
-        state.allowlistedContributors[details.id],
+        state.allowlists[details.id][LibAccounts.AllowlistType.AllowlistedContributors],
         msg.sender
       ) == 1);
     }
@@ -348,19 +349,23 @@ contract AccountsDepositWithdrawEndowments is
           bool beneficiaryAllowed;
           if (mature) {
             beneficiaryAllowed = (IterableMapping.get(
-              state.maturityAllowlist[id],
+              state.allowlists[id][LibAccounts.AllowlistType.MaturityAllowlist],
               beneficiaryAddress
             ) == 1);
             require(beneficiaryAllowed, "Beneficiary address is not listed in maturityAllowlist");
           } else {
             if (
-              state.allowlistedBeneficiaries[id].keys.length == 0 ||
+              state
+                .allowlists[id][LibAccounts.AllowlistType.AllowlistedBeneficiaries]
+                .keys
+                .length ==
+              0 ||
               beneficiaryAddress == tempEndowment.owner
             ) {
               beneficiaryAllowed = true;
             } else {
               beneficiaryAllowed = (IterableMapping.get(
-                state.allowlistedBeneficiaries[id],
+                state.allowlists[id][LibAccounts.AllowlistType.AllowlistedBeneficiaries],
                 beneficiaryAddress
               ) == 1);
             }

diff --git a/contracts/core/accounts/facets/AccountsQueryEndowments.sol b/contracts/core/accounts/facets/AccountsQueryEndowments.sol
@@ -75,9 +75,11 @@ contract AccountsQueryEndowments is IAccountsQueryEndowments, IterableMapping {
         splitToLiquid: endowment.splitToLiquid,
         referralId: endowment.referralId,
         gasFwd: endowment.gasFwd,
-        allowlistedBeneficiaries: state.allowlistedBeneficiaries[id].keys,
-        allowlistedContributors: state.allowlistedContributors[id].keys,
-        maturityAllowlist: state.maturityAllowlist[id].keys
+        allowlistedBeneficiaries: state
+        .allowlists[id][LibAccounts.AllowlistType.AllowlistedBeneficiaries].keys,
+        allowlistedContributors: state
+        .allowlists[id][LibAccounts.AllowlistType.AllowlistedContributors].keys,
+        maturityAllowlist: state.allowlists[id][LibAccounts.AllowlistType.MaturityAllowlist].keys
       });
   }
 

diff --git a/contracts/core/accounts/facets/AccountsUpdateEndowmentSettingsController.sol b/contracts/core/accounts/facets/AccountsUpdateEndowmentSettingsController.sol
@@ -58,7 +58,6 @@ contract AccountsUpdateEndowmentSettingsController is
         ),
         "Unauthorized"
       );
-      iterableAllowlist = state.allowlistedBeneficiaries[id];
     } else if (allowlistType == LibAccounts.AllowlistType.AllowlistedContributors) {
       require(
         Validator.canChange(
@@ -69,7 +68,6 @@ contract AccountsUpdateEndowmentSettingsController is
         ),
         "Unauthorized"
       );
-      iterableAllowlist = state.allowlistedContributors[id];
     } else if (allowlistType == LibAccounts.AllowlistType.MaturityAllowlist) {
       require(
         Validator.canChange(
@@ -80,19 +78,18 @@ contract AccountsUpdateEndowmentSettingsController is
         ),
         "Unauthorized"
       );
-      iterableAllowlist = state.maturityAllowlist[id];
     } else {
       revert("Invalid AllowlistType");
     }
 
     for (uint256 i = 0; i < add.length; i++) {
       require(add[i] != address(0), "Zero address passed");
-      IterableMapping.set(iterableAllowlist, add[i], 1);
+      IterableMapping.set(state.allowlists[id][allowlistType], add[i], 1);
     }
 
     for (uint256 i = 0; i < remove.length; i++) {
       require(add[i] != address(0), "Zero address passed");
-      IterableMapping.remove(iterableAllowlist, remove[i]);
+      IterableMapping.remove(state.allowlists[id][allowlistType], remove[i]);
     }
 
     emit EndowmentAllowlistUpdated(id, allowlistType, add, remove);

diff --git a/contracts/core/accounts/storage.sol b/contracts/core/accounts/storage.sol
@@ -68,10 +68,8 @@ library AccountStorage {
     mapping(uint32 => mapping(address => address)) PriceFeeds;
     // Endowments that a DAF can withdraw to, managed by contract Owner
     mapping(uint32 => bool) dafApprovedEndowments;
-    // Various Endowment linked AllowList mappings
-    mapping(uint32 => IterableMapping.Map) allowlistedContributors;
-    mapping(uint32 => IterableMapping.Map) allowlistedBeneficiaries;
-    mapping(uint32 => IterableMapping.Map) maturityAllowlist;
+    // Endowments AllowLists Iterable mappings
+    mapping(uint32 => mapping(LibAccounts.AllowlistType => IterableMapping.Map)) allowlists;
     Config config;
   }
 }

diff --git a/contracts/test/accounts/TestFacetProxyContract.sol b/contracts/test/accounts/TestFacetProxyContract.sol
@@ -204,20 +204,8 @@ contract TestFacetProxyContract is TransparentUpgradeableProxy, IterableMapping
     address[] memory allowlist
   ) external {
     AccountStorage.State storage state = LibAccounts.diamondStorage();
-    if (listType == LibAccounts.AllowlistType.AllowlistedBeneficiaries) {
-      for (uint256 i = 0; i < allowlist.length; i++) {
-        IterableMapping.set(state.allowlistedBeneficiaries[endowId], allowlist[i], 1);
-      }
-    } else if (listType == LibAccounts.AllowlistType.AllowlistedContributors) {
-      for (uint256 i = 0; i < allowlist.length; i++) {
-        IterableMapping.set(state.allowlistedContributors[endowId], allowlist[i], 1);
-      }
-    } else if (listType == LibAccounts.AllowlistType.MaturityAllowlist) {
-      for (uint256 i = 0; i < allowlist.length; i++) {
-        IterableMapping.set(state.maturityAllowlist[endowId], allowlist[i], 1);
-      }
-    } else {
-      revert("Invlaid AllowlistType");
+    for (uint256 i = 0; i < allowlist.length; i++) {
+      IterableMapping.set(state.allowlists[endowId][listType], allowlist[i], 1);
     }
   }
 
@@ -226,15 +214,7 @@ contract TestFacetProxyContract is TransparentUpgradeableProxy, IterableMapping
     LibAccounts.AllowlistType listType
   ) external view returns (address[] memory) {
     AccountStorage.State storage state = LibAccounts.diamondStorage();
-    if (listType == LibAccounts.AllowlistType.AllowlistedBeneficiaries) {
-      return state.allowlistedBeneficiaries[endowId].keys;
-    } else if (listType == LibAccounts.AllowlistType.AllowlistedContributors) {
-      return state.allowlistedContributors[endowId].keys;
-    } else if (listType == LibAccounts.AllowlistType.MaturityAllowlist) {
-      return state.maturityAllowlist[endowId].keys;
-    } else {
-      revert("Invlaid AllowlistType");
-    }
+    return state.allowlists[endowId][listType].keys;
   }
 
   function callSelf(uint256 value, bytes memory data) external {