From 3ee796b28d99639aac663e98519a64e7a7e67e3d Mon Sep 17 00:00:00 2001 From: Jake Kidd Date: Fri, 21 Jun 2024 13:00:28 -0400 Subject: [PATCH] feat: shares calc in create; tokens calc in close --- src/contracts/QWManager.sol | 120 ++++++++++++++++++----------- src/contracts/child/QWAaveV2.sol | 127 +++++++++++++++++++------------ src/interfaces/IQWChild.sol | 25 ++++-- src/interfaces/IQWManager.sol | 49 ++++++++---- 4 files changed, 207 insertions(+), 114 deletions(-) diff --git a/src/contracts/QWManager.sol b/src/contracts/QWManager.sol index 7ab376d..bfaf84e 100644 --- a/src/contracts/QWManager.sol +++ b/src/contracts/QWManager.sol @@ -7,16 +7,20 @@ import {IERC20} from '@openzeppelin/contracts/token/ERC20/IERC20.sol'; import {IQWChild} from 'interfaces/IQWChild.sol'; import {IQWManager} from 'interfaces/IQWManager.sol'; import {IQWRegistry} from 'interfaces/IQWRegistry.sol'; -import {QWShares} from './QWShares.sol'; /** * @title Quant Wealth Manager Contract * @notice This contract manages the execution, closing, and withdrawal of various strategies for Quant Wealth. */ -contract QWManager is IQWManager, Ownable, QWShares { +contract QWManager is IQWManager, Ownable { // Variables address public immutable REGISTRY; + // Mapping to store user shares for each protocol + mapping(address => mapping(address => uint256)) public shares; + // Mapping to store total shares for each protocol + mapping(address => uint256) public totalShares; + // Custom errors error InvalidInputLength(); // Error for mismatched input lengths error ContractNotWhitelisted(); // Error for contract not whitelisted @@ -32,100 +36,124 @@ contract QWManager is IQWManager, Ownable, QWShares { // External Functions /** - * @notice Execute a series of investments. + * @notice Execute a series of investments in batches for multiple protocols. * Transfers specified amounts of tokens and calls target contracts with provided calldata. - * @param _targetQwChild List of contract addresses to interact with. - * @param _callData Encoded function calls to be executed on the target contracts. - * @param _tokenAddress Token address to transfer. - * @param _amount Amount of tokens to transfer to each target contract. + * @param batches Array of ExecuteBatch data containing protocol, users, contributions, token, and amount. */ - function execute( - address[] memory _targetQwChild, - bytes[] memory _callData, - address _tokenAddress, - uint256 _amount - ) external onlyOwner { - if (_targetQwChild.length != _callData.length) { - revert InvalidInputLength(); - } + function execute(ExecuteBatch[] memory batches) external onlyOwner { + for (uint256 i = 0; i < batches.length; i++) { + ExecuteBatch memory batch = batches[i]; - for (uint256 i = 0; i < _targetQwChild.length; i++) { // Check if the target contract is whitelisted - if (!IQWRegistry(REGISTRY).whitelist(_targetQwChild[i])) { + if (!IQWRegistry(REGISTRY).whitelist(batch.protocol)) { revert ContractNotWhitelisted(); } // Approve the target contract to spend the specified amount of tokens - IERC20 token = IERC20(_tokenAddress); - token.approve(address(_targetQwChild[i]), _amount); + IERC20 token = IERC20(batch.token); + token.approve(address(batch.protocol), batch.amount); // Encode necessary data for child contract - bytes memory encodedData = abi.encode(totalShares[_targetQwChild[i]]); + bytes memory encodedData = abi.encode(totalShares[batch.protocol]); // Call the create function on the target contract with the provided calldata - (bool success, bytes memory result) = IQWChild(_targetQwChild[i]).create(encodedData, _tokenAddress, _amount); + (bool success, bytes memory result) = IQWChild(batch.protocol).create(encodedData, batch.amount); if (!success) { revert CallFailed(); } // Decode shares from result - (uint256 shares) = abi.decode(result, (uint256)); + (uint256 totalSharesReceived) = abi.decode(result, (uint256)); - // Update shares in QWManager - _updateSharesOnDeposit(msg.sender, shares, _targetQwChild[i]); + // Distribute the shares to users + for (uint256 j = 0; j < batch.users.length; j++) { + uint256 userShare = (totalSharesReceived * batch.contributions[j]) / 10000; + _updateSharesOnDeposit(batch.users[j], userShare, totalShares[batch.protocol], batch.protocol); + } } } /** - * @notice Close a series of investments. + * @notice Close a series of investments in batches for multiple protocols. * Calls target contracts with provided calldata to close positions. - * @param _targetQwChild List of contract addresses to interact with. - * @param _callData Encoded function calls to be executed on the target contracts. + * @param batches Array of CloseBatch data containing protocol, users, contributions, token, and shares. */ - function close(address[] memory _targetQwChild, bytes[] memory _callData) external onlyOwner { - if (_targetQwChild.length != _callData.length) { - revert InvalidInputLength(); - } - - for (uint256 i = 0; i < _targetQwChild.length; i++) { - // Decode the calldata to get the LP asset address and amount - (address _user, uint256 _sharesAmount) = abi.decode(_callData[i], (address, uint256)); + function close(CloseBatch[] memory batches) external onlyOwner { + for (uint256 i = 0; i < batches.length; i++) { + CloseBatch memory batch = batches[i]; // Encode necessary data for child contract - bytes memory encodedData = abi.encode(_user, _sharesAmount, totalShares[_targetQwChild[i]], _targetQwChild[i]); + bytes memory encodedData = abi.encode(totalShares[batch.protocol]); // Call the close function on the target contract with the provided calldata - (bool success) = IQWChild(_targetQwChild[i]).close(encodedData); + (bool success, bytes memory result) = IQWChild(batch.protocol).close( + encodedData, + batch.shares + ); if (!success) { revert CallFailed(); } - // Update shares in QWManager - _updateSharesOnWithdrawal(_user, _sharesAmount, _targetQwChild[i]); + // Decode tokens received from result + (uint256 tokens) = abi.decode(result, (uint256)); + + // Distribute the tokens to users + for (uint256 j = 0; j < batch.users.length; j++) { + // TODO: Handle potential leftover value due to division rounding + _updateSharesOnWithdrawal(batch.users[j], batch.contributions[j], batch.protocol); + uint256 userShares = (batch.shares * batch.contributions[j]) / 10000; + // TODO: transfer userTokens to user + // uint256 userTokens = (tokens * batch.contributions[j]) / 10000; + } } } /** * @notice Withdraw funds to a specified user. * Transfers a specified amount of funds to the user. - * @param user The address of the user to receive the funds. + * @param _user The address of the user to receive the funds. * @param _tokenAddress The address of the token to transfer. * @param _amount The amount of funds to transfer to the user. */ - function withdraw(address user, address _tokenAddress, uint256 _amount) external onlyOwner { + function withdraw(address _user, address _tokenAddress, uint256 _amount) external onlyOwner { IERC20 token = IERC20(_tokenAddress); - token.transfer(user, _amount); + token.transfer(_user, _amount); } /** * @notice Receive funds from a specified user. * Transfers a specified amount of funds from the user to this contract. - * @param user The address of the user sending the funds. + * @param _user The address of the user sending the funds. * @param _tokenAddress The address of the token to transfer. * @param _amount The amount of funds to transfer to this contract. */ - function receiveFunds(address user, address _tokenAddress, uint256 _amount) external { + function receiveFunds(address _user, address _tokenAddress, uint256 _amount) external { IERC20 token = IERC20(_tokenAddress); - token.transferFrom(user, address(this), _amount); + token.transferFrom(_user, address(this), _amount); + } + + // Internal Functions + + /** + * @notice Internal function to update shares on deposit. + * @param user The address of the user whose shares are being updated. + * @param sharesAmount The amount of shares to add to the user's balance. + * @param totalSharesProtocol The total shares in the protocol before the deposit. + * @param protocol The address of the protocol. + */ + function _updateSharesOnDeposit(address user, uint256 sharesAmount, uint256 totalSharesProtocol, address protocol) internal { + shares[protocol][user] += sharesAmount; + totalShares[protocol] = totalSharesProtocol + sharesAmount; + } + + /** + * @notice Internal function to update shares on withdrawal. + * @param user The address of the user whose shares are being updated. + * @param userShares The amount of shares to subtract from the user's balance. + * @param protocol The address of the protocol. + */ + function _updateSharesOnWithdrawal(address user, uint256 userShares, address protocol) internal { + shares[protocol][user] -= userShares; + totalShares[protocol] -= userShares; } } diff --git a/src/contracts/child/QWAaveV2.sol b/src/contracts/child/QWAaveV2.sol index b443396..9673c1a 100644 --- a/src/contracts/child/QWAaveV2.sol +++ b/src/contracts/child/QWAaveV2.sol @@ -4,6 +4,7 @@ pragma solidity 0.8.23; import {IERC20} from '@openzeppelin/contracts/token/ERC20/IERC20.sol'; import {IQWChild} from 'interfaces/IQWChild.sol'; import {ILendingPool} from 'interfaces/aave-v2/ILendingPool.sol'; +import {IAToken} from 'interfaces/aave-v2/IAToken.sol'; /** * @title AaveV2 Integration for Quant Wealth @@ -13,6 +14,8 @@ contract QWAaveV2 is IQWChild { // Variables address public immutable QW_MANAGER; address public immutable LENDING_POOL; + address public immutable INVESTMENT_TOKEN; + address public immutable A_INVESTMENT_TOKEN; // Custom errors error InvalidCallData(); // Error for invalid call data @@ -29,98 +32,126 @@ contract QWAaveV2 is IQWChild { * @dev Constructor to initialize the contract with required addresses. * @param _qwManager The address of the Quant Wealth Manager contract. * @param _lendingPool The address of the AaveV2 pool contract. + * @param _investmentToken The address of the investment token (e.g., USDT). + * @param _aInvestmentToken The address of the aToken (e.g., aUSDT). */ - constructor(address _qwManager, address _lendingPool) { + constructor( + address _qwManager, + address _lendingPool, + address _investmentToken, + address _aInvestmentToken + ) { QW_MANAGER = _qwManager; LENDING_POOL = _lendingPool; + INVESTMENT_TOKEN = _investmentToken; + A_INVESTMENT_TOKEN = _aInvestmentToken; } // Functions /** * @notice Executes a transaction on AaveV2 pool to deposit tokens. * @dev This function is called by the parent contract to deposit tokens into the AaveV2 pool. - * @param _callData Encoded function call data containing user address and total shares. - * @param _tokenAddress Address of the token to be deposited. - * @param _amount Amount of tokens to be deposited. + * @param _callData Encoded function call data containing total shares. + * @param _tokenAmount Amount of tokens to be deposited. * @return success boolean indicating the success of the transaction. + * @return shares Number of shares to be allocated to the user in return for investment created. */ function create( bytes memory _callData, - address _tokenAddress, - uint256 _amount - ) external override onlyQwManager returns (bool success) { + uint256 _tokenAmount + ) external override onlyQwManager returns (bool success, uint256 shares) { (uint256 _totalShares) = abi.decode(_callData, (uint256)); - IERC20 token = IERC20(_tokenAddress); - token.transferFrom(QW_MANAGER, address(this), _amount); - token.approve(LENDING_POOL, _amount); + // Transfer tokens from QWManager to this contract. + IERC20 token = IERC20(INVESTMENT_TOKEN); + token.transferFrom(QW_MANAGER, address(this), _tokenAmount); - ILendingPool(LENDING_POOL).deposit(_tokenAddress, _amount, address(this), 0); - uint256 aTokensReceived = IERC20(_tokenAddress).balanceOf(address(this)); + // Approve the Aave lending pool to spend the tokens. + token.approve(LENDING_POOL, _tokenAmount); - uint256 shares; - if (_totalShares == 0) { - shares = aTokensReceived; - } else { - uint256 totalPoolValue = getTotalPoolValue(_tokenAddress); - shares = (aTokensReceived * _totalShares) / totalPoolValue; - } + // Calculate price per share before new investment. This is the price that the investment is + // 'buying' shares of the pool at. + uint256 sharePrice = pricePerShare(_totalShares); - success = true; + // Deposit tokens into Aave. + ILendingPool(LENDING_POOL).deposit(INVESTMENT_TOKEN, _tokenAmount, address(this), 0); - // Encode shares to return back to QWManager - bytes memory returnData = abi.encode(success, shares); - assembly { - return(add(returnData, 32), mload(returnData)) - } + // Calculate shares to be issued for the new investment. + shares = _tokenAmount / sharePrice; + + success = true; } /** * @notice Executes a transaction on AaveV2 pool to withdraw tokens. * @dev This function is called by the parent contract to withdraw tokens from the AaveV2 pool. - * @param _callData Encoded function call data containing user address, shares amount, and total shares. + * @param _callData Encoded function call data containing total shares. + * @param _sharesAmount Amount of shares to be withdrawn. * @return success boolean indicating the success of the transaction. + * @return tokens Number of tokens to be returned to the user in exchange for shares withdrawn. */ function close( - bytes memory _callData - ) external override onlyQwManager returns (bool success) { - (address _user, uint256 _sharesAmount, uint256 _totalShares, address _tokenAddress) = abi.decode(_callData, (address, uint256, uint256, address)); + bytes memory _callData, + uint256 _sharesAmount + ) external override onlyQwManager returns (bool success, uint256 tokens) { + (uint256 _totalShares) = abi.decode(_callData, (uint256)); if (_sharesAmount > _totalShares) { revert InvalidCallData(); } - uint256 totalSharesValue = getTotalPoolValue(_tokenAddress); - uint256 amountToWithdraw = (_sharesAmount * totalSharesValue) / _totalShares; + // Calculate the amount of tokens to withdraw based on the shares. + uint256 totalInvestmentValue = getInvestmentValue(); + // If shares amount < total shares, then the token amount is share * price per share. + uint256 tokens = (_sharesAmount == _totalShares) ? + totalInvestmentValue + : (_sharesAmount * totalInvestmentValue) / _totalShares; - ILendingPool(LENDING_POOL).withdraw(_tokenAddress, amountToWithdraw, QW_MANAGER); - success = true; + // Withdraw the tokens from Aave. The number of aTokens to withdraw is equal to underlying tokens received. + ILendingPool(LENDING_POOL).withdraw(INVESTMENT_TOKEN, tokens, QW_MANAGER); - // Encode success to return back to QWManager - bytes memory returnData = abi.encode(success); - assembly { - return(add(returnData, 32), mload(returnData)) - } + // TODO: Send tokens back to QWManager. + + success = true; } /** * @notice Gets the price per share in terms of the specified token. * @dev This function calculates the value of one share in terms of the specified token. - * @param _tokenAddress The address of the token to get the price per share in. + * @param _totalShares The total shares. * @return pricePerShare uint256 representing the value of one share in the specified token. */ - function pricePerShare(address _tokenAddress) external view returns (uint256) { - uint256 totalSharesValue = getTotalPoolValue(_tokenAddress); - return totalSharesValue / IERC20(_tokenAddress).totalSupply(); + function pricePerShare(uint256 _totalShares) external view returns (uint256) { + return _totalShares == 0 ? + 1 * 10 ** token.decimals() + : getInvestmentValue() / _totalShares; + } + + /** + * @notice Gets the total investment value in terms of the specified token. + * @dev This function calculates the total value of the investment in terms of the specified token. + * @return investmentValue uint256 representing the total value of the investment in the specified token. + */ + function getInvestmentValue() public view returns (uint256 investmentValue) { + // Get the balance of aTokens, which will reflect the principle investment(s) + interest. + uint256 aTokenBalance = IAToken(A_INVESTMENT_TOKEN).balanceOf(address(this)); + return aTokenBalance; + } + + /** + * @notice Gets the address of the Quant Wealth Manager contract. + * @dev Returns the address of the Quant Wealth Manager contract. + */ + function QW_MANAGER() external view override returns (address) { + return qwManager; } /** - * @notice Gets the total pool value in terms of the specified token. - * @dev This function calculates the total value of the pool in terms of the specified token. - * @param _tokenAddress The address of the token to get the total pool value in. - * @return poolValue uint256 representing the total value of the pool in the specified token. + * @notice Gets the address of the investment token. + * @dev Returns the address of the token that is initially invested and received once the investment is withdrawn. + * @return The address of the investment token. */ - function getTotalPoolValue(address _tokenAddress) public view returns (uint256 poolValue) { - return IERC20(_tokenAddress).balanceOf(address(this)); + function INVESTMENT_TOKEN() external view override returns (address) { + return investmentToken; } } diff --git a/src/interfaces/IQWChild.sol b/src/interfaces/IQWChild.sol index d7f2b78..d460dd0 100644 --- a/src/interfaces/IQWChild.sol +++ b/src/interfaces/IQWChild.sol @@ -10,23 +10,38 @@ interface IQWChild { * @notice Executes a transaction on the child contract. * @dev This function is called by the parent contract to execute a transaction on the child contract. * @param _callData Encoded function call to be executed on the child contract. - * @param _tokenAddress Address of the token to be transferred. - * @param _amount Amount of tokens to be transferred to the child contract. + * @param _tokenAmount Amount of tokens to be transferred to the child contract. * @return success boolean indicating whether the transaction was successful. + * @return shares Number of shares to be allocated to the user in return for investment created. */ - function create(bytes memory _callData, address _tokenAddress, uint256 _amount) external returns (bool success); + function create( + bytes memory _callData, + uint256 _tokenAmount + ) external returns (bool success, uint256 shares); /** * @notice Closes a transaction on the child contract. * @dev This function is called by the parent contract to close a transaction on the child contract. * @param _callData Encoded function call to be executed on the child contract. + * @param _sharesAmount Amount of shares to be withdrawn, will determine tokens withdrawn. * @return success boolean indicating whether the transaction was successfully closed. + * @return tokens Number of tokens to be returned to the user in exchange for shares withdrawn. */ - function close(bytes memory _callData) external returns (bool success); + function close( + bytes memory _callData, + uint256 _sharesAmount + ) external returns (bool success, uint256 tokens); /** * @notice Gets the address of the Quant Wealth Manager contract. - * @dev Returns the address of the Quant Wealth Manager contract. + * @return address The address of the Quant Wealth Manager contract recorded in this child contract. */ function QW_MANAGER() external view returns (address); + + /** + * @notice Gets the address of the target investment token. + * @dev Returns the address of the token that is initially invested and received on withdraw. + * @return address The address of the investment token. + */ + function INVESTMENT_TOKEN() external view returns (address); } diff --git a/src/interfaces/IQWManager.sol b/src/interfaces/IQWManager.sol index cf56723..6c154ec 100644 --- a/src/interfaces/IQWManager.sol +++ b/src/interfaces/IQWManager.sol @@ -7,27 +7,18 @@ pragma solidity 0.8.23; */ interface IQWManager { /** - * @notice Execute a series of investments. + * @notice Execute a series of investments in batches for multiple protocols. * Transfers specified amounts of tokens and calls target contracts with provided calldata. - * @param _targetQwChild List of contract addresses to interact with. - * @param _callData Encoded function calls to be executed on the target contracts. - * @param _tokenAddress Token address to transfer. - * @param _amount Amount of tokens to transfer to each target contract. + * @param batches Array of ExecuteBatch data containing protocol, users, contributions, token, and amount. */ - function execute( - address[] memory _targetQwChild, - bytes[] memory _callData, - address _tokenAddress, - uint256 _amount - ) external; + function execute(ExecuteBatch[] memory batches) external; /** - * @notice Close a series of investments. + * @notice Close a series of investments in batches for multiple protocols. * Calls target contracts with provided calldata to close positions. - * @param _targetQwChild List of contract addresses to interact with. - * @param _callData Encoded function calls to be executed on the target contracts. + * @param batches Array of CloseBatch data containing protocol, users, contributions, token, and shares. */ - function close(address[] memory _targetQwChild, bytes[] memory _callData) external; + function close(CloseBatch[] memory batches) external; /** * @notice Withdraw funds to a specified user. @@ -52,4 +43,32 @@ interface IQWManager { * @return The address of the registry contract. */ function REGISTRY() external view returns (address); + + /** + * @notice ExecuteBatch struct to hold batch data for executing investments. + * @param protocol The protocol into which we are investing funds. + * @param users The users investing. + * @param contributions Contribution fractions in basis points (e.g., 1% = 100, 100% = 10000). + * @param amount The total amount being invested in the given token by all users into this protocol. + */ + struct ExecuteBatch { + address protocol; + address[] users; + uint256[] contributions; + uint256 amount; + } + + /** + * @notice CloseBatch struct to hold batch data for closing investments. + * @param protocol The protocol from which we are withdrawing funds. + * @param users The users withdrawing. + * @param contributions Contribution fractions in basis points (e.g., 1% = 100, 100% = 10000). + * @param shares The total shares being withdrawn from the given protocol by all users. + */ + struct CloseBatch { + address protocol; + address[] users; + uint256[] contributions; + uint256 shares; + } }