diff --git a/src/IPartyPool.sol b/src/IPartyPool.sol index 4e8418c..f63059d 100644 --- a/src/IPartyPool.sol +++ b/src/IPartyPool.sol @@ -81,6 +81,13 @@ interface IPartyPool is IERC20Metadata { // Initialization / Mint / Burn (LP token managed) + /// @notice Initial mint to set up pool for the first time. + /// @dev Assumes tokens have already been transferred to the pool prior to calling. + /// Can only be called when the pool is uninitialized (totalSupply() == 0 or lmsr.nAssets == 0). + /// @param receiver address that receives the LP tokens + /// @param lpTokens The number of LP tokens to issue for this mint. If 0, then the number of tokens returned will equal the LMSR internal q total + function initialMint(address receiver, uint256 lpTokens) external returns (uint256 lpMinted); + /// @notice Calculate the proportional deposit amounts required for a given LP token amount /// @dev Returns the minimum token amounts (rounded up) that must be supplied to receive lpTokenAmount /// LP tokens at current pool proportions. If the pool is empty (initial deposit) returns zeros @@ -114,7 +121,7 @@ interface IPartyPool is IERC20Metadata { /// @param receiver address that receives the withdrawn tokens /// @param lpAmount amount of LP tokens to burn (proportional withdrawal) /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore. - function burn(address payer, address receiver, uint256 lpAmount, uint256 deadline) external; + function burn(address payer, address receiver, uint256 lpAmount, uint256 deadline) external returns (uint256[] memory withdrawAmounts); // Swaps diff --git a/src/PartyPool.sol b/src/PartyPool.sol index 7f44752..caf14f6 100644 --- a/src/PartyPool.sol +++ b/src/PartyPool.sol @@ -126,40 +126,6 @@ contract PartyPool is PartyPoolBase, IPartyPool { ---------------------- */ /// @inheritdoc IPartyPool - function mintAmounts(uint256 lpTokenAmount) public view returns (uint256[] memory depositAmounts) { - return _mintAmounts(lpTokenAmount); - } - - function _mintAmounts(uint256 lpTokenAmount) internal view returns (uint256[] memory depositAmounts) { - uint256 n = tokens.length; - depositAmounts = new uint256[](n); - - // If this is the first mint or pool is empty, return zeros - // For first mint, tokens should already be transferred to the pool - if (totalSupply() == 0 || lmsr.nAssets == 0) { - return depositAmounts; // Return zeros, initial deposit handled differently - } - - // Calculate deposit based on current proportions - uint256 totalLpSupply = totalSupply(); - - // lpTokenAmount / totalLpSupply = depositAmount / currentBalance - // Therefore: depositAmount = (lpTokenAmount * currentBalance) / totalLpSupply - // We round up to protect the pool - for (uint i = 0; i < n; i++) { - uint256 currentBalance = cachedUintBalances[i]; - // Calculate with rounding up: (a * b + c - 1) / c - depositAmounts[i] = (lpTokenAmount * currentBalance + totalLpSupply - 1) / totalLpSupply; - } - - return depositAmounts; - } - - /// @notice Initial mint to set up pool for the first time. - /// @dev Assumes tokens have already been transferred to the pool prior to calling. - /// Can only be called when the pool is uninitialized (totalSupply() == 0 or lmsr.nAssets == 0). - /// @param receiver address that receives the LP tokens - /// @param lpTokens The number of LP tokens to issue for this mint. If 0, then the number of tokens returned will equal the LMSR internal q total function initialMint(address receiver, uint256 lpTokens) external nonReentrant returns (uint256 lpMinted) { uint256 n = tokens.length; @@ -195,6 +161,11 @@ contract PartyPool is PartyPoolBase, IPartyPool { emit Mint(address(0), receiver, depositAmounts, lpMinted); } + /// @inheritdoc IPartyPool + function mintAmounts(uint256 lpTokenAmount) public view returns (uint256[] memory depositAmounts) { + return MINT_IMPL.mintAmounts(lpTokenAmount, lmsr.nAssets, totalSupply()); + } + /// @notice Proportional mint for existing pool. /// @dev This function forwards the call to the mint implementation via delegatecall /// @param payer address that provides the input tokens @@ -210,35 +181,13 @@ contract PartyPool is PartyPoolBase, IPartyPool { lpTokenAmount, deadline ); - bytes memory result = Address.functionDelegateCall(address(MINT_IMPL), data); return abi.decode(result, (uint256)); } /// @inheritdoc IPartyPool function burnAmounts(uint256 lpTokenAmount) external view returns (uint256[] memory withdrawAmounts) { - return _burnAmounts(lpTokenAmount); - } - - function _burnAmounts(uint256 lpTokenAmount) internal view returns (uint256[] memory withdrawAmounts) { - uint256 n = tokens.length; - withdrawAmounts = new uint256[](n); - - // If supply is zero or pool uninitialized, return zeros - if (totalSupply() == 0 || lmsr.nAssets == 0) { - return withdrawAmounts; // Return zeros, nothing to withdraw - } - - // Calculate withdrawal amounts based on current proportions - uint256 totalLpSupply = totalSupply(); - - // withdrawAmount = floor(lpTokenAmount * currentBalance / totalLpSupply) - for (uint i = 0; i < n; i++) { - uint256 currentBalance = cachedUintBalances[i]; - withdrawAmounts[i] = (lpTokenAmount * currentBalance) / totalLpSupply; - } - - return withdrawAmounts; + return MINT_IMPL.burnAmounts(lpTokenAmount, lmsr.nAssets, totalSupply(), cachedUintBalances); } /// @notice Burn LP tokens and withdraw the proportional basket to receiver. @@ -247,7 +196,8 @@ contract PartyPool is PartyPoolBase, IPartyPool { /// @param receiver address that receives the withdrawn tokens /// @param lpAmount amount of LP tokens to burn (proportional withdrawal) /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore. - function burn(address payer, address receiver, uint256 lpAmount, uint256 deadline) external nonReentrant { + function burn(address payer, address receiver, uint256 lpAmount, uint256 deadline) external nonReentrant + returns (uint256[] memory withdrawAmounts) { bytes memory data = abi.encodeWithSignature( "burn(address,address,uint256,uint256)", payer, @@ -255,8 +205,8 @@ contract PartyPool is PartyPoolBase, IPartyPool { lpAmount, deadline ); - - Address.functionDelegateCall(address(MINT_IMPL), data); + bytes memory result = Address.functionDelegateCall(address(MINT_IMPL), data); + return abi.decode(result, (uint256[])); } /* ---------------------- diff --git a/src/PartyPoolMintImpl.sol b/src/PartyPoolMintImpl.sol index 3aa7ddc..36dbbcc 100644 --- a/src/PartyPoolMintImpl.sol +++ b/src/PartyPoolMintImpl.sol @@ -43,7 +43,7 @@ contract PartyPoolMintImpl is PartyPoolBase { uint256 oldScaled = ABDKMath64x64.mulu(oldTotal, LP_SCALE); // Calculate required deposit amounts for the desired LP tokens - uint256[] memory depositAmounts = _mintAmounts(lpTokenAmount, lmsr.nAssets, totalSupply()); + uint256[] memory depositAmounts = mintAmounts(lpTokenAmount, lmsr.nAssets, totalSupply()); // Transfer in all token amounts for (uint i = 0; i < n; ) { @@ -101,7 +101,8 @@ contract PartyPoolMintImpl is PartyPoolBase { /// @param receiver address that receives the withdrawn tokens /// @param lpAmount amount of LP tokens to burn (proportional withdrawal) /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore. - function burn(address payer, address receiver, uint256 lpAmount, uint256 deadline) external { + function burn(address payer, address receiver, uint256 lpAmount, uint256 deadline) external + returns (uint256[] memory withdrawAmounts) { require(deadline == 0 || block.timestamp <= deadline, "burn: deadline exceeded"); uint256 n = tokens.length; require(lpAmount > 0, "burn: zero lp"); @@ -119,7 +120,7 @@ contract PartyPoolMintImpl is PartyPoolBase { } // Compute proportional withdrawal amounts for the requested LP amount (rounded down) - uint256[] memory withdrawAmounts = _burnAmounts(lpAmount); + withdrawAmounts = burnAmounts(lpAmount, lmsr.nAssets, totalSupply(), cachedUintBalances); // Transfer underlying tokens out to receiver according to computed proportions for (uint i = 0; i < n; ) { @@ -165,11 +166,8 @@ contract PartyPoolMintImpl is PartyPoolBase { emit Burn(payer, receiver, withdrawAmounts, lpAmount); } - function mintAmounts(uint256 lpTokenAmount, uint256 numAssets, uint256 totalSupply) public view returns (uint256[] memory depositAmounts) { - return _mintAmounts(lpTokenAmount, numAssets, totalSupply); - } - - function _mintAmounts(uint256 lpTokenAmount, uint256 numAssets, uint256 totalSupply) internal view returns (uint256[] memory depositAmounts) { + function mintAmounts(uint256 lpTokenAmount, uint256 numAssets, uint256 totalSupply) public view + returns (uint256[] memory depositAmounts) { depositAmounts = new uint256[](numAssets); // If this is the first mint or pool is empty, return zeros @@ -190,23 +188,20 @@ contract PartyPoolMintImpl is PartyPoolBase { return depositAmounts; } - /// @notice Internal helper to calculate withdrawal amounts for burning LP tokens - function _burnAmounts(uint256 lpTokenAmount) internal view returns (uint256[] memory withdrawAmounts) { - uint256 n = tokens.length; - withdrawAmounts = new uint256[](n); + function burnAmounts(uint256 lpTokenAmount, + uint256 numAssets, uint256 totalSupply, uint256[] memory cachedUintBalances) public view + returns (uint256[] memory withdrawAmounts) { + withdrawAmounts = new uint256[](numAssets); // If supply is zero or pool uninitialized, return zeros - if (totalSupply() == 0 || lmsr.nAssets == 0) { + if (totalSupply == 0 || numAssets == 0) { return withdrawAmounts; // Return zeros, nothing to withdraw } - // Calculate withdrawal amounts based on current proportions - uint256 totalLpSupply = totalSupply(); - // withdrawAmount = floor(lpTokenAmount * currentBalance / totalLpSupply) - for (uint i = 0; i < n; i++) { + for (uint i = 0; i < numAssets; i++) { uint256 currentBalance = cachedUintBalances[i]; - withdrawAmounts[i] = (lpTokenAmount * currentBalance) / totalLpSupply; + withdrawAmounts[i] = (lpTokenAmount * currentBalance) / totalSupply; } return withdrawAmounts;