fix: custom flags

yearn · Oct 24, 2023 · 537ce5c · 537ce5c
1 parent f80664a
commit 537ce5c
Show file tree

Hide file tree

Showing 6 changed files with 1,100 additions and 486 deletions.
diff --git a/contracts/accountants/GenericAccountant.vy b/contracts/accountants/GenericAccountant.vy
@@ -124,7 +124,7 @@ vaults: public(HashMap[address, bool])
 # Default config to use unless a custom one is set.
 defaultConfig: public(Fee)
 # Mapping vault => strategy => custom Fee config
-fees: public(HashMap[address, HashMap[address, Fee]])
+customConfig: public(HashMap[address, HashMap[address, Fee]])
 # Mapping vault => strategy => flag to use a custom config.
 custom: public(HashMap[address, HashMap[address, bool]])
 
@@ -186,7 +186,7 @@ def report(strategy: address, gain: uint256, loss: uint256) -> (uint256, uint256
 
     # Check if it there is a custom config to use.
     if self.custom[msg.sender][strategy]:
-        fee = self.fees[msg.sender][strategy]
+        fee = self.customConfig[msg.sender][strategy]
     else:
         # Otherwise use the default.
         fee = self.defaultConfig
@@ -328,7 +328,7 @@ def setCustomConfig(
     assert custom_performance <= PERFORMANCE_FEE_THRESHOLD, "exceeds performance fee threshold"
 
     # Set this strategies custom config.
-    self.fees[vault][strategy] = Fee({
+    self.customConfig[vault][strategy] = Fee({
         managementFee: custom_management,
         performanceFee: custom_performance,
         refundRatio: custom_refund,
@@ -338,7 +338,7 @@ def setCustomConfig(
     # Make sure flag is declared as True.
     self.custom[vault][strategy] = True
 
-    log UpdateCustomFeeConfig(vault, strategy, self.fees[vault][strategy])
+    log UpdateCustomFeeConfig(vault, strategy, self.customConfig[vault][strategy])
 
 
 @external
@@ -351,7 +351,7 @@ def removeCustomConfig(vault: address, strategy: address):
     assert self.custom[vault][strategy], "No custom fees set"
 
     # Set all the strategies custom fees to 0.
-    self.fees[vault][strategy] = empty(Fee)
+    self.customConfig[vault][strategy] = empty(Fee)
     # Turn off the flag.
     self.custom[vault][strategy] = False
 

diff --git a/contracts/accountants/HelathCheckAccountant.sol b/contracts/accountants/HelathCheckAccountant.sol
@@ -100,10 +100,10 @@ contract HealthCheckAccountant {
     Fee public defaultConfig;
 
     /// @notice Mapping vault => strategy => custom Fee config if any.
-    mapping(address => mapping(address => Fee)) public fees;
+    mapping(address => mapping(address => Fee)) public customConfig;
 
     /// @notice Mapping vault => strategy => flag to use a custom config.
-    mapping(address => mapping(address => bool)) public custom;
+    mapping(address => mapping(address => uint256)) internal _custom;
 
     constructor(
         address _feeManager,
@@ -165,8 +165,8 @@ contract HealthCheckAccountant {
         Fee memory fee;
 
         // Check if there is a custom config to use.
-        if (custom[msg.sender][strategy]) {
-            fee = fees[msg.sender][strategy];
+        if (_custom[msg.sender][strategy] != 0) {
+            fee = customConfig[msg.sender][strategy];
         } else {
             // Otherwise use the default.
             fee = defaultConfig;
@@ -337,7 +337,7 @@ contract HealthCheckAccountant {
         require(customMaxLoss <= MAX_BPS, "too high");
 
         // Set the strategy's custom config.
-        fees[vault][strategy] = Fee({
+        customConfig[vault][strategy] = Fee({
             managementFee: customManagement,
             performanceFee: customPerformance,
             refundRatio: customRefund,
@@ -347,9 +347,13 @@ contract HealthCheckAccountant {
         });
 
         // Set the custom flag.
-        custom[vault][strategy] = true;
+        _custom[vault][strategy] = 1;
 
-        emit UpdateCustomFeeConfig(vault, strategy, fees[vault][strategy]);
+        emit UpdateCustomFeeConfig(
+            vault,
+            strategy,
+            customConfig[vault][strategy]
+        );
     }
 
     /**
@@ -362,18 +366,34 @@ contract HealthCheckAccountant {
         address strategy
     ) external onlyFeeManager {
         // Ensure custom fees are set for the specified vault and strategy.
-        require(custom[vault][strategy], "No custom fees set");
+        require(_custom[vault][strategy] != 0, "No custom fees set");
 
         // Set all the strategy's custom fees to 0.
-        delete fees[vault][strategy];
+        delete customConfig[vault][strategy];
 
         // Clear the custom flag.
-        custom[vault][strategy] = false;
+        _custom[vault][strategy] = 0;
 
         // Emit relevant event.
         emit RemovedCustomFeeConfig(vault, strategy);
     }
 
+    /**
+     * @notice Public getter to check for custom setting.
+     * @dev We use uint256 for the flag since its cheaper so this
+     *   will convert it to a bool for easy view functions.
+     *
+     * @param vault Address of the vault.
+     * @param strategy Address of the strategy
+     * @return If a custom fee config is set.
+     */
+    function custom(
+        address vault,
+        address strategy
+    ) external view returns (bool) {
+        return _custom[vault][strategy] != 0;
+    }
+
     /**
      * @notice Function to withdraw the underlying asset from a vault.
      * @param vault The vault to withdraw from.