Skip to content

Commit

Permalink
Improve state updates
Browse files Browse the repository at this point in the history
  • Loading branch information
@EridianAlpha
EridianAlpha committed Jun 3, 2024
1 parent ec030bf commit dc0b05b
Show file tree
Hide file tree
Showing 3 changed files with 128 additions and 38 deletions.
64 changes: 47 additions & 17 deletions src/AavePM.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -293,10 +293,20 @@ contract AavePM is
// Convert any existing tokens to wstETH and supply to Aave.
aaveSupply();

// TODO: Calculate what states to update
// During a rebalance, I want to know how much debt was repaid.
// It will initially take off the reinvested debt, then from collateral.
_rebalance();
// During a rebalance, I want to know how much debt was repaid and how much collateral was sold.
(uint256 repaymentAmountUSDC, uint256 soldCollateral) = _rebalance();

if (repaymentAmountUSDC > 0 && s_reinvestedDebtTotal >= repaymentAmountUSDC) {
s_reinvestedDebtTotal -= repaymentAmountUSDC;
} else {
s_reinvestedDebtTotal = 0;
}

if (soldCollateral > 0 && s_reinvestedCollateralTotal >= soldCollateral) {
s_reinvestedCollateralTotal -= soldCollateral;
} else {
s_reinvestedCollateralTotal = 0;
}
}

/// @notice // TODO: Add comment
Expand Down Expand Up @@ -328,10 +338,11 @@ contract AavePM is

