Refactor module interfaces and selectors

EridianAlpha · Jun 16, 2024 · d8d33d5 · d8d33d5
1 parent 442df97
commit d8d33d5
Show file tree

Hide file tree

Showing 13 changed files with 226 additions and 66 deletions.
diff --git a/src/AavePM.sol b/src/AavePM.sol
@@ -26,6 +26,7 @@ import {IPool} from "@aave/aave-v3-core/contracts/interfaces/IPool.sol";
 
 // Interface Imports
 import {IAavePM} from "./interfaces/IAavePM.sol";
+import {IAaveFunctionsModule} from "./interfaces/IAaveFunctionsModule.sol";
 
 // ================================================================
 // │                       AAVEPM CONTRACT                        │
@@ -370,7 +371,12 @@ contract AavePM is
     /// @notice // TODO: Add comment
     function aaveSupplyFromContractBalance() public onlyRole(MANAGER_ROLE) returns (uint256 suppliedCollateral) {
         suppliedCollateral = abi.decode(
-            delegateCallHelper("aaveFunctionsModule", "convertExistingBalanceToWstETHAndSupplyToAave()", new bytes(0)),
+            delegateCallHelper(
+                "aaveFunctionsModule",
+                abi.encodeWithSelector(
+                    IAaveFunctionsModule.convertExistingBalanceToWstETHAndSupplyToAave.selector, new bytes(0)
+                )
+            ),
             (uint256)
         );
         if (suppliedCollateral > 0) s_suppliedCollateralTotal += suppliedCollateral;
@@ -384,8 +390,12 @@ contract AavePM is
 
         delegateCallHelper(
             "aaveFunctionsModule",
-            "aaveRepayDebt(address,address,uint256)",
-            abi.encode(getContractAddress("aavePool"), getTokenAddress("USDC"), usdcBalance)
+            abi.encodeWithSelector(
+                IAaveFunctionsModule.aaveRepayDebt.selector,
+                getContractAddress("aavePool"),
+                getTokenAddress("USDC"),
+                usdcBalance
+            )
         );
 
         // When a debt repayment is made, s_withdrawnUSDCTotal should be updated to reflect the repayment, up to 0.
@@ -403,13 +413,9 @@ contract AavePM is
     }
 
     /// @notice // TODO: Add comment and move to different heading. How to protect this function? Does it need protecting?
-    function delegateCallHelper(string memory _targetIdentifier, string memory _functionSignature, bytes memory _args)
-        public
-        returns (bytes memory)
-    {
+    function delegateCallHelper(string memory _targetIdentifier, bytes memory _data) public returns (bytes memory) {
         address target = getContractAddress(_targetIdentifier);
-        bytes memory data = abi.encodePacked(abi.encodeWithSignature(_functionSignature), _args);
-        (bool success, bytes memory result) = target.delegatecall(data);
+        (bool success, bytes memory result) = target.delegatecall(_data);
         if (!success) revert("AavePM__DelegateCallFailed");
         return result;
     }
@@ -471,8 +477,9 @@ contract AavePM is
         address wstETHAddress = getTokenAddress("wstETH");
         delegateCallHelper(
             "aaveFunctionsModule",
-            "aaveWithdrawCollateral(address,address,uint256)",
-            abi.encode(aavePoolAddress, wstETHAddress, _amount)
+            abi.encodeWithSelector(
+                IAaveFunctionsModule.aaveWithdrawCollateral.selector, aavePoolAddress, wstETHAddress, _amount
+            )
         );
 
         IERC20(wstETHAddress).transfer(_owner, _amount);
@@ -496,7 +503,10 @@ contract AavePM is
         }
 
         // Check to ensure the health factor is above the minimum target.
-        delegateCallHelper("aaveFunctionsModule", "checkHealthFactorAboveMinimum()", new bytes(0));
+        delegateCallHelper(
+            "aaveFunctionsModule",
+            abi.encodeWithSelector(IAaveFunctionsModule.checkHealthFactorAboveMinimum.selector, new bytes(0))
+        );
 
         // TODO: Emit an event
     }
