Add Aave rebalance function

EridianAlpha · Apr 23, 2024 · e5ff4c8 · e5ff4c8
1 parent 90d7bf0
commit e5ff4c8
Show file tree

Hide file tree

Showing 3 changed files with 122 additions and 6 deletions.
diff --git a/src/AavePM.sol b/src/AavePM.sol
@@ -239,7 +239,7 @@ contract AavePM is IAavePM, Initializable, AccessControlUpgradeable, UUPSUpgrade
 
     /// @notice Borrow USDC from Aave.
     /// @dev Caller must have `MANAGER_ROLE`.
-    /// @param borrowAmount The amount of USDC to borrow. 8 decimal places to the cent.
+    /// @param borrowAmount The amount of USDC to borrow. 8 decimal places to the dollar. e.g. 100000000 = $1.00.
     function aaveBorrowUSDC(uint256 borrowAmount) public onlyRole(MANAGER_ROLE) {
         IPool(s_contractAddresses["aavePool"]).borrow(s_tokenAddresses["USDC"], borrowAmount, 2, 0, address(this));
     }
@@ -295,7 +295,7 @@ contract AavePM is IAavePM, Initializable, AccessControlUpgradeable, UUPSUpgrade
             sqrtPriceLimitX96: 0 // TODO: Calculate price limit
         });
 
-        // Approve the swapRouter to spend the tokenIn and swap the tokens
+        // Approve the swapRouter to spend the tokenIn and swap the tokens.
         TransferHelper.safeApprove(s_tokenAddresses[_tokenInIdentifier], address(swapRouter), currentBalance);
         amountOut = swapRouter.exactInputSingle(params);
         return (_tokenOutIdentifier, amountOut);
@@ -322,10 +322,10 @@ contract AavePM is IAavePM, Initializable, AccessControlUpgradeable, UUPSUpgrade
                 : _tokenOutIdentifier]
         ).decimals();
 
-        // Fetch current ratio from the pool
+        // Fetch current ratio from the pool.
         (uint160 sqrtRatioX96,,,,,,) = pool.slot0();
 
-        // Calculate the current ratio
+        // Calculate the current ratio.
         uint256 currentRatio = uint256(sqrtRatioX96) * (uint256(sqrtRatioX96)) * (10 ** _token0Decimals) >> (96 * 2);
 
         uint256 expectedOut = (_currentBalance * (10 ** _token0Decimals)) / currentRatio;
@@ -344,6 +344,74 @@ contract AavePM is IAavePM, Initializable, AccessControlUpgradeable, UUPSUpgrade
             : tokenIdentifier;
     }
 
+    // ================================================================
+    // │                    FUNCTIONS - CORE FEATURES                 │
+    // ================================================================
+    function rebalance() public onlyRole(MANAGER_ROLE) {
+        // Convert any ETH to WETH.
+        if (getContractBalance("ETH") > 0) wrapETHToWETH();
+
+        // Convert any WETH to wstETH.
+        if (getContractBalance("WETH") > 0) swapTokens("wstETH/ETH", "ETH", "wstETH");
+
+        // Deposit wstETH into Aave.
+        if (getContractBalance("wstETH") > 0) aaveSupplyWstETH();
+
+        // Check the current health factor
+        uint16 healthFactorTarget = getHealthFactorTarget();
+
+        // Get the current Aave account data
+        (
+            uint256 totalCollateralBase,
+            uint256 totalDebtBase,
+            ,
+            uint256 currentLiquidationThreshold,
+            ,
+            uint256 healthFactor
+        ) = getAaveAccountData();
+
+        // TODO: healthFactor and healthFactorTarget have different decimal places, how to compare properly?
+        if (healthFactor < healthFactorTarget) {
+            // If the health factor is below the target, repay debt to increase the health factor.
+            // TODO: Implement branch - Can this be combined with the calculations below?
+            //       It would show the max amount to borrow, but just repaying that amount - would it be a negative?
+        } else if (healthFactor > healthFactorTarget) {
+            // If the health factor is above the target, borrow more USDC and reinvest.
+            /* 
+            *   Calculate the maximum amount of USDC that can be borrowed.
+            *       - Minus totalDebtBase from totalCollateralBase to get the actual collateral not including reinvested debt.
+            *       - At the end, minus totalDebtBase to get the remaining amount to borrow to reach the target health factor.
+            *       - currentLiquidationThreshold is a percentage with 4 decimal places e.g. 8250 = 82.5%.
+            *       - healthFactorTarget is a value with 2 decimal places e.g. 200 = 2.00.
+            *       - totalCollateralBase is in USD base unit with 8 decimals to the dollar e.g. 100000000 = $1.00.
+            *       - totalDebtBase is in USD base unit with 8 decimals to the dollar e.g. 100000000 = $1.00.
+            *       - 1e2 used as healthFactorTarget has 2 decimal places.
+            *
+            *                  | ((totalCollateralBase - totalDebtBase) * currentLiquidationThreshold ) |
+            *  maxBorrowUSDC = |------------------------------------------------------------------------| - totalDebtBase
+            *                  |        ((healthFactorTarget * 1e2) - currentLiquidationThreshold)      |
+            */
+            uint256 maxBorrowUSDC = (
+                ((totalCollateralBase - totalDebtBase) * currentLiquidationThreshold)
+                    / ((healthFactorTarget * 1e2) - currentLiquidationThreshold)
+            ) - totalDebtBase;
+
+            // aaveBorrowUSDC input parameter is decimals to the dollar, so divide by 1e2 to get the correct amount.
+            aaveBorrowUSDC(maxBorrowUSDC / 1e2);
+
+            // Swap borrowed USDC to WETH.
+            swapTokens("USDC/ETH", "USDC", "ETH");
+
+            // Convert WETH to wstETH.
+            swapTokens("wstETH/ETH", "ETH", "wstETH");
+
+            // Deposit additional wstETH into Aave.
+            aaveSupplyWstETH();
+        }
+
+        // TODO: Should their be a check for the final health factor? Or should it be left to the user to monitor?
+    }
+
     // ================================================================
     // │                       FUNCTIONS - GETTERS                    │
     // ================================================================