// When a debt repayment is made, s_withdrawnUSDCTotal should be updated to reflect the repayment, up to 0.
// and any excess should be taken off the s_reinvestedDebtTotal.
if (s_withdrawnUSDCTotal >= usdcBalance) {
uint256 initialWithdrawnUSDCTotal = s_withdrawnUSDCTotal;
if (initialWithdrawnUSDCTotal >= usdcBalance) {
s_withdrawnUSDCTotal -= usdcBalance;
} else if (s_withdrawnUSDCTotal + s_reinvestedDebtTotal >= usdcBalance) {
s_reinvestedDebtTotal -= (usdcBalance - s_withdrawnUSDCTotal);
} else if (initialWithdrawnUSDCTotal + s_reinvestedDebtTotal >= usdcBalance) {
s_reinvestedDebtTotal -= (usdcBalance - initialWithdrawnUSDCTotal);
s_withdrawnUSDCTotal = 0;
} else {
s_reinvestedDebtTotal = 0;
Expand All @@ -341,18 +352,45 @@ contract AavePM is

/// @notice // TODO: Add comment
function aaveWithdrawWstETH(uint256 _amount, address _owner) public onlyRole(MANAGER_ROLE) {
address aavePoolAddress = getContractAddress("aavePool");

// Get the collateral base value before withdrawing.
(uint256 initialCollateralBase,,,,,) = IPool(aavePoolAddress).getUserAccountData(address(this));

address wstETHAddress = getTokenAddress("wstETH");
_aaveWithdrawCollateral(getContractAddress("aavePool"), wstETHAddress, _amount);
_aaveWithdrawCollateral(aavePoolAddress, wstETHAddress, _amount);
IERC20(wstETHAddress).transfer(_owner, _amount);

// TODO: Update s_suppliedCollateralTotal to reflect the withdrawal.
// Get the collateral base value after withdrawing.
(uint256 endCollateralBase,,,,,) = IPool(aavePoolAddress).getUserAccountData(address(this));

uint256 initialSuppliedCollateral = getSuppliedCollateralTotal();

if (endCollateralBase < initialCollateralBase) {
uint256 collateralDeltaBase = initialCollateralBase - endCollateralBase;
if (initialSuppliedCollateral >= collateralDeltaBase / 1e2) {
s_suppliedCollateralTotal -= collateralDeltaBase / 1e2;
} else if (initialSuppliedCollateral + s_reinvestedCollateralTotal >= collateralDeltaBase / 1e2) {
s_suppliedCollateralTotal = 0;
s_reinvestedCollateralTotal -= (collateralDeltaBase / 1e2 - initialSuppliedCollateral);
} else {
s_suppliedCollateralTotal = 0;
s_reinvestedCollateralTotal = 0;
}
}

// Check to ensure the health factor is above the minimum target.
_checkHealthFactorAboveMinimum();
}

/// @notice // TODO: Add comment
function borrowAndWithdrawUSDC(uint256 _amount, address _owner) public onlyRole(MANAGER_ROLE) {
// TODO: Add a return value to _borrowAndWithdrawUSDC so that it confirms the borrow was successful and the amount borrowed.
_borrowAndWithdrawUSDC(_amount, _owner);
s_withdrawnUSDCTotal += _amount;

// Check to ensure the health factor is above the minimum target.
_checkHealthFactorAboveMinimum();
}

// ================================================================
Expand Down Expand Up @@ -456,14 +494,6 @@ contract AavePM is
return members;
}

/// @notice Getter function to get the total debt interest.
/// @dev Public function to allow anyone to view the total debt interest.
/// @return totalDebtInterest The total debt interest.
function getTotalDebtInterest() public view returns (uint256 totalDebtInterest) {
(, uint256 totalDebtBase,,,,) = IPool(getContractAddress("aavePool")).getUserAccountData(address(this));
return _getTotalDebtInterest(totalDebtBase, getReinvestedDebtTotal(), getWithdrawnUSDCTotal());
}

/// @notice Getter function to get the total amount of USDC withdrawn.
/// @dev Public function to allow anyone to view the total amount of USDC withdrawn.
/// @return withdrawnUSDCTotal The total amount of USDC withdrawn.
Expand Down
29 changes: 20 additions & 9 deletions src/Rebalance.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,11 @@ contract Rebalance is TokenSwaps, AaveFunctions {
/// @dev Caller must have `MANAGER_ROLE`.
/// The function rebalances the Aave position.
/// If the health factor is below the target, it repays debt to increase the health factor.
function _rebalance() internal {
function _rebalance() internal returns (uint256 repaymentAmountUSDC, uint256 soldCollateral) {
IAavePM aavePM = IAavePM(address(this));

(
uint256 totalCollateralBase,
uint256 initialCollateralBase,
uint256 totalDebtBase,
uint256 currentLiquidationThreshold,
uint256 initialHealthFactorScaled,
Expand All @@ -45,11 +45,11 @@ contract Rebalance is TokenSwaps, AaveFunctions {

// If the health factor is below the target, repay debt to increase the health factor.
if (initialHealthFactorScaled < (healthFactorTarget - healthFactorTargetRange)) {
_repayDebt(
repaymentAmountUSDC = _repayDebt(
totalDebtBase,
aavePoolAddress,
usdcAddress,
totalCollateralBase,
initialCollateralBase,
currentLiquidationThreshold,
healthFactorTarget
);
Expand All @@ -59,8 +59,17 @@ contract Rebalance is TokenSwaps, AaveFunctions {

// Safety check to ensure the health factor is above the minimum target.
// It is also used to calculate the ?? by returning the updated position values.
/* (uint256 endCollateralBase,,,,,) = */
_checkHealthFactorAboveMinimum();
(uint256 endCollateralBase,,,,,) = _checkHealthFactorAboveMinimum();

// Calculate the sold collateral by comparing the initial and end collateral values.
if (endCollateralBase < initialCollateralBase) {
soldCollateral = (initialCollateralBase - endCollateralBase) / 1e2;
} else {
soldCollateral = 0;
}

// Return the reinvested debt and reinvested collateral so the state can be updated on the AavePM contract.
return (repaymentAmountUSDC, soldCollateral);
}

function _repayDebt(
Expand All @@ -70,16 +79,18 @@ contract Rebalance is TokenSwaps, AaveFunctions {
uint256 totalCollateralBase,
uint256 currentLiquidationThreshold,
uint16 healthFactorTarget
) private {
) private returns (uint256 repaymentAmountUSDC) {
// Calculate the maximum amount of USDC that can be borrowed.
uint256 maxBorrowUSDC =
_calculateMaxBorrowUSDC(totalCollateralBase, totalDebtBase, currentLiquidationThreshold, healthFactorTarget);

// Calculate the repayment amount required to reach the target health factor.
uint256 repaymentAmountUSDC = totalDebtBase - maxBorrowUSDC;
repaymentAmountUSDC = (totalDebtBase - maxBorrowUSDC) / 1e2;

// Take out a flash loan for the USDC amount needed to repay and rebalance the health factor.
// flashLoanSimple `amount` input parameter is decimals to the dollar, so divide by 1e2 to get the correct amount
IPool(aavePoolAddress).flashLoanSimple(address(this), usdcAddress, repaymentAmountUSDC / 1e2, bytes(""), 0);
IPool(aavePoolAddress).flashLoanSimple(address(this), usdcAddress, repaymentAmountUSDC, bytes(""), 0);

return repaymentAmountUSDC;
}
}
73 changes: 61 additions & 12 deletions test/unit/AavePMGetterTest.t.sol
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
// SPDX-License-Identifier: MIT
pragma solidity 0.8.24;

import {console} from "forge-std/Test.sol";

import {AavePMTestSetup} from "test/unit/AavePMTestSetupTest.t.sol";

// ================================================================
Expand Down Expand Up @@ -80,18 +82,65 @@ contract AavePMGetterTests is AavePMTestSetup {
assertEq(owners[0], owner1);
}

// TODO: Fix these tests
// function test_GetContractBalanceWstETH() public {
// vm.startPrank(manager1);
// // Send some ETH to the contract and wrap it to WETH
// sendEth(address(aavePM), SEND_VALUE);
// aavePM.wrapETHToWETH();
function test_GetWithdrawnUSDCTotal() public {
// Send ETH from manager1 to the contract
vm.startPrank(manager1);
sendEth(address(aavePM), SEND_VALUE);

// Supply ETH to Aave
aavePM.aaveSupply();

// Withdraw USDC from Aave
aavePM.borrowAndWithdrawUSDC(USDC_BORROW_AMOUNT, owner1);
vm.stopPrank();

assertEq(aavePM.getWithdrawnUSDCTotal(), USDC_BORROW_AMOUNT);
}

function test_GetReinvestedDebtTotal() public {
// Send ETH from manager1 to the contract
vm.startPrank(manager1);
sendEth(address(aavePM), SEND_VALUE);

// Supply ETH to Aave
aavePM.aaveSupply();

// // Call the swapTokens function to get wstETH
// aavePM.swapTokens("wstETH/ETH", "ETH", "wstETH");
// Reinvest
(uint256 reinvestedDebt,) = aavePM.reinvest();
vm.stopPrank();

// // Check the wstETH balance of the contract
// assertEq(aavePM.getContractBalance("wstETH"), wstETH.balanceOf(address(aavePM)));
// vm.stopPrank();
// }
assertEq(aavePM.getReinvestedDebtTotal(), reinvestedDebt);
}

function test_GetTotalCollateralDelta() public {
// Not a great test, but since it's a calculation that's based on the passage of time, but also on a mainnet fork
// it can't easily be tested in a unit test.
assertEq(aavePM.getReinvestedDebtTotal(), 0);
}

function test_GetSuppliedCollateralTotal() public {
// Send ETH from manager1 to the contract
vm.startPrank(manager1);
sendEth(address(aavePM), SEND_VALUE);

// Supply ETH to Aave
uint256 suppliedCollateral = aavePM.aaveSupply();

assertEq(aavePM.getSuppliedCollateralTotal(), suppliedCollateral);
}

function test_GetReinvestedCollateralTotal() public {
// Send ETH from manager1 to the contract
vm.startPrank(manager1);
sendEth(address(aavePM), SEND_VALUE);

// Supply ETH to Aave
aavePM.aaveSupply();

// Reinvest
(, uint256 reinvestedCollateral) = aavePM.reinvest();
vm.stopPrank();

assertEq(aavePM.getReinvestedCollateralTotal(), reinvestedCollateral);
}
}

0 comments on commit dc0b05b

Please sign in to comment.