add extrema cases

IndexCoop · Aug 30, 2024 · 4269d1e · 4269d1e
1 parent 01216fe
commit 4269d1e
Show file tree

Hide file tree

Showing 2 changed files with 355 additions and 34 deletions.
diff --git a/contracts/adapters/TargetWeightWrapExtension.sol b/contracts/adapters/TargetWeightWrapExtension.sol
@@ -139,7 +139,6 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
     {
         require(isRebalancingActive, "Rebalancing is not active");
         require(rebalanceInfo.targetAssets.contains(_targetAsset), "Invalid target asset");
-        require(isReserveOverweight(), "Reserve asset is not overweight");
 
         bytes memory data = abi.encodeWithSelector(
             wrapModule.wrap.selector,
@@ -153,8 +152,8 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
         invokeManager(address(wrapModule), data);
 
         (uint256 targetAssetWeight, uint256 reserveWeight) = getTargetAssetAndReserveWeight(_targetAsset);
-        require(targetAssetWeight < executionParams[_targetAsset].maxTargetWeight, "Target asset overweight post-wrap");
-        require(reserveWeight > rebalanceInfo.minReserveWeight, "Reserve asset underweight post-wrap");
+        require(targetAssetWeight <= executionParams[_targetAsset].maxTargetWeight, "Target asset overweight post-wrap");
+        require(reserveWeight >= rebalanceInfo.minReserveWeight, "Reserve asset underweight post-wrap");
     }
 
     /**
@@ -188,8 +187,8 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
         invokeManager(address(wrapModule), data);
 
         (uint256 targetAssetWeight, uint256 reserveWeight) = getTargetAssetAndReserveWeight(_targetAsset);
-        require(targetAssetWeight > executionParams[_targetAsset].minTargetWeight, "Target asset underweight post-unwrap");
-        require(reserveWeight < rebalanceInfo.maxReserveWeight, "Reserve asset overweight post-unwrap");
+        require(targetAssetWeight >= executionParams[_targetAsset].minTargetWeight, "Target asset underweight post-unwrap");
+        require(reserveWeight <= rebalanceInfo.maxReserveWeight, "Reserve asset overweight post-unwrap");
     }
 
     /* ========== Operator Functions ========== */
@@ -271,7 +270,9 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
      * @return reserveValuation The valuation of the reserve asset.
      */
     function getReserveValuation() public view returns(uint256 reserveValuation) {
-        reserveValuation = setValuer.calculateComponentValuation(setToken, rebalanceInfo.reserveAsset, rebalanceInfo.reserveAsset);
+        reserveValuation = setToken.isComponent(rebalanceInfo.reserveAsset)
+            ? setValuer.calculateComponentValuation(setToken, rebalanceInfo.reserveAsset, rebalanceInfo.reserveAsset)
+            : 0;
     }
 
     /**
@@ -280,7 +281,9 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
      * @return targetAssetValuation The valuation of the specified target asset.
      */
     function getTargetAssetValuation(address _targetAsset) public view returns(uint256 targetAssetValuation) {
-        targetAssetValuation = setValuer.calculateComponentValuation(setToken, _targetAsset, rebalanceInfo.reserveAsset);
+        targetAssetValuation = setToken.isComponent(_targetAsset)
+            ? setValuer.calculateComponentValuation(setToken, _targetAsset, rebalanceInfo.reserveAsset)
+            : 0;
     }
 
     /**
@@ -303,6 +306,33 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
     }
 
     /**
+     * @notice Gets the weight of a specific target asset relative to the total valuation of the SetToken.
+     * @dev The weight is returned as a percentage where 100% equals 1e18.
+     * @param _targetAsset The address of the target asset.
+     * @return targetAssetWeight The weight of the specified target asset relative to the SetToken's total valuation.
+     */
+    function getTargetAssetWeight(address _targetAsset) public view returns(uint256 targetAssetWeight) {
+        uint256 targetAssetValuation = getTargetAssetValuation(_targetAsset);
+        uint256 totalValuation = getTotalValuation();
+        targetAssetWeight = targetAssetValuation.preciseDiv(totalValuation);
+    }
+
+    /**
+     * @notice Gets the weights of both the target asset and the reserve asset relative to the total valuation of the SetToken.
+     * @dev The weights are returned as percentages where 100% equals 1e18.
+     * @param _targetAsset The address of the target asset.
+     * @return targetAssetWeight The weight of the target asset relative to the SetToken's total valuation.
+     * @return reserveWeight The weight of the reserve asset relative to the SetToken's total valuation.
+     */
+    function getTargetAssetAndReserveWeight(address _targetAsset) public view returns(uint256 targetAssetWeight, uint256 reserveWeight) {
+        uint256 targetAssetValuation = getTargetAssetValuation(_targetAsset);
+        uint256 reserveValuation = getReserveValuation();
+        uint256 totalValuation = getTotalValuation();
+        targetAssetWeight = targetAssetValuation.preciseDiv(totalValuation);
+        reserveWeight = reserveValuation.preciseDiv(totalValuation);
+    }
+
+        /**
      * @notice Checks if the reserve asset is overweight.
      */
     function isReserveOverweight() public view returns(bool) {
@@ -332,33 +362,6 @@ contract TargetWeightWrapExtension is BaseExtension, ReentrancyGuard {
         return getTargetAssetWeight(_targetAsset) < executionParams[_targetAsset].minTargetWeight;
     }
 
-    /**
-     * @notice Gets the weight of a specific target asset relative to the total valuation of the SetToken.
-     * @dev The weight is returned as a percentage where 100% equals 1e18.
-     * @param _targetAsset The address of the target asset.
-     * @return targetAssetWeight The weight of the specified target asset relative to the SetToken's total valuation.
-     */
-    function getTargetAssetWeight(address _targetAsset) public view returns(uint256 targetAssetWeight) {
-        uint256 targetAssetValuation = getTargetAssetValuation(_targetAsset);
-        uint256 totalValuation = getTotalValuation();
-        targetAssetWeight = targetAssetValuation.preciseDiv(totalValuation);
-    }
-
-    /**
-     * @notice Gets the weights of both the target asset and the reserve asset relative to the total valuation of the SetToken.
-     * @dev The weights are returned as percentages where 100% equals 1e18.
-     * @param _targetAsset The address of the target asset.
-     * @return targetAssetWeight The weight of the target asset relative to the SetToken's total valuation.
-     * @return reserveWeight The weight of the reserve asset relative to the SetToken's total valuation.
-     */
-    function getTargetAssetAndReserveWeight(address _targetAsset) public view returns(uint256 targetAssetWeight, uint256 reserveWeight) {
-        uint256 targetAssetValuation = getTargetAssetValuation(_targetAsset);
-        uint256 reserveValuation = getReserveValuation();
-        uint256 totalValuation = getTotalValuation();
-        targetAssetWeight = targetAssetValuation.preciseDiv(totalValuation);
-        reserveWeight = reserveValuation.preciseDiv(totalValuation);
-    }
-
     /**
      * @notice Gets the list of target assets that can be wrapped into or unwrapped from during rebalancing.
      * @return An array of addresses representing the target assets.