diff --git a/src/interfaces/IAavePM.sol b/src/interfaces/IAavePM.sol
@@ -106,6 +106,11 @@ interface IAavePM {
         string memory _tokenOutIdentifier
     ) external view returns (uint256 minOut);
 
+    // ================================================================
+    // │                    FUNCTIONS - CORE FEATURES                 │
+    // ================================================================
+    function rebalance() external;
+
     // ================================================================
     // │                       FUNCTIONS - GETTERS                    │
     // ================================================================

diff --git a/test/unit/AavePMTest.t.sol b/test/unit/AavePMTest.t.sol
@@ -40,9 +40,11 @@ contract AavePMTestSetup is Test {
     uint256 constant SEND_VALUE = 1 ether;
     uint256 constant STARTING_BALANCE = 10 ether;
     uint256 constant USDC_BORROW_AMOUNT = 100;
+    uint16 constant INCREASED_HEALTH_FACTOR_TARGET = 300;
+    uint16 constant DECREASED_HEALTH_FACTOR_TARGET = 200;
     uint16 constant INITIAL_HEALTH_FACTOR_TARGET_MINIMUM = 200;
-    uint16 constant UPDATED_HEALTH_FACTOR_TARGET_MINIMUM = 250;
     uint24 constant UPDATED_UNISWAPV3_POOL_FEE = 200;
+    uint256 constant AAVE_HEALTH_FACTOR_DIVISOR = 1e16; // Used to convert e.g. 2000003260332359246 into 200
 
     // Create users
     address owner1 = makeAddr("owner1");
@@ -230,7 +232,7 @@ contract AavePMUpdateTests is AavePMTestSetup {
     }
 
     function test_UpdateHealthFactorTarget() public {
-        uint16 newHealthFactorTarget = UPDATED_HEALTH_FACTOR_TARGET_MINIMUM;
+        uint16 newHealthFactorTarget = INCREASED_HEALTH_FACTOR_TARGET;
         uint16 previousHealthFactorTarget = aavePM.getHealthFactorTarget();
 
         vm.expectEmit();
@@ -500,6 +502,47 @@ contract AavePMTokenSwapTests is AavePMTestSetup {
     }
 }
 
+// ================================================================
+// │                        CORE FEATURE TESTS                    │
+// ================================================================
+contract CoreFeatureTests is AavePMTestSetup {
+    function test_Rebalance() public {
+        vm.startPrank(manager1);
+        // Send some ETH to the contract
+        (bool success,) = address(aavePM).call{value: SEND_VALUE}("");
+        require(success, "Failed to send ETH to AavePM contract");
+
+        aavePM.rebalance();
+
+        (,,,,, uint256 endHealthFactor) = aavePM.getAaveAccountData();
+        uint256 endHealthFactorScaled = endHealthFactor / AAVE_HEALTH_FACTOR_DIVISOR;
+
+        require(endHealthFactorScaled <= (aavePM.getHealthFactorTarget() + 1));
+        require(endHealthFactorScaled >= (aavePM.getHealthFactorTarget() - 1));
+        vm.stopPrank();
+    }
+
+    function test_RebalanceUpdateHealthFactor() public {
+        test_Rebalance();
+
+        // Update the health factor target
+        vm.prank(owner1);
+        aavePM.updateHealthFactorTarget(DECREASED_HEALTH_FACTOR_TARGET);
+
+        vm.startPrank(manager1);
+        aavePM.rebalance();
+
+        (,,,,, uint256 endHealthFactor) = aavePM.getAaveAccountData();
+        uint256 endHealthFactorScaled = endHealthFactor / AAVE_HEALTH_FACTOR_DIVISOR;
+
+        require(endHealthFactorScaled <= (aavePM.getHealthFactorTarget() + 1));
+        require(endHealthFactorScaled >= (aavePM.getHealthFactorTarget() - 1));
+        vm.stopPrank();
+    }
+
+    // TODO: Add additional tests for the rebalance function for non-empty Aave accounts
+}
+
 // ================================================================
 // │                         GETTER TESTS                         │
 // ================================================================