@@ -539,7 +549,10 @@ contract AavePM is
         }
 
         // Check to ensure the health factor is above the minimum target.
-        delegateCallHelper("aaveFunctionsModule", "checkHealthFactorAboveMinimum()", new bytes(0));
+        delegateCallHelper(
+            "aaveFunctionsModule",
+            abi.encodeWithSelector(IAaveFunctionsModule.checkHealthFactorAboveMinimum.selector, new bytes(0))
+        );
     }
 
     /// @notice // TODO: Add comment
@@ -704,8 +717,12 @@ contract AavePM is
         return abi.decode(
             delegateCallHelper(
                 "aaveFunctionsModule",
-                "getTotalCollateralDelta(uint256,uint256,uint256)",
-                abi.encode(totalCollateralBase, getReinvestedDebtTotal(), getSuppliedCollateralTotal())
+                abi.encodeWithSelector(
+                    IAaveFunctionsModule.getTotalCollateralDelta.selector,
+                    totalCollateralBase,
+                    getReinvestedDebtTotal(),
+                    getSuppliedCollateralTotal()
+                )
             ),
             (uint256, bool)
         );
@@ -744,8 +761,13 @@ contract AavePM is
         uint256 maxBorrowUSDC = abi.decode(
             delegateCallHelper(
                 "aaveFunctionsModule",
-                "calculateMaxBorrowUSDC(uint256,uint256,uint256,uint16)",
-                abi.encode(totalCollateralBase, totalDebtBase, currentLiquidationThreshold, getHealthFactorTarget())
+                abi.encodeWithSelector(
+                    IAaveFunctionsModule.calculateMaxBorrowUSDC.selector,
+                    totalCollateralBase,
+                    totalDebtBase,
+                    currentLiquidationThreshold,
+                    getHealthFactorTarget()
+                )
             ),
             (uint256)
         );

diff --git a/src/BorrowAndWithdrawUSDC.sol b/src/BorrowAndWithdrawUSDC.sol
@@ -11,6 +11,7 @@ import {IPool} from "@aave/aave-v3-core/contracts/interfaces/IPool.sol";
 // Interface Imports
 import {IAavePM} from "./interfaces/IAavePM.sol";
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {IAaveFunctionsModule} from "./interfaces/IAaveFunctionsModule.sol";
 
 // ================================================================
 // │               BORROW AND WITHDRAW USDC CONTRACT              │
