Skip to content

Commit

Permalink
feat: shares calc in create; tokens calc in close
Browse files Browse the repository at this point in the history
  • Loading branch information
@jakekidd
jakekidd committed Jun 21, 2024
1 parent 380976a commit 3ee796b
Show file tree
Hide file tree
Showing 4 changed files with 207 additions and 114 deletions.
120 changes: 74 additions & 46 deletions src/contracts/QWManager.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;
}
}
127 changes: 79 additions & 48 deletions src/contracts/child/QWAaveV2.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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;
}
}
Loading

0 comments on commit 3ee796b

Please sign in to comment.