gas test
This commit is contained in:
tim
400
test/GasTest.sol
Normal file
400
test/GasTest.sol
Normal file
@@ -0,0 +1,400 @@
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user