diff --git a/Cargo.lock b/Cargo.lock index 2ccac0e..4942ca4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2925,6 +2925,7 @@ dependencies = [ "serde", "serde_json", "solana-client", + "solana-commitment-config", "solana-compute-budget", "solana-sdk", "spl-associated-token-account", diff --git a/Cargo.toml b/Cargo.toml index ba60822..ce48708 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ secrecy = "0.10.3" serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.148" solana-client = "=3.0.0" +solana-commitment-config = "3.1.0" solana-compute-budget = "=3.0.0" solana-sdk = "3.0.0" spl-associated-token-account = "8.0.0" diff --git a/README.md b/README.md index 8cfef29..76ae7af 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # pmm-sim -Simulation & Benchmark environment for Solana's Proprietary AMMs. The setup relies on [Litesvm](https://crates.io/crates/litesvm) for local, consistent and expedited execution. Additionally, since some proprietary AMMs block swaps originating from direct offchain calls, we rely on a custom onchain router - [Magnus](https://github.com/limechain/magnus) - to facilitate the swap execution. +Simulation & Benchmark environment for Solana's Proprietary AMMs. The setup relies on [Litesvm](https://crates.io/crates/litesvm) for local, consistent and expedited execution. Additionally, since some proprietary AMMs block swaps originating from direct offchain calls, we rely on a custom router program - [Magnus](https://github.com/limechain/magnus) - to facilitate the swap execution. Supported Prop AMMs: @@ -11,14 +11,14 @@ Supported Prop AMMs: - [x] TesseraV - [x] Goonfi -The swaps can be done either with the local static accounts that can be found at [cfg/accounts](./cfg/accounts) or with the current live accounts (by fetching them on-the-go). By default all swaps & benchmark simulations are done with live accounts. +The swaps can be done either with the local static accounts that can be found at [cfg/accounts](./cfg/accounts) or with the current live accounts (by fetching them on-the-go). By default all swaps & benchmark simulations are done with live accounts. The markets are specified in [setup.toml](./setup.toml). Possible modes of execution include: - **single** - Run a single swap route across one or more Prop AMMs with specified weights. - **multi** - Execute swaps across nested Prop AMM routes. Each inner list represents a single route, each route possibly going through multiple Prop AMMs. - **fetch-accounts** - Fetch accounts for specified PMMs via RPC and save them locally (presumably for later usage). -- **benchmark** - Benchmark swaps for any of the implemented Prop AMMs by specifying, optionally, the accounts, src/dst tokens and step size. Furthermore benchmark data can be visualised with [plot.py](./scripts/plot.py). +- **benchmark** - Benchmark swaps for any of the implemented Prop AMMs by specifying, optionally, the accounts, src/dst tokens and steps size. Benchmark data can be visualised with [plot.py](./scripts/plot.py). Accounts are by default loaded (saved) from (at) [cfg/accounts](./cfg/accounts). Tweaking the source/destination is possible via `--accounts-path` or `ACCOUNTS_PATH` env variable. @@ -38,16 +38,22 @@ _Figure 2: Compute unit usage_ ### Single-route swaps -##### Swap 100 WSOL for USDC using Humidifi. +##### Swap 15K USDC for WSOL using Humidifi. ``` -cargo r -- single --amount-in=100 --pmms=humidifi --weights=100 +cargo r -- single --amount-in=15000 --pmms=humidifi --weights=100 --src-token=USDC --dst-token=WSOL ``` -##### Swap 150,000 USDC for WSOL using Tessera and SolfiV2, in a route, split evenly - 75000 USDC per Prop AMM. +##### Swap 375 WSOL for USDC using Tessera and SolfiV2, in one route, split evenly - 187,5 WSOL per Prop AMM. ``` -cargo r -- single --pmms=tessera,solfi-v2 --weights=50,50 --amount-in=150000 --src-token=USDC --dst-token=WSOL +cargo r -- single --pmms=tessera,solfi-v2 --weights=50,50 --amount-in=375 --src-token=WSOL --dst-token=USDC +``` + +##### Swap 100 WSOL for USDC using SolfiV2, Humidifi, and Tessera, in one route, split 33,33,34 WSOL per Prop AMM. + +``` +cargo r -- single --amount-in=100 --pmms=solfi-v2,humidifi,tessera --weights=33,33,34 --src-token=WSOL --dst-token=USDC --jit-accounts=false ``` ##### Swaps 10,000 USDC for USDT using ObricV2. @@ -61,7 +67,7 @@ cargo r -- single --amount-in=10000 --pmms=obric-v2 --weights=100 --src-token=US ##### Swap 103 WSOL for USDC in a multi-route swap, 100 WSOL via Humidifi and SolfiV2 (split 92%/8%) in one route, and 3 WSOL via SolfiV2 in another route. ``` -cargo r -- multi --pmms="[[humidifi,solfi-v2],[humidifi]]" --weights "[[92, 8],[100]]" --amount-in=100,3 +cargo r -- multi --pmms="[[humidifi,solfi-v2],[humidifi]]" --weights="[[92, 8],[100]]" --amount-in=100,3 ``` ##### Execute two routes, the first swapping 150,000 USDC for WSOL using Humidifi and SolfiV2 (split 25%/75%), the second swapping 1000 USDC for WSOL using Goonfi. Uses the static accounts (i.e the accounts found at [./cfg/accounts](./cfg/accounts)). @@ -72,24 +78,26 @@ RUST_LOG=debug cargo r -- multi --pmms="[[humidifi,solfi-v2],[goonfi]]" --weight ### Benchmark swaps -##### Benchmark swaps on Humidifi, from 1 to 4000 WSOL to USDC, in increments of 1 WSOL, and save the results at [./datasets](./datasets). +##### Benchmark swaps on Humidifi,Tessera,SolfiV2 and Goonfi, from 1 to 4000 WSOL to USDC, in increments of 1 WSOL, and save the results at [./datasets](./datasets). ``` -cargo r -- benchmark --step=1,4000,1 --pmms=humidifi --src-token=wsol --dst-token=usdc +cargo r -- benchmark --steps=1.0,4000.0,1.0 --pmms=humidifi,tessera,solfi-v2,goonfi --src-token=wsol --dst-token=usdc ``` -##### Benchmark swaps (USDC->WSOL) on Humidifi, Tessera, SolfiV2 and Goonfi, from 10K to 100K USDC, in increments of 100 USDC, and save the results at [./datasets](./datasets). +##### Benchmark swaps (USDC->WSOL) on Humidifi and SolfiV2, from 10K to 100K USDC, in increments of 100 USDC, and save the results at [./datasets](./datasets). ``` -cargo r -- benchmark --step=10000,100000,100 --pmms=humidifi,tessera,solfi-v2,goonfi --src-token=usdc --dst-token=wsol +cargo r -- benchmark --steps=10000,100000,100 --pmms=humidifi,solfi-v2 --src-token=usdc --dst-token=wsol ``` -Once generated, the results can be plotted through [./scripts/plot.py](./scripts/plot.py), i.e: +Generated benchmark data can be plotted through [./scripts/plot.py](./scripts/plot.py), like so: ##### Plots all the local datasets for slot `389141713`. ``` + ./scripts/plot.py ./datasets/389141713* + ``` ### Fetch live accounts @@ -97,13 +105,17 @@ Once generated, the results can be plotted through [./scripts/plot.py](./scripts ##### Locally sync the current (live) accounts for all supported Prop AMMs. ``` + cargo r -- fetch-accounts + ``` ##### Locally sync the current (live) accounts for Humidifi and SolfiV2. ``` + cargo r -- fetch-accounts --pmms=humidifi,solfi-v2 + ``` --- @@ -111,6 +123,7 @@ cargo r -- fetch-accounts --pmms=humidifi,solfi-v2 Check out the CLI subcommands for additional clues (i.e `pmm-sim single --help`) ``` + $ pmm-sim --help Simulation environment for Solana's Proprietary AMMs. @@ -128,4 +141,5 @@ Commands: Options: -h, --help Print help -V, --version Print version + ``` diff --git a/setup.toml b/setup.toml index 304fc3d..96d0c25 100644 --- a/setup.toml +++ b/setup.toml @@ -1,44 +1,44 @@ # https://solscan.io/account/FksffEqnBRixYGR791Qw2MgdU7zNCpHVFYBL4Fa4qVuH [humidifi] market = "FksffEqnBRixYGR791Qw2MgdU7zNCpHVFYBL4Fa4qVuH" -base_token_acc = "C3FzbX9n1YD2dow2dCmEv5uNyyf22Gb3TLAEqGBhw5fY" -quote_token_acc = "3RWFAQBRkNGq7CMGcTLK3kXDgFTe9jgMeFYqk8nHwcWh" +base_ta = "C3FzbX9n1YD2dow2dCmEv5uNyyf22Gb3TLAEqGBhw5fY" +quote_ta = "3RWFAQBRkNGq7CMGcTLK3kXDgFTe9jgMeFYqk8nHwcWh" ## other humidifi markets # https://solscan.io/account/DB3sUCP2H4icbeKmK6yb6nUxU5ogbcRHtGuq7W2RoRwW # market = "DB3sUCP2H4icbeKmK6yb6nUxU5ogbcRHtGuq7W2RoRwW" -# base_token_acc = "8BrVfsvzb1DZqCactbYWoKSv24AfsLBuXJqzpzYCwznF" -# quote_token_acc = "HsQcHFFNUVTp3MWrXYbuZchBNd4Pwk8636bKzLvpfYNR" +# base_ta = "8BrVfsvzb1DZqCactbYWoKSv24AfsLBuXJqzpzYCwznF" +# quote_ta = "HsQcHFFNUVTp3MWrXYbuZchBNd4Pwk8636bKzLvpfYNR" # https://solscan.io/account/6n9VhCwQ7EwK6NqFDjnHPzEk6wZdRBTfh43RFgHQWHuQ # market = "6n9VhCwQ7EwK6NqFDjnHPzEk6wZdRBTfh43RFgHQWHuQ" -# base_token_acc = "C3FzbX9n1YD2dow2dCmEv5uNyyf22Gb3TLAEqGBhw5fY" -# quote_token_acc = "3RWFAQBRkNGq7CMGcTLK3kXDgFTe9jgMeFYqk8nHwcWh" +# base_ta = "C3FzbX9n1YD2dow2dCmEv5uNyyf22Gb3TLAEqGBhw5fY" +# quote_ta = "3RWFAQBRkNGq7CMGcTLK3kXDgFTe9jgMeFYqk8nHwcWh" # market = "AvGeFw71N5sNfV97mZ1uNrHg4yfufRicCJUrS9j2ehTX" -# base_token_acc = "ECEPWwZJ1U1Vjsj1X5sUbZYETKMSCjYHuoTMVitCn64t" -# quote_token_acc = "FBWtVVvzsRuAAzVX8ua1hden9KmgPrC2rFijuwEn1ngJ" +# base_ta = "ECEPWwZJ1U1Vjsj1X5sUbZYETKMSCjYHuoTMVitCn64t" +# quote_ta = "FBWtVVvzsRuAAzVX8ua1hden9KmgPrC2rFijuwEn1ngJ" # https://solscan.io/account/FLckHLGMJy5gEoXWwcE68Nprde1D4araK4TGLw4pQq2n [tessera] market = "FLckHLGMJy5gEoXWwcE68Nprde1D4araK4TGLw4pQq2n" -base_token_acc = "5pVN5XZB8cYBjNLFrsBCPWkCQBan5K5Mq2dWGzwPgGJV" -quote_token_acc = "9t4P5wMwfFkyn92Z7hf463qYKEZf8ERVZsGBEPNp8uJx" +base_ta = "5pVN5XZB8cYBjNLFrsBCPWkCQBan5K5Mq2dWGzwPgGJV" +quote_ta = "9t4P5wMwfFkyn92Z7hf463qYKEZf8ERVZsGBEPNp8uJx" global_state = "8ekCy2jHHUbW2yeNGFWYJT9Hm9FW7SvZcZK66dSZCDiF" # https://solscan.io/account/4uWuh9fC7rrZKrN8ZdJf69MN1e2S7FPpMqcsyY1aof6K [goonfi] market = "4uWuh9fC7rrZKrN8ZdJf69MN1e2S7FPpMqcsyY1aof6K" -base_token_acc = "pKiUC9hDXv52xqU1p3BKypV9AQjAMgfZUGRnoBsdkKm" -quote_token_acc = "Gsy5Zr7Vxn5KckAbduPHHGR1qzPJ4w3GSYmcinWAkhrC" +base_ta = "pKiUC9hDXv52xqU1p3BKypV9AQjAMgfZUGRnoBsdkKm" +quote_ta = "Gsy5Zr7Vxn5KckAbduPHHGR1qzPJ4w3GSYmcinWAkhrC" blacklist = "7XqYD6DEGmDXooB1E8NNRWV9pWAmm1z6WYpsfjnABTUz" # https://solscan.io/account/65ZHSArs5XxPseKQbB1B4r16vDxMWnCxHMzogDAqiDUc [solfi-v2] market = "65ZHSArs5XxPseKQbB1B4r16vDxMWnCxHMzogDAqiDUc" -base_token_acc = "CRo8DBwrmd97DJfAnvCv96tZPL5Mktf2NZy2ZnhDer1A" -quote_token_acc = "GhFfLFSprPpfoRaWakPMmJTMJBHuz6C694jYwxy2dAic" +base_ta = "CRo8DBwrmd97DJfAnvCv96tZPL5Mktf2NZy2ZnhDer1A" +quote_ta = "GhFfLFSprPpfoRaWakPMmJTMJBHuz6C694jYwxy2dAic" cfg = "FmxXDSR9WvpJTCh738D1LEDuhMoA8geCtZgHb3isy7Dp" oracle = "2ny7eGyZCoeEVTkNLf5HcnJFBKkyA4p4gcrtb3b8y8ou" diff --git a/src/main.rs b/src/main.rs index f41e225..1f22483 100644 --- a/src/main.rs +++ b/src/main.rs @@ -27,6 +27,7 @@ use magnus_shared::{Dex, Route}; use secrecy::{ExposeSecret, SecretString}; use serde::{Deserialize, Serialize}; use solana_client::rpc_client::RpcClient; +use solana_commitment_config::CommitmentConfig; use solana_compute_budget::compute_budget::ComputeBudget; use solana_sdk::{ account::Account, message::AccountMeta, program_pack::Pack, pubkey::Pubkey, rent::Rent, signature::Keypair, signer::Signer, sysvar, @@ -37,7 +38,7 @@ use tracing::{debug, info, warn}; use tracing_subscriber::{EnvFilter, fmt::time::UtcTime}; /// Constants used throughout the simulation environment. -/// Holds the CFG file paths and swappable token accounts; +/// Holds the CFG file paths, swappable token accounts and more; pub mod consts { use solana_sdk::{pubkey, pubkey::Pubkey}; @@ -57,6 +58,12 @@ pub mod consts { pub const USDT_DECIMALS: u8 = 6; pub const PROGRESS_TEMPLATE: &str = "{prefix:>12.bold} [{bar:40.cyan/blue}] {pos:>6}/{len:<6} ({percent}%)"; + pub const PROGRESS_CHARS: &str = "█▓░"; + + // used to pay for tx fees + pub const AIRDROP_AMOUNT: u64 = 100_000_000_000; + // the maximum number of compute units a tx can consume + pub const COMPUTE_UNITS_LIMIT: u64 = 20_000_000; } /// Macro to generate dex configuration structs and their associated functions. @@ -144,25 +151,25 @@ macro_rules! define_dex_configs { define_dex_configs! { Humidifi => HumidifiCfg : humidifi ("humidifi") { market, - base_token_acc, - quote_token_acc, + base_ta, + quote_ta, }, Tessera => TesseraCfg : tessera ("tessera") { market, - base_token_acc, - quote_token_acc, + base_ta, + quote_ta, global_state, }, Goonfi => GoonfiCfg : goonfi ("goonfi") { market, - base_token_acc, - quote_token_acc, + base_ta, + quote_ta, blacklist, }, SolfiV2 => SolfiV2Cfg : solfi_v2 ("solfi-v2") { market, - base_token_acc, - quote_token_acc, + base_ta, + quote_ta, cfg, oracle, }, @@ -217,7 +224,7 @@ impl CliArgs { serde_json::from_str(s).map_err(|e| format!("invalid format: {}", e)) } - fn parse_step(s: &str) -> Result<[f64; 3], String> { + fn parse_steps(s: &str) -> Result<[f64; 3], String> { let parts: Vec = s.split(',').map(|p| p.trim().parse::().map_err(|e| e.to_string())).collect::, _>>()?; let parts: [f64; 3] = @@ -364,11 +371,11 @@ pub enum Command { after_help = "Examples: # Benchmark Humidifi swaps (WSOL->USDC) with the current AMM state, stepping from 1 to 100 with a step size of 1. The resulting CSV # will be saved in the ./datasets directory - pmm-sim benchmark --pmms=humidifi --step=1.0,100.0,1.0 + pmm-sim benchmark --pmms=humidifi --steps=1.0,100.0,1.0 - # Benchmark SolfiV2 and Tessera swaps (USDC->USDT) with the current AMM state, stepping from 10 to 1000 with a step size of 5. The + # Benchmark SolfiV2 and Tessera swaps (WSOL->USDC) with the current AMM state, stepping from 10 to 1000 with a step size of 5. The # resulting CSVs will be saved in the ./datasets directory - pmm-sim benchmark --pmms=solfi-v2,tessera --src-token=USDC --dst-token=USDT --step=10.0,1000.0,5.0 + pmm-sim benchmark --pmms=solfi-v2,tessera --src-token=WSOL --dst-token=USDC --steps=10.0,1000.0,5.0 " )] Benchmark { @@ -381,8 +388,8 @@ pub enum Command { #[arg(long, env = "PROP_AMMS", value_delimiter = ',', default_value = "humidifi", help = "The Prop AMMs to benchmark")] pmms: Vec, - #[arg(long, env = "STEP", default_value = "1.0,100.0,1.0", value_parser = CliArgs::parse_step, help = "Comma-separated step parameters: start, end, step")] - step: [f64; 3], + #[arg(long, env = "STEPS", default_value = "1.0,100.0,1.0", value_parser = CliArgs::parse_steps, help = "Comma-separated step parameters: start, end, step")] + steps: [f64; 3], }, } @@ -490,7 +497,7 @@ impl<'a, P: Into + Display + Clone + Debug> Environment<'a, P> { slot: Option, ) -> eyre::Result> { let mut budget = ComputeBudget::new_with_defaults(false); - budget.compute_unit_limit = 20_000_000; + budget.compute_unit_limit = consts::COMPUTE_UNITS_LIMIT; let wallet = Keypair::new(); let mut svm = LiteSVM::new().with_default_programs().with_sysvars().with_sigverify(true).with_compute_budget(budget); @@ -501,36 +508,52 @@ impl<'a, P: Into + Display + Clone + Debug> Environment<'a, P> { } } - Ok(Environment { svm, slot, wallet, programs_path, accounts_path, mints, cfg }) - } - - fn setup_wallet(&mut self, mint: &Pubkey, mint_amount: u64, airdrop_amount: u64) -> eyre::Result<()> { - // create the ATAs for all initialised mints - if let Some(mints) = self.mints { - for (mint, _) in mints { - let ata = get_associated_token_address(&self.wallet_pubkey(), mint); - self.svm.set_account(ata, self.mk_ata(mint, &self.wallet_pubkey(), 0))?; - } + if let Some(slot) = slot { + svm.warp_to_slot(slot); } - let ata = get_associated_token_address(&self.wallet_pubkey(), mint); - self.svm.set_account(ata, self.mk_ata(mint, &self.wallet_pubkey(), mint_amount))?; + Ok(Environment { svm, slot, wallet, programs_path, accounts_path, mints, cfg }) + } + /// Sets up the wallet for the simulation environment. + /// + /// This function initializes the wallet's Associated Token Accounts (ATAs) for all + /// configured mints and funds the wallet with SOL for transaction fees. + /// + /// # Arguments + /// * `src_mint` - The mint address of the source token to fund + /// * `src_amount` - The amount of source tokens to mint to the wallet's ATA + /// * `airdrop_amount` - The amount of SOL (in lamports) to airdrop for transaction fees + /// + /// # Behavior + /// 1. Creates ATAs with zero balance for all mints in `self.mints` + /// 2. Sets the source mint's ATA balance to `src_amount` + /// 3. Airdrops SOL to the wallet for fees + fn setup_wallet(&mut self, src_mint: &Pubkey, src_amount: u64, airdrop_amount: u64) -> eyre::Result<()> { + self.reset_wallet(src_mint, src_amount)?; self.svm.airdrop(&self.wallet_pubkey(), airdrop_amount).expect("airdrop failed"); Ok(()) } - fn reset_wallet(&mut self, mint: &Pubkey, amount: u64) -> eyre::Result<()> { - let src_ata = self.wallet_ata(mint); - self.svm.set_account(src_ata, self.mk_ata(mint, &self.wallet_pubkey(), amount))?; - + /// Resets the wallet's token balances between simulation iterations. + /// + /// This function is used in benchmarking to restore the wallet to a known state + /// before each swap iteration, ensuring consistent and reproducible results. + /// + /// # Arguments + /// * `src_mint` - The mint address of the source token + /// * `src_amount` - The amount of source tokens to set in the wallet's ATA + /// + /// # Behavior + /// 1. Sets the source mint's ATA balance to `src_amount` + /// 2. Resets all other mint ATAs to zero balance + fn reset_wallet(&mut self, src_mint: &Pubkey, src_amount: u64) -> eyre::Result<()> { if let Some(mints) = self.mints { - for (m, _) in mints { - if m != mint { - let dst_ata = self.wallet_ata(m); - self.svm.set_account(dst_ata, self.mk_ata(m, &self.wallet_pubkey(), 0))?; - } + for (mint, _) in mints { + let ata = self.wallet_ata(mint); + let amount = if mint == src_mint { src_amount } else { 0 }; + self.svm.set_account(ata, Misc::mk_ata(mint, &self.wallet_pubkey(), amount))?; } } @@ -545,10 +568,13 @@ impl<'a, P: Into + Display + Clone + Debug> Environment<'a, P> { get_associated_token_address(&self.wallet.pubkey(), mint) } + /// Loads the router program and all required PMM programs into the SVM. + /// + /// Programs are loaded from `.so` files in the configured `programs_path` directory. + /// The router program is always loaded, plus any unique PMM programs from the provided list. fn load_programs(&mut self, pmms: &[Dex]) -> eyre::Result<()> { // mandatory load - self.svm - .add_program_from_file(magnus_router_client::programs::ROUTER_ID, format!("{}/{}.so", self.programs_path, consts::ROUTER))?; + self.load_program_router()?; let unique_pmms: HashSet<_> = pmms.iter().collect(); for dex in unique_pmms { @@ -557,11 +583,20 @@ impl<'a, P: Into + Display + Clone + Debug> Environment<'a, P> { self.svm.add_program_from_file(Pubkey::new_from_array(program_id.to_bytes()), format!("{}/{}.so", self.programs_path, dex))?; } - info!("loaded {pmms:?} programs"); + info!("loaded {pmms:?} program(s)"); + + Ok(()) + } + + /// Loads the router program into the SVM. + fn load_program_router(&mut self) -> eyre::Result<()> { + self.svm + .add_program_from_file(magnus_router_client::programs::ROUTER_ID, format!("{}/{}.so", self.programs_path, consts::ROUTER))?; Ok(()) } + /// Sets multiple accounts in the SVM state. fn load_accounts(&mut self, accs: &Vec<(Pubkey, Account)>) -> eyre::Result<()> { for (pubkey, acc) in accs { self.svm.set_account(*pubkey, acc.clone())?; @@ -570,6 +605,7 @@ impl<'a, P: Into + Display + Clone + Debug> Environment<'a, P> { Ok(()) } + /// Fetches and loads PMM accounts either from RPC (JIT) or from disk cache. fn fetch_and_load_accounts(&mut self, pmms: &[Dex], jit: bool, client: Option<&RpcClient>) -> eyre::Result<()> { match jit { true => { @@ -584,14 +620,12 @@ impl<'a, P: Into + Display + Clone + Debug> Environment<'a, P> { Ok(()) } + /// Loads PMM accounts from disk cache and warps to the cached slot. fn static_accounts(&mut self, pmms: &[Dex]) -> eyre::Result<()> { let (slot, accs_map) = Misc::read_accounts_disk(pmms, &self.accounts_path.to_string())?; - for (dex, accounts) in accs_map { - for (pubkey, account) in accounts { - self.svm.set_account(pubkey, account)?; - debug!("loaded account {pubkey} for {dex}"); - } + for (_, accs) in accs_map { + self.load_accounts(&accs)?; } if let Some(s) = slot { @@ -602,58 +636,33 @@ impl<'a, P: Into + Display + Clone + Debug> Environment<'a, P> { Ok(()) } + /// Fetches PMM accounts from RPC and warps to the fetched slot. fn jit_accounts(&mut self, pmms: &[Dex], client: &RpcClient) -> eyre::Result<()> { let (slot, fetched) = Misc::fetch_pmm_accounts(pmms, client, &self.cfg)?; - for (dex, accounts) in fetched { - for (pubkey, account) in accounts { - self.svm.set_account(pubkey, account)?; - debug!("loaded account {pubkey} for {dex}"); - } + for (_, accs) in fetched { + self.load_accounts(&accs)?; } - info!("loaded {pmms:?} accounts"); self.svm.warp_to_slot(slot); self.slot = Some(slot); Ok(()) } - fn save_account_to_disk(&self, dex: &Dex, pubkey: &Pubkey, account: &Account, slot: u64) -> eyre::Result<()> { - let filename = format!("{}_{}.json", dex, pubkey); - let accounts_path = format!("{}", self.accounts_path); - let data_dir = Path::new(&accounts_path); - - if !data_dir.exists() { - fs::create_dir_all(data_dir)?; - } - - let file_path = data_dir.join(filename); - - let value = serde_json::json!({ - "pubkey": pubkey.to_string(), - "slot": slot, - "account": { - "lamports": account.lamports, - "data": [general_purpose::STANDARD.encode(&account.data), "base64"], - "owner": account.owner.to_string(), - "executable": account.executable, - "rentEpoch": account.rent_epoch, - } - }); - - let mut file = File::create(file_path)?; - file.write_all(serde_json::to_string_pretty(&value)?.as_bytes())?; - - Ok(()) - } - + /// Returns the token balance for the wallet's ATA of the given mint. fn token_balance(&self, mint: &Pubkey) -> u64 { let ata = self.wallet_ata(mint); let account = self.svm.get_account(&ata).unwrap_or_default(); spl_token::state::Account::unpack(&account.data).map(|a| a.amount).unwrap_or(0) } + fn token_balance_norm(&self, mint: &Pubkey, decimals: u8) -> f64 { + let balance = self.token_balance(mint); + + balance as f64 / 10_f64.powi(decimals as i32) + } + fn latest_blockhash(&self) -> solana_sdk::hash::Hash { self.svm.latest_blockhash() } @@ -662,13 +671,17 @@ impl<'a, P: Into + Display + Clone + Debug> Environment<'a, P> { self.svm.send_transaction(tx) } + /// Extracts the output amount from a swap transaction's logs. + /// + /// Parses the `SwapEvent` log emitted by the router program to find the + /// `amount_out` value. Panics if the event is not found in the logs. fn get_event_amount_out(&self, metadata: &TransactionMetadata) -> u64 { let amount_out: u64 = metadata .logs .iter() .find_map(|log| { if log.contains("SwapEvent") { - // Log format: "Program log: SwapEvent { dex: Humidifi, amount_in: 1000000000, amount_out: 121518066 }" + // i.e.: "Program log: SwapEvent { dex: Humidifi, amount_in: 1000000000, amount_out: 121518066 }" log.split("amount_out: ").nth(1)?.split(|c: char| !c.is_ascii_digit()).next()?.parse().ok() } else { None @@ -678,27 +691,6 @@ impl<'a, P: Into + Display + Clone + Debug> Environment<'a, P> { amount_out } - - fn mk_ata(&self, mint: &Pubkey, user: &Pubkey, amount: u64) -> Account { - let ata = spl_token::state::Account { - mint: *mint, - owner: *user, - amount, - state: spl_token::state::AccountState::Initialized, - ..Default::default() - }; - - let mut data = vec![0u8; spl_token::state::Account::LEN]; - ata.pack_into_slice(&mut data); - - Account { - lamports: Rent::default().minimum_balance(data.len()), - data, - owner: spl_token::id(), - executable: false, - rent_epoch: u64::MAX, - } - } } /// A helper struct to construct swap instructions with the required accounts @@ -708,12 +700,14 @@ impl<'a, P: Into + Display + Clone + Debug> Environment<'a, P> { /// which in turn calls the respective Prop AMM program. Therefore, the swap /// instruction is built using the `SwapBuilder` from the `magnus-router-client` /// crate, and then the required accounts for the specific Prop AMM are attached. +/// +/// The order of the remaining_accounts matters. pub struct ConstructSwap<'a> { cfg: PMMCfg, builder: &'a mut SwapBuilder, payer: Pubkey, - sta: Pubkey, - dta: Pubkey, + src_ta: Pubkey, + dst_ta: Pubkey, src_mint: Pubkey, dst_mint: Pubkey, } @@ -723,6 +717,11 @@ impl<'a> ConstructSwap<'a> { self.builder.instruction() } + /// Attaches the required remaining accounts for the specified PMM to the swap instruction. + /// + /// Each Prop AMM program expects a specific set of accounts in a precise order as + /// "remaining accounts" on the swap instruction. This method dispatches to the + /// appropriate PMM-specific attachment function based on the DEX type. fn attach_pmm_accs(&mut self, pmm: &Dex) { match pmm { Dex::Humidifi => self.attach_humidifi_accs(), @@ -737,140 +736,133 @@ impl<'a> ConstructSwap<'a> { }; } - pub fn attach_solfiv2_accs(&mut self) { + fn attach_solfiv2_accs(&mut self) { if let Some(cfg) = &self.cfg.solfi_v2 { - self.builder - .add_remaining_account(AccountMeta::new_readonly( - Pubkey::new_from_array(magnus_shared::pmm_solfi_v2::id().to_bytes()), - false, - )) - .add_remaining_account(AccountMeta::new(self.payer, true)) - .add_remaining_account(AccountMeta::new(self.sta, false)) - .add_remaining_account(AccountMeta::new(self.dta, false)) - .add_remaining_account(AccountMeta::new(cfg.market, false)) - .add_remaining_account(AccountMeta::new_readonly(cfg.oracle, false)) - .add_remaining_account(AccountMeta::new_readonly(cfg.cfg, false)) - .add_remaining_account(AccountMeta::new(cfg.base_token_acc, false)) - .add_remaining_account(AccountMeta::new(cfg.quote_token_acc, false)) - .add_remaining_account(AccountMeta::new_readonly(consts::WSOL, false)) - .add_remaining_account(AccountMeta::new_readonly(consts::USDC, false)) - .add_remaining_account(AccountMeta::new_readonly(spl_token::id(), false)) - .add_remaining_account(AccountMeta::new_readonly(spl_token::id(), false)) - .add_remaining_account(AccountMeta::new_readonly(sysvar::instructions::id(), false)); + self.builder.add_remaining_accounts(&[ + AccountMeta::new_readonly(Pubkey::new_from_array(magnus_shared::pmm_solfi_v2::id().to_bytes()), false), + AccountMeta::new(self.payer, true), + AccountMeta::new(self.src_ta, false), + AccountMeta::new(self.dst_ta, false), + AccountMeta::new(cfg.market, false), + AccountMeta::new_readonly(cfg.oracle, false), + AccountMeta::new_readonly(cfg.cfg, false), + AccountMeta::new(cfg.base_ta, false), + AccountMeta::new(cfg.quote_ta, false), + AccountMeta::new_readonly(consts::WSOL, false), + AccountMeta::new_readonly(consts::USDC, false), + AccountMeta::new_readonly(spl_token::id(), false), + AccountMeta::new_readonly(spl_token::id(), false), + AccountMeta::new_readonly(sysvar::instructions::id(), false), + ]); } else { panic!("SolfiV2 config is missing, cannot attach accounts."); } } - pub fn attach_humidifi_accs(&mut self) { - if let Some(cfg) = &self.cfg.humidifi { - self.builder - .add_remaining_account(AccountMeta::new_readonly( - Pubkey::new_from_array(magnus_shared::pmm_humidifi::id().to_bytes()), - false, - )) - .add_remaining_account(AccountMeta::new(self.payer, true)) - .add_remaining_account(AccountMeta::new(self.sta, false)) - .add_remaining_account(AccountMeta::new(self.dta, false)) - .add_remaining_account(AccountMeta::new_readonly(Misc::create_humidifi_param(1500), false)) - .add_remaining_account(AccountMeta::new(cfg.market, false)) - .add_remaining_account(AccountMeta::new(cfg.base_token_acc, false)) - .add_remaining_account(AccountMeta::new(cfg.quote_token_acc, false)) - .add_remaining_account(AccountMeta::new_readonly(sysvar::clock::id(), false)) - .add_remaining_account(AccountMeta::new_readonly(spl_token::id(), false)) - .add_remaining_account(AccountMeta::new_readonly(sysvar::instructions::id(), false)); - } else { + fn attach_humidifi_accs(&mut self) { + let Some(cfg) = &self.cfg.humidifi else { panic!("Humidifi config is missing, cannot attach accounts."); - } - } + }; - pub fn attach_zerofi_accs(&mut self) { - if let Some(cfg) = &self.cfg.zerofi { - self.builder - .add_remaining_account(AccountMeta::new_readonly(Pubkey::new_from_array(magnus_shared::pmm_zerofi::id().to_bytes()), false)) - .add_remaining_account(AccountMeta::new(self.payer, true)) - .add_remaining_account(AccountMeta::new(self.sta, false)) - .add_remaining_account(AccountMeta::new(self.dta, false)) - .add_remaining_account(AccountMeta::new(cfg.market, false)) - .add_remaining_account(AccountMeta::new(cfg.vault_info_base, false)) - .add_remaining_account(AccountMeta::new(cfg.vault_base, false)) - .add_remaining_account(AccountMeta::new(cfg.vault_info_quote, false)) - .add_remaining_account(AccountMeta::new(cfg.vault_quote, false)) - .add_remaining_account(AccountMeta::new_readonly(spl_token::id(), false)) - .add_remaining_account(AccountMeta::new_readonly(sysvar::instructions::id(), false)); - } else { + self.builder.add_remaining_accounts(&[ + AccountMeta::new_readonly(Pubkey::new_from_array(magnus_shared::pmm_humidifi::id().to_bytes()), false), + AccountMeta::new(self.payer, true), + AccountMeta::new(self.src_ta, false), + AccountMeta::new(self.dst_ta, false), + AccountMeta::new_readonly(Misc::create_humidifi_param(1500), false), + AccountMeta::new(cfg.market, false), + AccountMeta::new(cfg.base_ta, false), + AccountMeta::new(cfg.quote_ta, false), + AccountMeta::new_readonly(sysvar::clock::id(), false), + AccountMeta::new_readonly(spl_token::id(), false), + AccountMeta::new_readonly(sysvar::instructions::id(), false), + ]); + } + + fn attach_zerofi_accs(&mut self) { + let Some(cfg) = &self.cfg.zerofi else { panic!("Zerofi config is missing, cannot attach accounts."); - } - } + }; - pub fn attach_obric_v2_accs(&mut self) { - if let Some(cfg) = &self.cfg.obric_v2 { - self.builder - .add_remaining_account(AccountMeta::new_readonly( - Pubkey::new_from_array(magnus_shared::pmm_obric_v2::id().to_bytes()), - false, - )) - .add_remaining_account(AccountMeta::new(self.payer, true)) - .add_remaining_account(AccountMeta::new(self.sta, false)) - .add_remaining_account(AccountMeta::new(self.dta, false)) - .add_remaining_account(AccountMeta::new(cfg.market, false)) - .add_remaining_account(AccountMeta::new_readonly(cfg.second_ref_oracle, false)) - .add_remaining_account(AccountMeta::new_readonly(cfg.third_ref_oracle, false)) - .add_remaining_account(AccountMeta::new(cfg.reserve_x, false)) - .add_remaining_account(AccountMeta::new(cfg.reserve_y, false)) - .add_remaining_account(AccountMeta::new(cfg.ref_oracle, false)) - .add_remaining_account(AccountMeta::new_readonly(cfg.x_price_feed, false)) - .add_remaining_account(AccountMeta::new_readonly(cfg.y_price_feed, false)) - .add_remaining_account(AccountMeta::new_readonly(spl_token::id(), false)); - } else { + self.builder.add_remaining_accounts(&[ + AccountMeta::new_readonly(Pubkey::new_from_array(magnus_shared::pmm_zerofi::id().to_bytes()), false), + AccountMeta::new(self.payer, true), + AccountMeta::new(self.src_ta, false), + AccountMeta::new(self.dst_ta, false), + AccountMeta::new(cfg.market, false), + AccountMeta::new(cfg.vault_info_base, false), + AccountMeta::new(cfg.vault_base, false), + AccountMeta::new(cfg.vault_info_quote, false), + AccountMeta::new(cfg.vault_quote, false), + AccountMeta::new_readonly(spl_token::id(), false), + AccountMeta::new_readonly(sysvar::instructions::id(), false), + ]); + } + + fn attach_obric_v2_accs(&mut self) { + let Some(cfg) = &self.cfg.obric_v2 else { panic!("ObricV2 config is missing, cannot attach accounts."); - } - } + }; - pub fn attach_tessera_accs(&mut self) { - if let Some(cfg) = &self.cfg.tessera { - self.builder - .add_remaining_account(AccountMeta::new_readonly( - Pubkey::new_from_array(magnus_shared::pmm_tessera::id().to_bytes()), - false, - )) - .add_remaining_account(AccountMeta::new(self.payer, true)) - .add_remaining_account(AccountMeta::new(self.sta, false)) - .add_remaining_account(AccountMeta::new(self.dta, false)) - .add_remaining_account(AccountMeta::new_readonly(cfg.global_state, false)) - .add_remaining_account(AccountMeta::new(cfg.market, false)) - .add_remaining_account(AccountMeta::new(cfg.base_token_acc, false)) - .add_remaining_account(AccountMeta::new(cfg.quote_token_acc, false)) - .add_remaining_account(AccountMeta::new_readonly(self.src_mint, false)) - .add_remaining_account(AccountMeta::new_readonly(self.dst_mint, false)) - .add_remaining_account(AccountMeta::new_readonly(spl_token::id(), false)) - .add_remaining_account(AccountMeta::new_readonly(spl_token::id(), false)) - .add_remaining_account(AccountMeta::new_readonly(sysvar::instructions::id(), false)); - } else { + self.builder.add_remaining_accounts(&[ + AccountMeta::new_readonly(Pubkey::new_from_array(magnus_shared::pmm_obric_v2::id().to_bytes()), false), + AccountMeta::new(self.payer, true), + AccountMeta::new(self.src_ta, false), + AccountMeta::new(self.dst_ta, false), + AccountMeta::new(cfg.market, false), + AccountMeta::new_readonly(cfg.second_ref_oracle, false), + AccountMeta::new_readonly(cfg.third_ref_oracle, false), + AccountMeta::new(cfg.reserve_x, false), + AccountMeta::new(cfg.reserve_y, false), + AccountMeta::new(cfg.ref_oracle, false), + AccountMeta::new_readonly(cfg.x_price_feed, false), + AccountMeta::new_readonly(cfg.y_price_feed, false), + AccountMeta::new_readonly(spl_token::id(), false), + ]); + } + + fn attach_tessera_accs(&mut self) { + let Some(cfg) = &self.cfg.tessera else { panic!("Tessera config is missing, cannot attach accounts."); - } - } + }; - pub fn attach_goonfi_accs(&mut self) { - if let Some(cfg) = &self.cfg.goonfi { - let goonfi_param_bytes = [0u8; 32]; - let goonfi_param = Pubkey::new_from_array(goonfi_param_bytes); - - self.builder - .add_remaining_account(AccountMeta::new_readonly(Pubkey::new_from_array(magnus_shared::pmm_goonfi::id().to_bytes()), false)) - .add_remaining_account(AccountMeta::new(self.payer, true)) - .add_remaining_account(AccountMeta::new(self.sta, false)) - .add_remaining_account(AccountMeta::new(self.dta, false)) - .add_remaining_account(AccountMeta::new_readonly(goonfi_param, false)) - .add_remaining_account(AccountMeta::new(cfg.market, false)) - .add_remaining_account(AccountMeta::new(cfg.base_token_acc, false)) - .add_remaining_account(AccountMeta::new(cfg.quote_token_acc, false)) - .add_remaining_account(AccountMeta::new_readonly(cfg.blacklist, false)) - .add_remaining_account(AccountMeta::new_readonly(sysvar::instructions::id(), false)) - .add_remaining_account(AccountMeta::new_readonly(spl_token::id(), false)); - } else { + self.builder.add_remaining_accounts(&[ + AccountMeta::new_readonly(Pubkey::new_from_array(magnus_shared::pmm_tessera::id().to_bytes()), false), + AccountMeta::new(self.payer, true), + AccountMeta::new(self.src_ta, false), + AccountMeta::new(self.dst_ta, false), + AccountMeta::new_readonly(cfg.global_state, false), + AccountMeta::new(cfg.market, false), + AccountMeta::new(cfg.base_ta, false), + AccountMeta::new(cfg.quote_ta, false), + AccountMeta::new_readonly(self.src_mint, false), + AccountMeta::new_readonly(self.dst_mint, false), + AccountMeta::new_readonly(spl_token::id(), false), + AccountMeta::new_readonly(spl_token::id(), false), + AccountMeta::new_readonly(sysvar::instructions::id(), false), + ]); + } + + fn attach_goonfi_accs(&mut self) { + let Some(cfg) = &self.cfg.goonfi else { panic!("Goonfi config is missing, cannot attach accounts."); - } + }; + + let goonfi_param = Pubkey::new_from_array([0u8; 32]); + + self.builder.add_remaining_accounts(&[ + AccountMeta::new_readonly(Pubkey::new_from_array(magnus_shared::pmm_goonfi::id().to_bytes()), false), + AccountMeta::new(self.payer, true), + AccountMeta::new(self.src_ta, false), + AccountMeta::new(self.dst_ta, false), + AccountMeta::new_readonly(goonfi_param, false), + AccountMeta::new(cfg.market, false), + AccountMeta::new(cfg.base_ta, false), + AccountMeta::new(cfg.quote_ta, false), + AccountMeta::new_readonly(cfg.blacklist, false), + AccountMeta::new_readonly(sysvar::instructions::id(), false), + AccountMeta::new_readonly(spl_token::id(), false), + ]); } } @@ -882,6 +874,7 @@ impl Misc { Pubkey::new_from_array(bytes) } + /// Creates fully initialised mint account suitable for use in LiteSVM simulations. fn mk_mint_acc(decimals: u8) -> Account { let mint = spl_token::state::Mint { mint_authority: solana_sdk::program_option::COption::None, @@ -903,18 +896,44 @@ impl Misc { } } + /// Creates a mock SPL Token Account (ATA) with the specified balance. + fn mk_ata(mint: &Pubkey, user: &Pubkey, amount: u64) -> Account { + let ata = spl_token::state::Account { + mint: *mint, + owner: *user, + amount, + state: spl_token::state::AccountState::Initialized, + ..Default::default() + }; + + let mut data = vec![0u8; spl_token::state::Account::LEN]; + ata.pack_into_slice(&mut data); + + Account { + lamports: Rent::default().minimum_balance(data.len()), + data, + owner: spl_token::id(), + executable: false, + rent_epoch: u64::MAX, + } + } + + /// Reads previously saved PMM accounts from disk. + /// + /// Searches the `accounts_path` directory for JSON files matching each DEX's prefix + /// (e.g., `humidifi_*.json`) and deserialises them into account data. fn read_accounts_disk(pmms: &[Dex], accounts_path: &str) -> eyre::Result<(Option, HashMap>)> { - let unique_pmms: HashSet<_> = pmms.iter().collect(); - let mut results = HashMap::new(); + let pmms: HashSet<_> = pmms.iter().collect(); + let mut res = HashMap::new(); let mut all_slots: Vec = vec![]; let data_dir = Path::new(accounts_path); if !data_dir.exists() { - return Ok((None, results)); + return Ok((None, res)); } - for dex in unique_pmms { - let prefix = dex.to_string(); + for pmm in pmms { + let prefix = pmm.to_string(); let mut dex_accounts = vec![]; let mut slots = vec![]; @@ -925,7 +944,7 @@ impl Misc { if path.is_file() && path.file_name().and_then(|n| n.to_str()).is_some_and(|name| name.starts_with(&prefix) && name.ends_with(".json")) { - let (pubkey, account, slot) = Misc::parse_account_from_file(&path)?; + let (slot, pubkey, account) = Misc::parse_account_from_file(&path)?; dex_accounts.push((pubkey, account)); if let Some(s) = slot { slots.push(s); @@ -937,8 +956,8 @@ impl Misc { all_slots.extend(&slots); } - results.insert(*dex, dex_accounts); - info!("loaded accounts for {dex} from disk"); + res.insert(*pmm, dex_accounts); + info!("loaded accounts for {pmm} from disk"); } let slot = if all_slots.is_empty() { @@ -953,38 +972,106 @@ impl Misc { Some(first_slot) }; - Ok((slot, results)) + Ok((slot, res)) } + /// Persists an account to disk in JSON for later reuse. + fn save_account_to_disk(accounts_path: &str, dex: &Dex, pubkey: &Pubkey, account: &Account, slot: u64) -> eyre::Result<()> { + let filename = format!("{}_{}.json", dex, pubkey); + let data_dir = Path::new(&accounts_path); + + if !data_dir.exists() { + fs::create_dir_all(data_dir)?; + } + + let file_path = data_dir.join(filename); + + let value = serde_json::json!({ + "pubkey": pubkey.to_string(), + "slot": slot, + "account": { + "lamports": account.lamports, + "data": [general_purpose::STANDARD.encode(&account.data), "base64"], + "owner": account.owner.to_string(), + "executable": account.executable, + "rentEpoch": account.rent_epoch, + } + }); + + let mut file = File::create(file_path)?; + file.write_all(serde_json::to_string_pretty(&value)?.as_bytes())?; + + Ok(()) + } + + /// Fetches PMM accounts from an RPC node in a single atomic request. + /// + /// Collects all account pubkeys from the provided DEX configurations and fetches + /// them in one `get_multiple_accounts_with_commitment` call to ensure all accounts + /// are read at the same slot. fn fetch_pmm_accounts(pmms: &[Dex], client: &RpcClient, cfg: &PMMCfg) -> eyre::Result<(u64, HashMap>)> { - let slot = client.get_slot()?; - let unique_pmms: HashSet<_> = pmms.iter().collect(); - let mut results = HashMap::new(); + let pmms: HashSet<_> = pmms.iter().collect(); + let mut res = HashMap::new(); - info!("fetching accounts for {pmms:?} at slot {slot}"); - for dex in unique_pmms { - let Some(accounts) = cfg.get_accounts(dex) else { - warn!("skipping unsupported prop amms: {dex}"); + // track which dex the accounts belong to + let mut all_pubkeys: Vec = vec![]; + let mut dex_ranges: Vec<(Dex, std::ops::Range)> = vec![]; + + for pmm in &pmms { + let Some(accounts) = cfg.get_accounts(pmm) else { + warn!("skipping unsupported prop amms: {pmm}"); continue; }; - let fetched = client.get_multiple_accounts(&accounts)?; + let start = all_pubkeys.len(); + all_pubkeys.extend(accounts.iter()); + let end = all_pubkeys.len(); + dex_ranges.push((**pmm, start..end)); + } + + let response = client.get_multiple_accounts_with_commitment(&all_pubkeys, CommitmentConfig::confirmed())?; + let slot = response.context.slot; + let all_accounts = response.value; + + info!("fetched {} accounts for {pmms:?} at slot {slot}", all_pubkeys.len()); + + // reconstruct per-dex account maps + for (dex, range) in dex_ranges { let mut dex_accounts = vec![]; - for (pubkey, account) in accounts.iter().zip(fetched.into_iter()) { - if let Some(acc) = account { - dex_accounts.push((*pubkey, acc)); + for (i, pubkey) in all_pubkeys[range.clone()].iter().enumerate() { + let idx = range.start + i; + if let Some(acc) = &all_accounts[idx] { + dex_accounts.push((*pubkey, acc.clone())); } else { warn!("account {pubkey} not found for {dex}"); } } - - results.insert(*dex, dex_accounts); + res.insert(dex, dex_accounts); } - Ok((slot, results)) - } - - fn parse_account_from_file(path: &Path) -> eyre::Result<(Pubkey, Account, Option)> { + Ok((slot, res)) + } + + /// Parses a Solana account from a JSON file. + /// + /// Expected JSON format (matches Solana CLI `account` command output): + /// ```json + /// { + /// "slot": 12345678, + /// "pubkey": "Base58EncodedPubkey", + /// "account": { + /// "lamports": 1000000, + /// "data": ["Base64EncodedData", "base64"], + /// "owner": "Base58EncodedOwner", + /// "executable": false, + /// "rentEpoch": 0 + /// } + /// } + /// ``` + /// + /// # Returns + /// A tuple of (slot, pubkey, account). + fn parse_account_from_file(path: &Path) -> eyre::Result<(Option, Pubkey, Account)> { let contents = fs::read_to_string(path)?; let value: serde_json::Value = serde_json::from_str(&contents)?; @@ -997,9 +1084,13 @@ impl Misc { let rent_epoch = value["account"]["rentEpoch"].as_u64().ok_or_else(|| eyre::eyre!("missing rentEpoch"))?; let slot = value["slot"].as_u64(); - Ok((pubkey, Account { lamports, data, owner, executable, rent_epoch }, slot)) + Ok((slot, pubkey, Account { lamports, data, owner, executable, rent_epoch })) } + /// Custom serde deserializer for `Pubkey` from a base58-encoded string. + /// + /// Used with `#[serde(deserialize_with = "Misc::deserialize_pubkey")]` attribute + /// on struct fields that should be deserialized as Solana pubkeys. fn deserialize_pubkey<'de, D>(deserializer: D) -> Result where D: serde::Deserializer<'de>, @@ -1007,9 +1098,19 @@ impl Misc { let s = String::deserialize(deserializer)?; Pubkey::from_str(&s).map_err(serde::de::Error::custom) } + + /// Converts a raw token amount (in base units) to a human-readable decimal value. + fn to_human(amount: u64, dec: u8) -> f64 { + amount as f64 / 10f64.powi(dec as i32) + } + + /// Converts a human-readable decimal value to raw token amount (in base units). + fn to_raw(amount: f64, dec: u8) -> u64 { + (amount * 10f64.powi(dec as i32)) as u64 + } } -#[derive(Debug, Serialize)] +#[derive(Debug, Serialize, Clone)] struct BenchmarkRecord<'a> { slot: u64, pmm: Dex, @@ -1025,7 +1126,6 @@ struct BenchmarkRecord<'a> { struct Benchmark<'a> { records: Vec>, writer: csv::Writer, - save_path: String, } impl<'a> Benchmark<'a> { @@ -1037,7 +1137,7 @@ impl<'a> Benchmark<'a> { }; let writer = Writer::from_path(&save_path)?; - Ok(Benchmark { records, writer, save_path }) + Ok(Benchmark { records, writer }) } pub fn save(&mut self) -> eyre::Result<()> { @@ -1045,13 +1145,51 @@ impl<'a> Benchmark<'a> { self.writer.serialize(record)?; } - info!("saved benchmark records @ {}", self.save_path); self.writer.flush()?; Ok(()) } } +/// Benchmark step configuration with normalized values. +#[derive(Debug, Clone, Copy)] +struct BenchmarkSteps { + start: u64, + end: u64, + step: u64, +} + +impl BenchmarkSteps { + /// Creates normalized benchmark steps from human-readable values. + /// + /// Converts floating-point token amounts to their base unit representation + /// using the token's decimal places. + /// + /// # Arguments + /// * `steps` - Array of [start, end, step] in human-readable token amounts + /// * `dec` - Number of decimal places for the token (e.g., 9 for SOL, 6 for USDC) + /// + /// # Example + /// ``` + /// // For WSOL (9 decimals): 1.0 to 100.0 with step 0.5 + /// let steps = BenchmarkSteps::from_human([1.0, 100.0, 0.5], 9); + /// assert_eq!(steps.start, 1_000_000_000); // 1 SOL in lamports + /// ``` + fn from_human(steps: [f64; 3], dec: u8) -> Self { + Self { start: Misc::to_raw(steps[0], dec), end: Misc::to_raw(steps[1], dec), step: Misc::to_raw(steps[2], dec) } + } + + /// Returns the total number of iterations for this step configuration. + fn count(&self) -> u64 { + (self.end - self.start) / self.step + 1 + } + + /// Returns an iterator over all step values from start to end (inclusive). + fn iter(&self) -> impl Iterator { + (self.start..=self.end).step_by(self.step as usize) + } +} + pub struct Run { args: CliArgs, cfg: PMMCfg, @@ -1074,12 +1212,10 @@ impl Run { let Command::FetchAccounts { http_url, accounts_path, pmms, .. } = &self.args.command else { unreachable!() }; let rpc_client = RpcClient::new(http_url.expose_secret().to_string()); - let env = Environment::new("", accounts_path, None, self.cfg.clone(), None)?; - let (slot, fetched) = Misc::fetch_pmm_accounts(pmms, &rpc_client, &self.cfg)?; for (dex, accounts) in fetched { for (pubkey, account) in accounts { - env.save_account_to_disk(&dex, &pubkey, &account, slot)?; + Misc::save_account_to_disk(accounts_path, &dex, &pubkey, &account, slot)?; info!("saved account {pubkey} for {dex}"); } } @@ -1089,21 +1225,15 @@ impl Run { } fn benchmark(&self) -> eyre::Result<()> { - let Command::Benchmark { common, datasets_path, pmms, step } = &self.args.command else { unreachable!() }; + let Command::Benchmark { common, datasets_path, pmms, steps } = &self.args.command else { unreachable!() }; let rpc_client = RpcClient::new(common.http_url.expose_secret().to_string()); let (src_mint, src_dec, src_name) = (common.src_token.get_addr(), common.src_token.get_decimals(), common.src_token.to_string()); let (dst_mint, dst_dec, dst_name) = (common.dst_token.get_addr(), common.dst_token.get_decimals(), common.dst_token.to_string()); let mints = vec![(src_mint, src_dec), (dst_mint, dst_dec)]; - let norm_step = [ - (step[0] * 10f64.powi(src_dec as i32)) as u64, - (step[1] * 10f64.powi(src_dec as i32)) as u64, - (step[2] * 10f64.powi(src_dec as i32)) as u64, - ]; - let steps_count = ((norm_step[1] - norm_step[0]) / norm_step[2] + 1) as u64; - - let time = Local::now().format("%Y%m%d-%H%M%S").to_string(); + let steps = BenchmarkSteps::from_human(*steps, src_dec); + let time_fmt = Local::now().format("%Y%m%d-%H%M%S").to_string(); let multi = MultiProgress::new(); let (slot, accs_map) = if common.jit_accounts { @@ -1117,37 +1247,40 @@ impl Run { let handles: Vec<_> = pmms .iter() .map(|pmm| { - let (cfg, multi, mints, time) = (&self.cfg, &multi, &mints, &time); + let (cfg, multi, mints, time) = (&self.cfg, &multi, &mints, &time_fmt); let (src_name, dst_name) = (&src_name, &dst_name); let pmm_accounts = accs_map.get(pmm).cloned().unwrap_or_default(); s.spawn(move || -> eyre::Result<()> { // start up the progress bar only when all the spawned threads - // have finished bootstrapping so there's no CLI race - let (mut env, src_ata, dst_ata, original_accounts) = multi.suspend(|| -> eyre::Result<_> { + // have finished bootstrapping so there's no CLI progress bar race cond + let (mut env, src_ata, dst_ata) = multi.suspend(|| -> eyre::Result<_> { let mut env = Environment::new(&common.programs_path, &common.accounts_path, Some(mints), cfg.clone(), slot)?; env.load_programs(&[*pmm])?; - env.load_accounts(&pmm_accounts.clone())?; - env.setup_wallet(&src_mint, norm_step[1], 10_000_000_000)?; + env.setup_wallet(&src_mint, steps.end, consts::AIRDROP_AMOUNT)?; let (src_ata, dst_ata) = (env.wallet_ata(&src_mint), env.wallet_ata(&dst_mint)); - Ok((env, src_ata, dst_ata, pmm_accounts)) + Ok((env, src_ata, dst_ata)) })?; let market = cfg.get_market(pmm).unwrap_or_else(|| panic!("{} not configured", pmm)).to_string(); - let pb = multi.add(ProgressBar::new(steps_count)); - pb.set_style(ProgressStyle::default_bar().template(consts::PROGRESS_TEMPLATE)?.progress_chars("█▓░")); + let pb = multi.add(ProgressBar::new(steps.count())); + pb.set_style( + ProgressStyle::default_bar().template(consts::PROGRESS_TEMPLATE)?.progress_chars(consts::PROGRESS_CHARS), + ); pb.set_prefix(format!("{}", pmm)); - let mut r = vec![]; - let mut warn_count = 0u64; - let route: Vec = vec![Route { dexes: vec![*pmm], weights: vec![100] }.into()]; - for amount_in in (norm_step[0]..=norm_step[1]).step_by(norm_step[2] as usize) { - let order_id = SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(); + let (mut records, mut warn_cnt) = (vec![], u64::default()); + let routes: Vec> = + vec![vec![Route { dexes: vec![*pmm], weights: vec![100] }.into()]]; + + for amount_in in steps.iter() { env.reset_wallet(&src_mint, amount_in)?; - env.load_accounts(&original_accounts)?; + env.load_accounts(&pmm_accounts)?; + + let order_id = SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(); let mut swap_builder = SwapBuilder::new(); let swap = swap_builder @@ -1160,15 +1293,15 @@ impl Run { .expect_amount_out(1) .min_return(1) .amounts(vec![amount_in]) - .routes(vec![route.clone()]) + .routes(routes.clone()) .order_id(order_id); let mut construct = ConstructSwap { cfg: cfg.clone(), builder: swap, payer: env.wallet_pubkey(), - sta: src_ata, - dta: dst_ata, + src_ta: src_ata, + dst_ta: dst_ata, src_mint, dst_mint, }; @@ -1186,10 +1319,8 @@ impl Run { let res = match env.send_transaction(tx) { Ok(res) => res, Err(e) => { - if warn_count == 0 { - pb.println(format!("[WARN] {}: {:?}", pmm, e)); - } - warn_count += 1; + (warn_cnt == 0).then(|| pb.println(format!("[WARN] {}: {:?}", pmm, e))); + warn_cnt += 1; pb.inc(1); continue; } @@ -1197,27 +1328,27 @@ impl Run { let amount_out = env.get_event_amount_out(&res); - r.push(BenchmarkRecord { + records.push(BenchmarkRecord { slot: env.slot.unwrap_or_default(), pmm: *pmm, market: &market, src_token: src_name, dst_token: dst_name, - amount_in: amount_in as f64 / 10f64.powi(src_dec as i32), - amount_out: amount_out as f64 / 10f64.powi(dst_dec as i32), + amount_in: Misc::to_human(amount_in, src_dec), + amount_out: Misc::to_human(amount_out, dst_dec), compute_units: res.compute_units_consumed, }); - pb.set_message(format!("in: {:.2}", amount_in as f64 / 10f64.powi(src_dec as i32))); + pb.set_message(format!("in: {:.2}", Misc::to_human(amount_in, src_dec))); pb.inc(1); } - if warn_count > 0 { - pb.println(format!("[WARN] {}: {} total failures", pmm, warn_count)); - } + (warn_cnt != 0).then(|| pb.println(format!("[WARN] {}: {} total failures", pmm, warn_cnt))); let filename = format!("{}/{}_{}_{}_{}.csv", datasets_path, env.slot.unwrap_or_default(), pmm, market, time); - let _ = Benchmark::new(r, &filename)?.save(); + Benchmark::new(records.clone(), &filename)?.save().is_ok().then(|| { + pb.println(format!("[{}] saved {} records to {}", pmm, records.len(), filename)); + }); Ok(()) }) @@ -1263,20 +1394,18 @@ impl Run { env.load_programs(&flat_pmms)?; env.fetch_and_load_accounts(&flat_pmms, common.jit_accounts, Some(&rpc_client))?; - let norm_amount_in: Vec = amount_in.iter().map(|amount| amount * 10f64.powi(src_dec as i32)).map(|a| a as u64).collect(); - let norm_amount_in_sum: u64 = norm_amount_in.iter().sum(); + let amount_in: Vec = amount_in.iter().map(|amount| Misc::to_raw(*amount, src_dec)).collect(); + let amount_in_sum: u64 = amount_in.iter().sum(); // - mint only the source token's desired amount (i.e the amount we're going to swap) // - airdrop some SOL to cover fees - env.setup_wallet(&src_mint, norm_amount_in_sum, 10_000_000_000)?; + env.setup_wallet(&src_mint, amount_in_sum, consts::AIRDROP_AMOUNT)?; info!(?env); - let (src_ata, dst_ata) = (env.wallet_ata(&src_mint), env.wallet_ata(&dst_mint)); - let (src_before, dst_before) = ( - env.token_balance(&src_mint) as f64 / 10_f64.powi(src_dec as i32), - env.token_balance(&dst_mint) as f64 / 10_f64.powi(dst_dec as i32), - ); - info!("before: {} = {} | {} = {}", src_name, src_before, dst_name, dst_before); + let (src_ata, src_before) = (env.wallet_ata(&src_mint), env.token_balance_norm(&src_mint, src_dec)); + let (dst_ata, dst_before) = (env.wallet_ata(&dst_mint), env.token_balance_norm(&dst_mint, dst_dec)); + + debug!(?src_name, ?src_before, ?dst_name, ?dst_before); let routes: Vec> = pmms .iter() @@ -1284,8 +1413,7 @@ impl Run { .map(|(dex_grp, weight_group)| vec![Route { dexes: dex_grp.clone(), weights: weight_group.clone() }.into()]) .collect(); - info!("swapping {:?} {} via routes: {:?}", norm_amount_in, src_name, routes); - + let order_id = SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(); let mut swap_builder = SwapBuilder::new(); let swap = swap_builder .payer(env.wallet_pubkey()) @@ -1293,19 +1421,19 @@ impl Run { .destination_token_account(dst_ata) .source_mint(src_mint) .destination_mint(dst_mint) - .amount_in(norm_amount_in_sum) + .amount_in(amount_in_sum) .expect_amount_out(1) .min_return(1) - .amounts(norm_amount_in) - .routes(routes) - .order_id(SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs()); + .amounts(amount_in) + .routes(routes.clone()) + .order_id(order_id); let mut construct = ConstructSwap { cfg: self.cfg.clone(), builder: swap, payer: env.wallet_pubkey(), - sta: src_ata, - dta: dst_ata, + src_ta: src_ata, + dst_ta: dst_ata, src_mint, dst_mint, }; @@ -1316,21 +1444,19 @@ impl Run { }); let swap_ix = construct.instruction(); - debug!("router program id: {}", swap_ix.program_id); - let tx = Transaction::new_signed_with_payer(&[swap_ix], Some(&env.wallet_pubkey()), &[&env.wallet], env.latest_blockhash()); let res = env.send_transaction(tx).expect("failed to exec tx"); let amount_out = env.get_event_amount_out(&res); - let (src_after, dst_after) = ( - env.token_balance(&src_mint) as f64 / 10_f64.powi(src_dec as i32), - env.token_balance(&dst_mint) as f64 / 10_f64.powi(dst_dec as i32), + info!( + src_token = %src_name, + dst_token = %dst_name, + routes = ?routes, + amount_in = ?Misc::to_human(amount_in_sum, src_dec), + amount_out = ?Misc::to_human(amount_out, dst_dec), + cu = res.compute_units_consumed ); - info!("|SWAP EXECUTED| compute units consumed: {:?} | amount_out: {}", res.compute_units_consumed, amount_out); - info!("after: {} = {:.6} | {} = {:.6} | ", src_name, src_after, dst_name, dst_after); - info!("diff: {} spent = {:.6} | {} received = {:.6}", src_name, src_before - src_after, dst_name, dst_after - dst_before); - Ok(()) } } @@ -1358,134 +1484,410 @@ fn main() -> eyre::Result<()> { mod tests { use super::*; - #[test] - fn test_parse_nested_pmms_json_single() { - let input = r#"[["humidifi"]]"#; - let result = CliArgs::parse_nested_pmms(input).unwrap(); - assert_eq!(result, vec![vec![Dex::Humidifi]]); - } + mod cli { + use super::*; + #[test] + fn test_parse_nested_pmms_json_single() { + let input = r#"[["humidifi"]]"#; + let result = CliArgs::parse_nested_pmms(input).unwrap(); + assert_eq!(result, vec![vec![Dex::Humidifi]]); + } - #[test] - fn test_parse_nested_pmms_json_multiple() { - let input = r#"[["humidifi","obric-v2"],["zerofi"]]"#; - let result = CliArgs::parse_nested_pmms(input).unwrap(); - assert_eq!(result, vec![vec![Dex::Humidifi, Dex::ObricV2], vec![Dex::Zerofi]]); - } + #[test] + fn test_parse_nested_pmms_json_multiple() { + let input = r#"[["humidifi","obric-v2"],["zerofi"]]"#; + let result = CliArgs::parse_nested_pmms(input).unwrap(); + assert_eq!(result, vec![vec![Dex::Humidifi, Dex::ObricV2], vec![Dex::Zerofi]]); + } - #[test] - fn test_parse_nested_pmms_no_quotes_single() { - let input = "[[humidifi]]"; - let result = CliArgs::parse_nested_pmms(input).unwrap(); - assert_eq!(result, vec![vec![Dex::Humidifi]]); - } + #[test] + fn test_parse_nested_pmms_no_quotes_single() { + let input = "[[humidifi]]"; + let result = CliArgs::parse_nested_pmms(input).unwrap(); + assert_eq!(result, vec![vec![Dex::Humidifi]]); + } - #[test] - fn test_parse_nested_pmms_no_quotes_single_route_multiple_pmms() { - let input = "[[humidifi,obric-v2]]"; - let result = CliArgs::parse_nested_pmms(input).unwrap(); - assert_eq!(result, vec![vec![Dex::Humidifi, Dex::ObricV2]]); - } + #[test] + fn test_parse_nested_pmms_no_quotes_single_route_multiple_pmms() { + let input = "[[humidifi,obric-v2]]"; + let result = CliArgs::parse_nested_pmms(input).unwrap(); + assert_eq!(result, vec![vec![Dex::Humidifi, Dex::ObricV2]]); + } - #[test] - fn test_parse_nested_pmms_no_quotes_multiple_routes() { - let input = "[[humidifi,obric-v2],[zerofi]]"; - let result = CliArgs::parse_nested_pmms(input).unwrap(); - assert_eq!(result, vec![vec![Dex::Humidifi, Dex::ObricV2], vec![Dex::Zerofi]]); - } + #[test] + fn test_parse_nested_pmms_no_quotes_multiple_routes() { + let input = "[[humidifi,obric-v2],[zerofi]]"; + let result = CliArgs::parse_nested_pmms(input).unwrap(); + assert_eq!(result, vec![vec![Dex::Humidifi, Dex::ObricV2], vec![Dex::Zerofi]]); + } - #[test] - fn test_parse_nested_pmms_no_quotes_three_routes() { - let input = "[[humidifi],[obric-v2,solfi-v2],[zerofi]]"; - let result = CliArgs::parse_nested_pmms(input).unwrap(); - assert_eq!(result, vec![vec![Dex::Humidifi], vec![Dex::ObricV2, Dex::SolfiV2], vec![Dex::Zerofi],]); - } + #[test] + fn test_parse_nested_pmms_no_quotes_three_routes() { + let input = "[[humidifi],[obric-v2,solfi-v2],[zerofi]]"; + let result = CliArgs::parse_nested_pmms(input).unwrap(); + assert_eq!(result, vec![vec![Dex::Humidifi], vec![Dex::ObricV2, Dex::SolfiV2], vec![Dex::Zerofi],]); + } - #[test] - fn test_parse_nested_pmms_no_quotes_all_pmms() { - let input = "[[raydium-cl-v2,raydium-cp],[obric-v2,solfi-v2,zerofi,humidifi]]"; - let result = CliArgs::parse_nested_pmms(input).unwrap(); - assert_eq!(result, vec![vec![Dex::RaydiumClV2, Dex::RaydiumCp], vec![Dex::ObricV2, Dex::SolfiV2, Dex::Zerofi, Dex::Humidifi],]); - } + #[test] + fn test_parse_nested_pmms_no_quotes_all_pmms() { + let input = "[[raydium-cl-v2,raydium-cp],[obric-v2,solfi-v2,zerofi,humidifi]]"; + let result = CliArgs::parse_nested_pmms(input).unwrap(); + assert_eq!(result, vec![vec![Dex::RaydiumClV2, Dex::RaydiumCp], vec![Dex::ObricV2, Dex::SolfiV2, Dex::Zerofi, Dex::Humidifi],]); + } - #[test] - fn test_parse_nested_pmms_no_quotes_with_spaces() { - let input = "[[ humidifi , obric-v2 ],[ zerofi ]]"; - let result = CliArgs::parse_nested_pmms(input).unwrap(); - assert_eq!(result, vec![vec![Dex::Humidifi, Dex::ObricV2], vec![Dex::Zerofi]]); - } + #[test] + fn test_parse_nested_pmms_no_quotes_with_spaces() { + let input = "[[ humidifi , obric-v2 ],[ zerofi ]]"; + let result = CliArgs::parse_nested_pmms(input).unwrap(); + assert_eq!(result, vec![vec![Dex::Humidifi, Dex::ObricV2], vec![Dex::Zerofi]]); + } - #[test] - fn test_parse_nested_pmms_invalid_pmm() { - let input = "[[humidifi,invalid-dex]]"; - let result = CliArgs::parse_nested_pmms(input); - assert!(result.is_err()); - } + #[test] + fn test_parse_nested_pmms_invalid_pmm() { + let input = "[[humidifi,invalid-dex]]"; + let result = CliArgs::parse_nested_pmms(input); + assert!(result.is_err()); + } - #[test] - fn test_parse_nested_pmms_invalid_format() { - let input = "[humidifi]"; // not nested - let result = CliArgs::parse_nested_pmms(input); - assert!(result.is_err()); - } + #[test] + fn test_parse_nested_pmms_invalid_format() { + let input = "[humidifi]"; // not nested + let result = CliArgs::parse_nested_pmms(input); + assert!(result.is_err()); + } - #[test] - fn test_parse_nested_weights_single() { - let input = "[[100]]"; - let result = CliArgs::parse_nested_weights(input).unwrap(); - assert_eq!(result, vec![vec![100u8]]); - } + #[test] + fn test_parse_nested_weights_single() { + let input = "[[100]]"; + let result = CliArgs::parse_nested_weights(input).unwrap(); + assert_eq!(result, vec![vec![100u8]]); + } - #[test] - fn test_parse_nested_weights_multiple() { - let input = "[[50,50],[100]]"; - let result = CliArgs::parse_nested_weights(input).unwrap(); - assert_eq!(result, vec![vec![50u8, 50u8], vec![100u8]]); - } + #[test] + fn test_parse_nested_weights_multiple() { + let input = "[[50,50],[100]]"; + let result = CliArgs::parse_nested_weights(input).unwrap(); + assert_eq!(result, vec![vec![50u8, 50u8], vec![100u8]]); + } - #[test] - fn test_parse_nested_weights_complex() { - let input = "[[30,30,40],[60,40],[100]]"; - let result = CliArgs::parse_nested_weights(input).unwrap(); - assert_eq!(result, vec![vec![30u8, 30u8, 40u8], vec![60u8, 40u8], vec![100u8]]); - } + #[test] + fn test_parse_nested_weights_complex() { + let input = "[[30,30,40],[60,40],[100]]"; + let result = CliArgs::parse_nested_weights(input).unwrap(); + assert_eq!(result, vec![vec![30u8, 30u8, 40u8], vec![60u8, 40u8], vec![100u8]]); + } - #[test] - fn test_pmms_and_weights_match() { - let pmms_input = "[[humidifi,obric-v2],[zerofi]]"; - let weights_input = "[[50,50],[100]]"; + #[test] + fn test_pmms_and_weights_match() { + let pmms_input = "[[humidifi,obric-v2],[zerofi]]"; + let weights_input = "[[50,50],[100]]"; - let pmms = CliArgs::parse_nested_pmms(pmms_input).unwrap(); - let weights = CliArgs::parse_nested_weights(weights_input).unwrap(); + let pmms = CliArgs::parse_nested_pmms(pmms_input).unwrap(); + let weights = CliArgs::parse_nested_weights(weights_input).unwrap(); - assert_eq!(pmms.len(), weights.len()); - for (d, w) in pmms.iter().zip(weights.iter()) { - assert_eq!(d.len(), w.len()); + assert_eq!(pmms.len(), weights.len()); + for (d, w) in pmms.iter().zip(weights.iter()) { + assert_eq!(d.len(), w.len()); + } } - } - #[test] - fn test_parse_step_valid() { - let result = CliArgs::parse_step("1.0,100.0,0.5").unwrap(); - assert_eq!(result, [1.0, 100.0, 0.5]); - } + #[test] + fn test_parse_step_valid() { + let result = CliArgs::parse_steps("1.0,100.0,0.5").unwrap(); + assert_eq!(result, [1.0, 100.0, 0.5]); + } - #[test] - fn test_parse_step_with_spaces() { - let result = CliArgs::parse_step("1.0, 100.0, 0.5").unwrap(); - assert_eq!(result, [1.0, 100.0, 0.5]); + #[test] + fn test_parse_step_with_spaces() { + let result = CliArgs::parse_steps("1.0, 100.0, 0.5").unwrap(); + assert_eq!(result, [1.0, 100.0, 0.5]); + } + + #[test] + fn test_parse_step_start_gte_end() { + let result = CliArgs::parse_steps("100.0,50.0,1.0"); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("start must be less than end")); + } + + #[test] + fn test_parse_step_negative_step() { + let result = CliArgs::parse_steps("1.0,100.0,-1.0"); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("step must be positive")); + } } - #[test] - fn test_parse_step_start_gte_end() { - let result = CliArgs::parse_step("100.0,50.0,1.0"); - assert!(result.is_err()); - assert!(result.unwrap_err().contains("start must be less than end")); + mod environment { + use super::*; + + fn default_cfg() -> PMMCfg { + PMMCfg::default() + } + + #[test] + fn test_new_generates_unique_wallet() { + let env1 = Environment::new("", "", None, default_cfg(), None).unwrap(); + let env2 = Environment::new("", "", None, default_cfg(), None).unwrap(); + + assert_ne!(env1.wallet_pubkey(), env2.wallet_pubkey()); + } + + #[test] + fn test_new_without_slot_leaves_slot_none() { + let env = Environment::new("", "", None, default_cfg(), None).unwrap(); + + assert!(env.slot.is_none()); + } + + #[test] + fn test_new_with_slot_sets_slot() { + let slot = 12345678u64; + let env = Environment::new("", "", None, default_cfg(), Some(slot)).unwrap(); + + assert_eq!(env.slot, Some(slot)); + } + + #[test] + fn test_new_without_mints_leaves_mints_none() { + let env = Environment::new("", "", None, default_cfg(), None).unwrap(); + + assert!(env.mints.is_none()); + } + + #[test] + fn test_new_with_mints_creates_mint_accounts() { + let mints = vec![(consts::WSOL, consts::WSOL_DECIMALS), (consts::USDC, consts::USDC_DECIMALS)]; + + let env = Environment::new("", "", Some(&mints), default_cfg(), None).unwrap(); + + assert!(env.mints.is_some()); + assert_eq!(env.mints.unwrap().len(), 2); + + // verify mint accounts exist in SVM + let wsol_account = env.svm.get_account(&consts::WSOL); + let usdc_account = env.svm.get_account(&consts::USDC); + + assert!(wsol_account.is_some()); + assert!(usdc_account.is_some()); + + // verify mint account data is valid + let wsol_mint = spl_token::state::Mint::unpack(&wsol_account.unwrap().data).unwrap(); + let usdc_mint = spl_token::state::Mint::unpack(&usdc_account.unwrap().data).unwrap(); + + assert_eq!(wsol_mint.decimals, consts::WSOL_DECIMALS); + assert_eq!(usdc_mint.decimals, consts::USDC_DECIMALS); + assert!(wsol_mint.is_initialized); + assert!(usdc_mint.is_initialized); + } + + #[test] + fn test_wallet_ata_derives_correct_address() { + let mints = vec![(consts::WSOL, consts::WSOL_DECIMALS)]; + let env = Environment::new("", "", Some(&mints), default_cfg(), None).unwrap(); + + let expected_ata = get_associated_token_address(&env.wallet_pubkey(), &consts::WSOL); + let actual_ata = env.wallet_ata(&consts::WSOL); + + assert_eq!(actual_ata, expected_ata); + } + + #[test] + fn test_setup_wallet_creates_atas_and_funds() { + let mints = vec![(consts::WSOL, consts::WSOL_DECIMALS), (consts::USDC, consts::USDC_DECIMALS)]; + let mut env = Environment::new("", "", Some(&mints), default_cfg(), None).unwrap(); + + let src_amount = 1_000_000_000u64; // 1 SOL + let airdrop = 10_000_000_000u64; // 10 SOL for fees + + env.setup_wallet(&consts::WSOL, src_amount, airdrop).unwrap(); + + // verify src token balance is correct + assert_eq!(env.token_balance(&consts::WSOL), src_amount); + + // verify dst token balance is zero + assert_eq!(env.token_balance(&consts::USDC), 0); + + // verify SOL was airdropped + let wallet_account = env.svm.get_account(&env.wallet_pubkey()); + assert!(wallet_account.is_some()); + assert!(wallet_account.unwrap().lamports >= airdrop); + } + + #[test] + fn test_reset_wallet_restores_balances() { + let mints = vec![(consts::WSOL, consts::WSOL_DECIMALS), (consts::USDC, consts::USDC_DECIMALS)]; + let mut env = Environment::new("", "", Some(&mints), default_cfg(), None).unwrap(); + + env.setup_wallet(&consts::WSOL, 1_000_000_000, 10_000_000_000).unwrap(); + + // simulate a swap by manually changing balances + let wsol_ata = env.wallet_ata(&consts::WSOL); + let usdc_ata = env.wallet_ata(&consts::USDC); + env.svm.set_account(wsol_ata, Misc::mk_ata(&consts::WSOL, &env.wallet_pubkey(), 500_000_000)).unwrap(); + env.svm.set_account(usdc_ata, Misc::mk_ata(&consts::USDC, &env.wallet_pubkey(), 100_000_000)).unwrap(); + + // seset wallet + let new_amount = 2_000_000_000u64; + env.reset_wallet(&consts::WSOL, new_amount).unwrap(); + + // verify balances are reset + assert_eq!(env.token_balance(&consts::WSOL), new_amount); + assert_eq!(env.token_balance(&consts::USDC), 0); + } + + #[test] + fn test_token_balance_returns_zero_for_nonexistent_ata() { + let env = Environment::new("", "", None, default_cfg(), None).unwrap(); + + assert_eq!(env.token_balance(&consts::WSOL), 0); + } + + #[test] + fn test_token_balance_norm_converts_correctly() { + let mints = vec![(consts::WSOL, consts::WSOL_DECIMALS)]; + let mut env = Environment::new("", "", Some(&mints), default_cfg(), None).unwrap(); + + let raw_amount = 1_500_000_000u64; // 1.5 SOL in lamports + env.setup_wallet(&consts::WSOL, raw_amount, 10_000_000_000).unwrap(); + + let normalized = env.token_balance_norm(&consts::WSOL, consts::WSOL_DECIMALS); + + assert!((normalized - 1.5).abs() < f64::EPSILON); + } + + #[test] + fn test_latest_blockhash_returns_valid_hash() { + let env = Environment::new("", "", None, default_cfg(), None).unwrap(); + + let blockhash = env.latest_blockhash(); + + // should not be the default/zero hash + assert_ne!(blockhash, solana_sdk::hash::Hash::default()); + } + + #[test] + fn test_load_accounts_sets_accounts_in_svm() { + let mut env = Environment::new("", "", None, default_cfg(), None).unwrap(); + + let pubkey = Pubkey::new_unique(); + let account = + Account { lamports: 1_000_000, data: vec![1, 2, 3, 4], owner: Pubkey::new_unique(), executable: false, rent_epoch: 0 }; + + env.load_accounts(&vec![(pubkey, account.clone())]).unwrap(); + + let loaded = env.svm.get_account(&pubkey).unwrap(); + assert_eq!(loaded.lamports, account.lamports); + assert_eq!(loaded.data, account.data); + assert_eq!(loaded.owner, account.owner); + } } - #[test] - fn test_parse_step_negative_step() { - let result = CliArgs::parse_step("1.0,100.0,-1.0"); - assert!(result.is_err()); - assert!(result.unwrap_err().contains("step must be positive")); + mod misc { + use super::*; + + #[test] + fn test_to_human_wsol_decimals() { + let lamports = 1_500_000_000u64; + let sol = Misc::to_human(lamports, 9); + assert_eq!(sol, 1.5); + } + + #[test] + fn test_to_human_usdc_decimals() { + let raw = 1_500_000u64; + let usdc = Misc::to_human(raw, 6); + assert_eq!(usdc, 1.5); + } + + #[test] + fn test_to_human_zero() { + assert_eq!(Misc::to_human(0, 9), 0.0); + } + + #[test] + fn test_to_raw_wsol_decimals() { + let sol = 1.5f64; + let lamports = Misc::to_raw(sol, 9); + assert_eq!(lamports, 1_500_000_000); + } + + #[test] + fn test_to_raw_usdc_decimals() { + let usdc = 1.5f64; + let raw = Misc::to_raw(usdc, 6); + assert_eq!(raw, 1_500_000); + } + + #[test] + fn test_to_raw_zero() { + assert_eq!(Misc::to_raw(0.0, 9), 0); + } + + #[test] + fn test_to_human_to_raw_roundtrip() { + let original = 123_456_789u64; + let decimals = 9u8; + + let human = Misc::to_human(original, decimals); + let back = Misc::to_raw(human, decimals); + + assert_eq!(back, original); + } + + #[test] + fn test_mk_mint_acc_creates_valid_mint() { + let decimals = 9u8; + let account = Misc::mk_mint_acc(decimals); + + assert_eq!(account.owner, spl_token::id()); + assert!(!account.executable); + + let mint = spl_token::state::Mint::unpack(&account.data).unwrap(); + assert_eq!(mint.decimals, decimals); + assert!(mint.is_initialized); + assert_eq!(mint.supply, u64::MAX); + assert!(mint.mint_authority.is_none()); + assert!(mint.freeze_authority.is_none()); + } + + #[test] + fn test_mk_mint_acc_different_decimals() { + for decimals in [0, 6, 9, 18] { + let account = Misc::mk_mint_acc(decimals); + let mint = spl_token::state::Mint::unpack(&account.data).unwrap(); + assert_eq!(mint.decimals, decimals); + } + } + + #[test] + fn test_mk_ata_creates_valid_token_account() { + let mint = Pubkey::new_unique(); + let owner = Pubkey::new_unique(); + let amount = 1_000_000u64; + + let account = Misc::mk_ata(&mint, &owner, amount); + + // verify owner is token program + assert_eq!(account.owner, spl_token::id()); + + // verify data is valid token account + let token_acc = spl_token::state::Account::unpack(&account.data).unwrap(); + assert_eq!(token_acc.mint, mint); + assert_eq!(token_acc.owner, owner); + assert_eq!(token_acc.amount, amount); + assert_eq!(token_acc.state, spl_token::state::AccountState::Initialized); + } + + #[test] + fn test_mk_ata_zero_balance() { + let mint = Pubkey::new_unique(); + let owner = Pubkey::new_unique(); + + let account = Misc::mk_ata(&mint, &owner, 0); + let token_acc = spl_token::state::Account::unpack(&account.data).unwrap(); + + assert_eq!(token_acc.amount, 0); + } } }