@@ -46,8 +47,9 @@ contract BorrowAndWithdrawUSDC {
             // so the USDC can be borrowed without repaying reinvested debt
             aavePM.delegateCallHelper(
                 "aaveFunctionsModule",
-                "aaveBorrow(address,address,uint256)",
-                abi.encode(aavePoolAddress, usdcAddress, borrowAmountUSDC)
+                abi.encodeWithSelector(
+                    IAaveFunctionsModule.aaveBorrow.selector, aavePoolAddress, usdcAddress, borrowAmountUSDC
+                )
             );
         } else if (aavePM.getReinvestedDebtTotal() > 0) {
             // The requested borrow amount would put the HF below the target
@@ -62,8 +64,9 @@ contract BorrowAndWithdrawUSDC {
             // Borrow the requested amount of USDC
             aavePM.delegateCallHelper(
                 "aaveFunctionsModule",
-                "aaveBorrow(address,address,uint256)",
-                abi.encode(aavePoolAddress, usdcAddress, borrowAmountUSDC)
+                abi.encodeWithSelector(
+                    IAaveFunctionsModule.aaveBorrow.selector, aavePoolAddress, usdcAddress, borrowAmountUSDC
+                )
             );
         }
 

diff --git a/src/FlashLoan.sol b/src/FlashLoan.sol
@@ -15,6 +15,8 @@ import {TransferHelper} from "@uniswap/v3-periphery/contracts/libraries/Transfer
 // Interface Imports
 import {IAavePM} from "./interfaces/IAavePM.sol";
 import {IWETH9} from "./interfaces/IWETH9.sol";
+import {ITokenSwapsModule} from "./interfaces/ITokenSwapsModule.sol";
+import {IAaveFunctionsModule} from "./interfaces/IAaveFunctionsModule.sol";
 
 // ================================================================
 // │                   AAVEFUNCTIONS CONTRACT                     │
@@ -53,8 +55,9 @@ contract FlashLoan {
         // Use the flash loan USDC to repay the debt.
         aavePM.delegateCallHelper(
             "aaveFunctionsModule",
-            "aaveRepayDebt(address,address,uint256)",
-            abi.encode(aavePoolAddress, aavePM.getTokenAddress("USDC"), amount)
+            abi.encodeWithSelector(
+                IAaveFunctionsModule.aaveRepayDebt.selector, aavePoolAddress, aavePM.getTokenAddress("USDC"), amount
+            )
         );
 
         // Now the HF is higher, withdraw the corresponding amount of wstETH from collateral.
@@ -72,16 +75,21 @@ contract FlashLoan {
         // Withdraw the wstETH from Aave.
         aavePM.delegateCallHelper(
             "aaveFunctionsModule",
-            "aaveWithdrawCollateral(address,address,uint256)",
-            abi.encode(aavePoolAddress, wstETHAddress, wstETHToWithdrawSlippageAllowance)
+            abi.encodeWithSelector(
+                IAaveFunctionsModule.aaveWithdrawCollateral.selector,
+                aavePoolAddress,
+                wstETHAddress,
+                wstETHToWithdrawSlippageAllowance
+            )
         );
 
         // Convert the wstETH to USDC.
         aavePM.delegateCallHelper(
-            "tokenSwapsModule", "swapTokens(string,string,string)", abi.encode("wstETH/ETH", "wstETH", "ETH")
+            "tokenSwapsModule",
+            abi.encodeWithSelector(ITokenSwapsModule.swapTokens.selector, "wstETH/ETH", "wstETH", "ETH")
         );
         aavePM.delegateCallHelper(
-            "tokenSwapsModule", "swapTokens(string,string,string)", abi.encode("USDC/ETH", "ETH", "USDC")
+            "tokenSwapsModule", abi.encodeWithSelector(ITokenSwapsModule.swapTokens.selector, "USDC/ETH", "ETH", "USDC")
         );
         TransferHelper.safeApprove(asset, aavePoolAddress, repaymentAmountTotalUSDC);
         return true;

diff --git a/src/Rebalance.sol b/src/Rebalance.sol
@@ -10,6 +10,7 @@ import {IPool} from "@aave/aave-v3-core/contracts/interfaces/IPool.sol";
 
 // Interface Imports
 import {IAavePM} from "./interfaces/IAavePM.sol";
+import {IAaveFunctionsModule} from "./interfaces/IAaveFunctionsModule.sol";
 
 // ================================================================
 // │                       REBALANCE CONTRACT                     │
@@ -35,7 +36,10 @@ contract Rebalance {
             ,
             address usdcAddress
         ) = abi.decode(
-            aavePM.delegateCallHelper("aaveFunctionsModule", "getCurrentPositionValues(address)", abi.encode(aavePM)),
+            aavePM.delegateCallHelper(
+                "aaveFunctionsModule",
+                abi.encodeWithSelector(IAaveFunctionsModule.getCurrentPositionValues.selector, aavePM)
+            ),
             (uint256, uint256, uint256, uint256, uint16, address, address, address)
         );
 
@@ -57,7 +61,10 @@ contract Rebalance {
         }
 
         // Safety check to ensure the health factor is above the minimum target.
-        aavePM.delegateCallHelper("aaveFunctionsModule", "checkHealthFactorAboveMinimum()", new bytes(0));
+        aavePM.delegateCallHelper(
+            "aaveFunctionsModule",
+            abi.encodeWithSelector(IAaveFunctionsModule.checkHealthFactorAboveMinimum.selector, new bytes(0))
+        );
 
         // Return the reinvested debt and reinvested collateral so the state can be updated on the AavePM contract.
         return (repaymentAmountUSDC);
@@ -81,8 +88,13 @@ contract Rebalance {
             uint256 maxBorrowUSDC = abi.decode(
                 IAavePM(address(this)).delegateCallHelper(
                     "aaveFunctionsModule",
-                    "calculateMaxBorrowUSDC(uint256,uint256,uint256,uint16)",
-                    abi.encode(totalCollateralBase, totalDebtBase, currentLiquidationThreshold, healthFactorTarget)
+                    abi.encodeWithSelector(
+                        IAaveFunctionsModule.calculateMaxBorrowUSDC.selector,
+                        totalCollateralBase,
+                        totalDebtBase,
+                        currentLiquidationThreshold,
+                        healthFactorTarget
+                    )
                 ),
                 (uint256)
             );

diff --git a/src/Reinvest.sol b/src/Reinvest.sol
@@ -10,6 +10,8 @@ import {IPool} from "@aave/aave-v3-core/contracts/interfaces/IPool.sol";
 
 // Interface Imports
 import {IAavePM} from "./interfaces/IAavePM.sol";
+import {ITokenSwapsModule} from "./interfaces/ITokenSwapsModule.sol";
+import {IAaveFunctionsModule} from "./interfaces/IAaveFunctionsModule.sol";
 
 // ================================================================
 // │                        REINVEST CONTRACT                     │
@@ -34,7 +36,10 @@ contract Reinvest {
             address wstETHAddress,
             address usdcAddress
         ) = abi.decode(
-            aavePM.delegateCallHelper("aaveFunctionsModule", "getCurrentPositionValues(address)", abi.encode(aavePM)),
+            aavePM.delegateCallHelper(
+                "aaveFunctionsModule",
+                abi.encodeWithSelector(IAaveFunctionsModule.getCurrentPositionValues.selector, aavePM)
+            ),
             (uint256, uint256, uint256, uint256, uint16, address, address, address)
         );
 
@@ -58,7 +63,10 @@ contract Reinvest {
         }
 
         // Safety check to ensure the health factor is above the minimum target.
-        aavePM.delegateCallHelper("aaveFunctionsModule", "checkHealthFactorAboveMinimum()", new bytes(0));
+        aavePM.delegateCallHelper(
+            "aaveFunctionsModule",
+            abi.encodeWithSelector(IAaveFunctionsModule.checkHealthFactorAboveMinimum.selector, new bytes(0))
+        );
 
         // Return the reinvested debt and reinvested collateral so the state can be updated on the AavePM contract.
         return (reinvestedDebt);
@@ -79,8 +87,13 @@ contract Reinvest {
         uint256 maxBorrowUSDC = abi.decode(
             IAavePM(address(this)).delegateCallHelper(
                 "aaveFunctionsModule",
-                "calculateMaxBorrowUSDC(uint256,uint256,uint256,uint16)",
-                abi.encode(initialCollateralBase, totalDebtBase, currentLiquidationThreshold, healthFactorTarget)
+                abi.encodeWithSelector(
+                    IAaveFunctionsModule.calculateMaxBorrowUSDC.selector,
+                    initialCollateralBase,
+                    totalDebtBase,
+                    currentLiquidationThreshold,
+                    healthFactorTarget
+                )
             ),
             (uint256)
         );
@@ -92,22 +105,28 @@ contract Reinvest {
         borrowAmountUSDC = additionalBorrowUSDC / 1e2;
         aavePM.delegateCallHelper(
             "aaveFunctionsModule",
-            "aaveBorrow(address,address,uint256)",
-            abi.encode(aavePoolAddress, usdcAddress, borrowAmountUSDC)
+            abi.encodeWithSelector(
+                IAaveFunctionsModule.aaveBorrow.selector, aavePoolAddress, usdcAddress, borrowAmountUSDC
+            )
         );
 
         // Swap borrowed USDC ➜ WETH ➜ wstETH then supply to Aave.
         aavePM.delegateCallHelper(
-            "tokenSwapsModule", "swapTokens(string,string,string)", abi.encode("USDC/ETH", "USDC", "ETH")
+            "tokenSwapsModule", abi.encodeWithSelector(ITokenSwapsModule.swapTokens.selector, "USDC/ETH", "USDC", "ETH")
         );
         aavePM.delegateCallHelper(
-            "tokenSwapsModule", "swapTokens(string,string,string)", abi.encode("wstETH/ETH", "ETH", "wstETH")
+            "tokenSwapsModule",
+            abi.encodeWithSelector(ITokenSwapsModule.swapTokens.selector, "wstETH/ETH", "ETH", "wstETH")
         );
 
         aavePM.delegateCallHelper(
             "aaveFunctionsModule",
-            "aaveSupply(address,address,uint256)",
-            abi.encode(aavePoolAddress, wstETHAddress, aavePM.getContractBalance("wstETH"))
+            abi.encodeWithSelector(
+                IAaveFunctionsModule.aaveSupply.selector,
+                aavePoolAddress,
+                wstETHAddress,
+                aavePM.getContractBalance("wstETH")
+            )
         );
 
         return borrowAmountUSDC;

diff --git a/src/interfaces/IAaveFunctionsModule.sol b/src/interfaces/IAaveFunctionsModule.sol
@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.24;
+
+import {IAavePM} from "./IAavePM.sol";
+
+interface IAaveFunctionsModule {
+    function aaveSupply(address aavePoolAddress, address tokenAddress, uint256 tokenBalance) external;
+    function aaveWithdrawCollateral(address aavePoolAddress, address tokenAddress, uint256 withdrawAmount) external;
+    function aaveBorrow(address aavePoolAddress, address tokenAddress, uint256 borrowAmount) external;
+    function aaveRepayDebt(address aavePoolAddress, address tokenAddress, uint256 repayAmount) external;
+    function getCurrentPositionValues(IAavePM aavePM)
+        external
+        view
+        returns (
+            uint256 initialCollateralBase,
+            uint256 totalDebtBase,
+            uint256 currentLiquidationThreshold,
+            uint256 initialHealthFactorScaled,
+            uint16 healthFactorTarget,
+            address aavePoolAddress,
+            address wstETHAddress,
+            address usdcAddress
+        );
+    function checkHealthFactorAboveMinimum()
+        external
+        view
+        returns (
+            uint256 totalCollateralBase,
+            uint256 totalDebtBase,
+            uint256 availableBorrowsBase,
+            uint256 currentLiquidationThreshold,
+            uint256 ltv,
+            uint256 healthFactor
+        );
+    function getTotalCollateralDelta(
+        uint256 totalCollateralBase,
+        uint256 reinvestedDebtTotal,
+        uint256 suppliedCollateralTotal
+    ) external pure returns (uint256 delta, bool isPositive);
+    function convertExistingBalanceToWstETHAndSupplyToAave() external returns (uint256 suppliedCollateral);
+    function calculateMaxBorrowUSDC(
+        uint256 totalCollateralBase,
+        uint256 totalDebtBase,
+        uint256 currentLiquidationThreshold,
+        uint16 healthFactorTarget
+    ) external pure returns (uint256 maxBorrowUSDC);
+}
diff --git a/src/interfaces/IAavePM.sol b/src/interfaces/IAavePM.sol
@@ -106,9 +106,7 @@ interface IAavePM {
     function aaveRepayUSDCFromContractBalance() external;
 
     /// TODO: Move to different heading
-    function delegateCallHelper(string memory _targetIdentifier, string memory _functionSignature, bytes memory _args)
-        external
-        returns (bytes memory);
+    function delegateCallHelper(string memory _targetIdentifier, bytes memory _data) external returns (bytes memory);
 
     // ================================================================
     // │                FUNCTIONS - WITHDRAW FUNCTIONS                │