gasbench/
main.rs

1use alloy_primitives::{Address, BlockNumber, Bytes};
2use alloy_provider::Provider;
3use alloy_rpc_types::{BlockNumberOrTag, TransactionInput, TransactionRequest};
4use clap::Parser;
5use colored::*;
6use eyre::{OptionExt, Result};
7use serde::{Deserialize, Serialize};
8use std::collections::{BTreeMap, HashMap};
9use std::env;
10use std::sync::Arc;
11use tokio::fs::File;
12use tokio::io::{AsyncReadExt, AsyncWriteExt};
13use tracing::{error, info};
14
15use crate::cli::Cli;
16use crate::dto::SwapLineDTO;
17use crate::preloader::preload_pools;
18use crate::soltest::create_sol_test;
19use kabu_node_debug_provider::AnvilDebugProviderFactory;
20
21use kabu_defi_address_book::UniswapV2PoolAddress;
22use kabu_types_entities::{Market, MarketState, PoolId, PoolWrapper, Swap, SwapAmountType, SwapDirection, SwapLine};
23
24use kabu_core_actors::SharedState;
25use kabu_defi_preloader::preload_market_state;
26use kabu_evm_db::KabuDBType;
27use kabu_evm_utils::{BalanceCheater, NWETH};
28use kabu_execution_multicaller::pool_opcodes_encoder::ProtocolSwapOpcodesEncoderV2;
29use kabu_execution_multicaller::{
30    MulticallerDeployer, MulticallerEncoder, MulticallerSwapEncoder, ProtocolABIEncoderV2, SwapLineEncoder, SwapStepEncoder,
31};
32use revm::database::CacheDB;
33
34mod cli;
35mod dto;
36mod preloader;
37mod soltest;
38
39#[derive(Clone, Debug, Serialize, Deserialize)]
40struct SwapPathDTO {
41    tokens: Vec<Address>,
42    pools: Vec<Address>,
43}
44
45#[tokio::main]
46async fn main() -> Result<()> {
47    let cli: Cli = Cli::try_parse()?;
48
49    env_logger::init_from_env(env_logger::Env::default().default_filter_or("info"));
50    let block_number = 20089277u64;
51
52    println!("Hello, block {block_number}!");
53
54    let node_url = env::var("MAINNET_WS")?;
55
56    let client = AnvilDebugProviderFactory::from_node_on_block(node_url, BlockNumber::from(block_number)).await?;
57
58    let _block_header = client.get_block_by_number(BlockNumberOrTag::Number(block_number)).await?.unwrap().header;
59
60    let operator_address = Address::repeat_byte(0x12);
61    let multicaller_address = Address::repeat_byte(0x78);
62
63    // Set Multicaller code
64    let multicaller_address =
65        MulticallerDeployer::new().set_code(client.clone(), multicaller_address).await?.address().ok_or_eyre("MULTICALLER_NOT_DEPLOYED")?;
66    info!("Multicaller deployed at {:?}", multicaller_address);
67
68    // Cheat : SET Multicaller WETH balance
69    BalanceCheater::set_anvil_token_balance_float(client.clone(), NWETH::ADDRESS, multicaller_address, 1.0).await?;
70
71    // Initialization
72    let cache_db = KabuDBType::default();
73
74    let market_instance = Market::default();
75
76    let market_state_instance = MarketState::new(cache_db.clone());
77
78    let market_instance = SharedState::new(market_instance);
79
80    let market_state_instance = SharedState::new(market_state_instance);
81
82    let abi_encoder = ProtocolABIEncoderV2::default();
83
84    let swap_opcodes_encoder = ProtocolSwapOpcodesEncoderV2::default();
85
86    let swap_line_encoder = SwapLineEncoder::new(multicaller_address, Arc::new(abi_encoder), Arc::new(swap_opcodes_encoder));
87
88    let swap_step_encoder = SwapStepEncoder::new(multicaller_address, swap_line_encoder);
89
90    let swap_encoder = Arc::new(MulticallerSwapEncoder::new(multicaller_address, swap_step_encoder));
91
92    //preload state
93    preload_market_state(client.clone(), vec![multicaller_address], vec![], vec![], market_state_instance.clone(), None).await?;
94
95    //Preloading market
96    preload_pools::<_, _>(client.clone(), market_instance.clone(), market_state_instance.clone()).await?;
97
98    let market = market_instance.read().await;
99
100    // Getting swap directions
101    let pool_address: Address = UniswapV2PoolAddress::WETH_USDT;
102
103    let pool = market.get_pool(&PoolId::Address(pool_address)).ok_or_eyre("POOL_NOT_FOUND")?;
104
105    let swap_directions = pool.get_swap_directions();
106
107    let mut btree_map: BTreeMap<PoolWrapper, Vec<SwapDirection>> = BTreeMap::new();
108
109    btree_map.insert(pool.clone(), swap_directions);
110
111    //let swap_paths = market.build_swap_path_vec(&btree_map)?;
112
113    let swap_paths = market.swap_paths_vec();
114
115    let db = CacheDB::new(market_state_instance.read().await.state_db.clone());
116
117    let in_amount_f64 = 1.0;
118    let in_amount = NWETH::from_float(in_amount_f64);
119
120    let mut gas_used_map: HashMap<SwapLineDTO, u64> = HashMap::new();
121    let mut calldata_map: HashMap<SwapLineDTO, Bytes> = HashMap::new();
122
123    // Make tests
124
125    for swap_path in swap_paths.iter() {
126        if !swap_path.tokens[0].is_weth() {
127            continue;
128        }
129
130        let sp = swap_path.clone();
131        let sp_dto: SwapLineDTO = (&sp).into();
132        println!("Checking {sp_dto}");
133        if let Some(filter) = &cli.filter.clone() {
134            if !format!("{sp_dto}").contains(filter) {
135                println!("Skipping {sp_dto}");
136                continue;
137            }
138        }
139
140        let mut swapline = SwapLine { path: sp, amount_in: SwapAmountType::Set(in_amount), ..SwapLine::default() };
141
142        match swapline.calculate_with_in_amount(&db, in_amount) {
143            Ok((out_amount, gas_used, _)) => {
144                println!("{} gas: {}  amount {} -> {}", sp_dto, gas_used, in_amount_f64, NWETH::to_float(out_amount));
145                swapline.amount_out = SwapAmountType::Set(out_amount)
146            }
147            Err(e) => {
148                error!("calculate_with_in_amount error : {:?}", e);
149            }
150        }
151        let swap = Swap::BackrunSwapLine(swapline);
152
153        let calls = swap_encoder.make_calls(&swap)?;
154        let (to, payload) = swap_encoder.encode_calls(calls)?;
155
156        calldata_map.insert(swap_path.into(), payload.clone());
157
158        let tx_request = TransactionRequest::default().to(to).from(operator_address).input(TransactionInput::new(payload));
159
160        let gas_used = match client.estimate_gas(tx_request).await {
161            Ok(gas_needed) => {
162                //info!("Gas required:  {gas_needed}");
163                gas_needed
164            }
165            Err(e) => {
166                error!("Gas estimation error for {sp_dto}, err={e}");
167                0
168            }
169        };
170
171        gas_used_map.insert(swap_path.into(), gas_used);
172    }
173
174    if let Some(bench_file) = cli.file {
175        if cli.anvil {
176            // Save anvil test data
177            let mut calldata_vec: Vec<(SwapLineDTO, Bytes)> = calldata_map.into_iter().collect();
178            calldata_vec.sort_by(|a, b| a.0.cmp(&b.0));
179            let calldata_vec: Vec<(String, Bytes)> = calldata_vec.into_iter().map(|(k, v)| (k.to_string(), v)).collect();
180            let test_data = create_sol_test(calldata_vec);
181            println!("{test_data}");
182            let mut file = File::create(bench_file).await?;
183            file.write_all(test_data.as_bytes()).await?;
184        } else if cli.save {
185            // Save benchmark results
186            let results: Vec<(SwapLineDTO, u64)> = gas_used_map.into_iter().collect();
187
188            let json_string = serde_json::to_string_pretty(&results)?;
189
190            let mut file = File::create(bench_file).await?;
191            file.write_all(json_string.as_bytes()).await?;
192        } else {
193            // Compare benchmark results
194            let mut file = File::open(bench_file).await?;
195            let mut json_string = String::new();
196            file.read_to_string(&mut json_string).await?;
197
198            let stored_results: Vec<(SwapLineDTO, u64)> = serde_json::from_str(&json_string)?;
199
200            let stored_gas_map: HashMap<SwapLineDTO, u64> = stored_results.clone().into_iter().collect();
201
202            for (current_entry, gas) in gas_used_map.iter() {
203                match stored_gas_map.get(current_entry) {
204                    Some(stored_gas) => {
205                        let change_i: i64 = *gas as i64 - *stored_gas as i64;
206                        let change = format!("{change_i}");
207
208                        let change = if *gas < 40000 {
209                            change.red()
210                        } else {
211                            match change_i {
212                                i if i > 0 => change.red(),
213                                i if i < 0 => change.green(),
214                                _ => change.normal(),
215                            }
216                        };
217
218                        println!("{change} : {current_entry} {gas} - {stored_gas} ",);
219                    }
220                    None => {
221                        if *gas < 40000 {
222                            println!("{} : {} {}", "FAILED".red(), current_entry, gas,);
223                        } else {
224                            println!("{} : {} {}", "NO_DATA".green(), current_entry, gas,);
225                        }
226                    }
227                }
228            }
229        }
230    }
231
232    Ok(())
233}
gasbench/main.rs

gasbench/
main.rs
