401 lines
15 KiB
Solidity
401 lines
15 KiB
Solidity
// SPDX-License-Identifier: UNLICENSED
|
|
pragma solidity ^0.8.30;
|
|
|
|
import "forge-std/Test.sol";
|
|
import "@abdk/ABDKMath64x64.sol";
|
|
import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
|
|
import "../src/PartyPool.sol";
|
|
|
|
// Import the flash callback interface
|
|
import "../src/IPartyFlashCallback.sol";
|
|
|
|
/// @notice Test contract that implements the flash callback for testing flash loans
|
|
contract FlashBorrower is IPartyFlashCallback {
|
|
enum Action {
|
|
NORMAL, // Normal repayment
|
|
REPAY_NONE, // Don't repay anything
|
|
REPAY_PARTIAL, // Repay less than required
|
|
REPAY_NO_FEE, // Repay only the principal without fee
|
|
REPAY_EXACT, // Repay exactly the required amount
|
|
REPAY_EXTRA // Repay more than required (donation)
|
|
}
|
|
|
|
Action public action;
|
|
address public pool;
|
|
address public recipient;
|
|
address[] public tokens;
|
|
|
|
constructor(address _pool, address[] memory _tokens) {
|
|
pool = _pool;
|
|
tokens = _tokens;
|
|
}
|
|
|
|
function setAction(Action _action, address _recipient) external {
|
|
action = _action;
|
|
recipient = _recipient;
|
|
}
|
|
|
|
function flash(uint256[] memory amounts) external {
|
|
PartyPool(pool).flash(recipient, amounts, "");
|
|
}
|
|
|
|
function partyFlashCallback(
|
|
uint256[] memory loanAmounts,
|
|
uint256[] memory repaymentAmounts,
|
|
bytes calldata /* data */
|
|
) external override {
|
|
require(msg.sender == pool, "Callback not called by pool");
|
|
|
|
if (action == Action.NORMAL || action == Action.REPAY_EXTRA) {
|
|
// Normal or extra repayment - transfer required amounts back to pool
|
|
for (uint256 i = 0; i < loanAmounts.length; i++) {
|
|
if (loanAmounts[i] > 0) {
|
|
uint256 repaymentAmount = repaymentAmounts[i];
|
|
|
|
// For REPAY_EXTRA, add 1 to each repayment
|
|
if (action == Action.REPAY_EXTRA) {
|
|
repaymentAmount += 1;
|
|
}
|
|
|
|
// Transfer from recipient back to pool
|
|
TestERC20(tokens[i]).transferFrom(
|
|
recipient,
|
|
pool,
|
|
repaymentAmount
|
|
);
|
|
}
|
|
}
|
|
} else if (action == Action.REPAY_PARTIAL) {
|
|
// Repay half of the required amounts
|
|
for (uint256 i = 0; i < loanAmounts.length; i++) {
|
|
if (loanAmounts[i] > 0) {
|
|
uint256 partialRepayment = repaymentAmounts[i] / 2;
|
|
TestERC20(tokens[i]).transferFrom(
|
|
recipient,
|
|
pool,
|
|
partialRepayment
|
|
);
|
|
}
|
|
}
|
|
} else if (action == Action.REPAY_NO_FEE) {
|
|
// Repay only the principal without fee
|
|
for (uint256 i = 0; i < loanAmounts.length; i++) {
|
|
if (loanAmounts[i] > 0) {
|
|
TestERC20(tokens[i]).transferFrom(
|
|
recipient,
|
|
pool,
|
|
loanAmounts[i]
|
|
);
|
|
}
|
|
}
|
|
} else if (action == Action.REPAY_EXACT) {
|
|
// Repay exactly what was required
|
|
for (uint256 i = 0; i < loanAmounts.length; i++) {
|
|
if (loanAmounts[i] > 0) {
|
|
TestERC20(tokens[i]).transferFrom(
|
|
recipient,
|
|
pool,
|
|
repaymentAmounts[i]
|
|
);
|
|
}
|
|
}
|
|
}
|
|
// For REPAY_NONE, do nothing (don't repay)
|
|
}
|
|
}
|
|
|
|
/// @notice Minimal ERC20 token for tests with an external mint function.
|
|
contract TestERC20 is ERC20 {
|
|
constructor(string memory name_, string memory symbol_, uint256 initialSupply) ERC20(name_, symbol_) {
|
|
if (initialSupply > 0) {
|
|
_mint(msg.sender, initialSupply);
|
|
}
|
|
}
|
|
|
|
function mint(address to, uint256 amount) external {
|
|
_mint(to, amount);
|
|
}
|
|
|
|
// Expose convenient approve helper for tests (not necessary but handy)
|
|
function approveMax(address spender) external {
|
|
_approve(msg.sender, spender, type(uint256).max);
|
|
}
|
|
}
|
|
|
|
/// @notice Gas testing contract for PartyPool - contains all gas measurement tests
|
|
contract GasTest is Test {
|
|
using ABDKMath64x64 for int128;
|
|
|
|
PartyPool pool2;
|
|
PartyPool pool10;
|
|
PartyPool pool20;
|
|
PartyPool pool50;
|
|
|
|
address alice;
|
|
address bob;
|
|
|
|
// Common parameters
|
|
int128 tradeFrac;
|
|
int128 targetSlippage;
|
|
|
|
uint256 constant INIT_BAL = 1_000_000; // initial token units for each token (internal==amount when base==1)
|
|
uint256 constant BASE = 1; // use base=1 so internal amounts correspond to raw integers (Q64.64 units)
|
|
|
|
/// @notice Helper function to create a pool with the specified number of tokens
|
|
function createPool(uint256 numTokens) internal returns (PartyPool) {
|
|
// Deploy tokens dynamically
|
|
address[] memory tokens = new address[](numTokens);
|
|
uint256[] memory bases = new uint256[](numTokens);
|
|
|
|
for (uint256 i = 0; i < numTokens; i++) {
|
|
string memory name = string(abi.encodePacked("T", vm.toString(i)));
|
|
TestERC20 token = new TestERC20(name, name, 0);
|
|
tokens[i] = address(token);
|
|
bases[i] = BASE;
|
|
|
|
// Mint initial balances for pool initialization and test users
|
|
token.mint(address(this), INIT_BAL);
|
|
token.mint(alice, INIT_BAL);
|
|
token.mint(bob, INIT_BAL);
|
|
}
|
|
|
|
// Deploy pool with a small fee to test fee-handling paths (use 1000 ppm = 0.1%)
|
|
uint256 feePpm = 1000;
|
|
string memory poolName = string(abi.encodePacked("LP", vm.toString(numTokens)));
|
|
PartyPool newPool = new PartyPool(poolName, poolName, tokens, bases, tradeFrac, targetSlippage, feePpm, feePpm);
|
|
|
|
// Transfer initial deposit amounts into pool before initial mint
|
|
for (uint256 i = 0; i < numTokens; i++) {
|
|
TestERC20(tokens[i]).transfer(address(newPool), INIT_BAL);
|
|
}
|
|
|
|
// Perform initial mint (initial deposit); receiver is this contract
|
|
newPool.mint(address(0), address(this), 0, 0);
|
|
|
|
return newPool;
|
|
}
|
|
|
|
function setUp() public {
|
|
alice = address(0xA11ce);
|
|
bob = address(0xB0b);
|
|
|
|
// Configure LMSR parameters similar to other tests: trade size 1% of asset -> 0.01, slippage 0.001
|
|
tradeFrac = ABDKMath64x64.divu(100, 10_000); // 0.01
|
|
targetSlippage = ABDKMath64x64.divu(10, 10_000); // 0.001
|
|
|
|
// Create pools of different sizes
|
|
pool2 = createPool(2);
|
|
pool10 = createPool(10);
|
|
pool20 = createPool(20);
|
|
pool50 = createPool(50);
|
|
}
|
|
|
|
/// @notice Setup a flash borrower for testing
|
|
function setupFlashBorrower() internal returns (FlashBorrower borrower) {
|
|
// Get token addresses from the 2-token pool
|
|
address[] memory tokenAddresses = pool2.allTokens();
|
|
|
|
// Deploy the borrower contract
|
|
borrower = new FlashBorrower(address(pool2), tokenAddresses);
|
|
|
|
// Mint tokens to alice to be used for repayments and approve borrower
|
|
vm.startPrank(alice);
|
|
for (uint256 i = 0; i < tokenAddresses.length; i++) {
|
|
TestERC20(tokenAddresses[i]).mint(alice, INIT_BAL * 2);
|
|
TestERC20(tokenAddresses[i]).approve(address(borrower), type(uint256).max);
|
|
}
|
|
vm.stopPrank();
|
|
}
|
|
|
|
/// @notice Helper function: perform 10 swaps back-and-forth between the first two tokens.
|
|
function _performSwapGasTest(PartyPool testPool) internal {
|
|
address[] memory tokens = testPool.allTokens();
|
|
require(tokens.length >= 2, "Pool must have at least 2 tokens");
|
|
|
|
// Ensure alice approves pool for both tokens
|
|
vm.prank(alice);
|
|
TestERC20(tokens[0]).approve(address(testPool), type(uint256).max);
|
|
vm.prank(alice);
|
|
TestERC20(tokens[1]).approve(address(testPool), type(uint256).max);
|
|
|
|
uint256 maxIn = 1_000;
|
|
|
|
// Perform 10 swaps alternating directions to avoid large imbalance
|
|
for (uint256 i = 0; i < 10; i++) {
|
|
vm.prank(alice);
|
|
if (i % 2 == 0) {
|
|
// swap token0 -> token1
|
|
testPool.swap(alice, alice, 0, 1, maxIn, 0, 0);
|
|
} else {
|
|
// swap token1 -> token0
|
|
testPool.swap(alice, alice, 1, 0, maxIn, 0, 0);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// @notice Gas measurement: perform 10 swaps back-and-forth between first two tokens in the 2-token pool.
|
|
function testSwapGasPair() public {
|
|
_performSwapGasTest(pool2);
|
|
}
|
|
|
|
/// @notice Gas measurement: perform 10 swaps back-and-forth between first two tokens in the 10-token pool.
|
|
function testSwapGasTen() public {
|
|
_performSwapGasTest(pool10);
|
|
}
|
|
|
|
/// @notice Gas measurement: perform 10 swaps back-and-forth between first two tokens in the 20-token pool.
|
|
function testSwapGasTwenty() public {
|
|
_performSwapGasTest(pool20);
|
|
}
|
|
|
|
/// @notice Gas measurement: perform 10 swaps back-and-forth between first two tokens in the 100-token pool.
|
|
function testSwapGasFifty() public {
|
|
_performSwapGasTest(pool50);
|
|
}
|
|
|
|
/// @notice Helper function: alternate swapMint then burnSwap to keep pool size roughly stable.
|
|
function _performSwapMintBurnSwapGasTest(PartyPool testPool) internal {
|
|
uint256 iterations = 10;
|
|
uint256 input = 1_000;
|
|
address[] memory tokens = testPool.allTokens();
|
|
|
|
// Top up alice so repeated operations won't fail
|
|
TestERC20(tokens[0]).mint(alice, iterations * input * 2);
|
|
|
|
vm.startPrank(alice);
|
|
TestERC20(tokens[0]).approve(address(testPool), type(uint256).max);
|
|
|
|
for (uint256 k = 0; k < iterations; k++) {
|
|
// Mint LP by providing single-token input; receive LP minted
|
|
uint256 minted = testPool.swapMint(alice, alice, 0, input, 0);
|
|
// If nothing minted (numerical edge), skip burn step
|
|
if (minted == 0) continue;
|
|
// Immediately burn the minted LP back to tokens, targeting the same token index
|
|
testPool.burnSwap(alice, alice, minted, 0, 0);
|
|
}
|
|
|
|
vm.stopPrank();
|
|
}
|
|
|
|
/// @notice Gas-style test: alternate swapMint then burnSwap on the 2-token pool to keep pool size roughly stable.
|
|
function testSwapMintBurnSwapGasPair() public {
|
|
_performSwapMintBurnSwapGasTest(pool2);
|
|
}
|
|
|
|
/// @notice Gas-style test: alternate swapMint then burnSwap on the 10-token pool to keep pool size roughly stable.
|
|
function testSwapMintBurnSwapGasTen() public {
|
|
_performSwapMintBurnSwapGasTest(pool10);
|
|
}
|
|
|
|
/// @notice Gas-style test: alternate swapMint then burnSwap on the 20-token pool to keep pool size roughly stable.
|
|
function testSwapMintBurnSwapGasTwenty() public {
|
|
_performSwapMintBurnSwapGasTest(pool20);
|
|
}
|
|
|
|
/// @notice Gas-style test: alternate swapMint then burnSwap on the 100-token pool to keep pool size roughly stable.
|
|
function testSwapMintBurnSwapGasFifty() public {
|
|
_performSwapMintBurnSwapGasTest(pool50);
|
|
}
|
|
|
|
/// @notice Helper function: combined gas test (mint then burn) using mint() and burn().
|
|
/// Alternates minting a tiny LP amount and immediately burning the actual minted LP back to avoid net pool depletion.
|
|
function _performMintBurnGasTest(PartyPool testPool) internal {
|
|
uint256 iterations = 50;
|
|
uint256 input = 1_000;
|
|
address[] memory poolTokens = testPool.allTokens();
|
|
|
|
vm.startPrank(alice);
|
|
|
|
// Mint additional tokens to alice and approve pool to transfer tokens for proportional mint
|
|
for (uint256 i = 0; i < poolTokens.length; i++) {
|
|
TestERC20(poolTokens[i]).mint(alice, iterations * input * 2);
|
|
TestERC20(poolTokens[i]).approve(address(testPool), type(uint256).max);
|
|
}
|
|
|
|
for (uint256 k = 0; k < iterations; k++) {
|
|
// Request a tiny LP mint (1 wei) - pool will compute deposits and transfer from alice
|
|
uint256 lpRequest = 1;
|
|
|
|
// Snapshot alice LP before to compute actual minted
|
|
uint256 lpBefore = testPool.balanceOf(alice);
|
|
|
|
// Perform mint; this will transfer underlying from alice into pool
|
|
testPool.mint(alice, alice, lpRequest, 0);
|
|
|
|
uint256 lpAfter = testPool.balanceOf(alice);
|
|
uint256 actualMinted = lpAfter - lpBefore;
|
|
|
|
// If nothing minted due to rounding edge, skip burn
|
|
if (actualMinted == 0) {
|
|
continue;
|
|
}
|
|
|
|
// Burn via plain burn() which will transfer underlying back to alice and burn LP
|
|
testPool.burn(alice, alice, actualMinted, 0);
|
|
}
|
|
|
|
vm.stopPrank();
|
|
}
|
|
|
|
/// @notice Combined gas test (mint then burn) on 2-token pool using mint() and burn().
|
|
/// Alternates minting a tiny LP amount and immediately burning the actual minted LP back to avoid net pool depletion.
|
|
function testMintBurnGasPair() public {
|
|
_performMintBurnGasTest(pool2);
|
|
}
|
|
|
|
/// @notice Combined gas test (mint then burn) on 10-token pool using mint() and burn().
|
|
/// Alternates small mints and burns to keep the pool size roughly stable.
|
|
function testMintBurnGasTen() public {
|
|
_performMintBurnGasTest(pool10);
|
|
}
|
|
|
|
/// @notice Combined gas test (mint then burn) on 20-token pool using mint() and burn().
|
|
/// Alternates small mints and burns to keep the pool size roughly stable.
|
|
function testMintBurnGasTwenty() public {
|
|
_performMintBurnGasTest(pool20);
|
|
}
|
|
|
|
/// @notice Combined gas test (mint then burn) on 100-token pool using mint() and burn().
|
|
/// Alternates small mints and burns to keep the pool size roughly stable.
|
|
function testMintBurnGasFifty() public {
|
|
_performMintBurnGasTest(pool50);
|
|
}
|
|
|
|
/// @notice Gas measurement: flash with single token
|
|
function testFlashGasSingleToken() public {
|
|
FlashBorrower borrower = setupFlashBorrower();
|
|
|
|
// Configure borrower
|
|
borrower.setAction(FlashBorrower.Action.NORMAL, alice);
|
|
|
|
// Create loan request for single token (get array size from pool)
|
|
address[] memory poolTokens = pool2.allTokens();
|
|
uint256[] memory amounts = new uint256[](poolTokens.length);
|
|
amounts[0] = 1000;
|
|
|
|
// Execute flash loan 10 times to measure gas
|
|
for (uint256 i = 0; i < 10; i++) {
|
|
borrower.flash(amounts);
|
|
}
|
|
}
|
|
|
|
/// @notice Gas measurement: flash with multiple tokens
|
|
function testFlashGasMultipleTokens() public {
|
|
FlashBorrower borrower = setupFlashBorrower();
|
|
|
|
// Configure borrower
|
|
borrower.setAction(FlashBorrower.Action.NORMAL, alice);
|
|
|
|
// Create loan request for multiple tokens (get array size from pool)
|
|
address[] memory poolTokens = pool2.allTokens();
|
|
uint256[] memory amounts = new uint256[](poolTokens.length);
|
|
amounts[0] = 1000;
|
|
amounts[1] = 2000;
|
|
|
|
// Execute flash loan 10 times to measure gas
|
|
for (uint256 i = 0; i < 10; i++) {
|
|
borrower.flash(amounts);
|
|
}
|
|
}
|
|
}
|