kabu_strategy_backrun/
affected_pools_code.rs

1use alloy_eips::BlockNumberOrTag;
2use alloy_evm::EvmEnv;
3use alloy_network::Network;
4use alloy_provider::Provider;
5use eyre::eyre;
6use tokio::sync::RwLock;
7
8use kabu_defi_pools::protocols::{UniswapV2Protocol, UniswapV3Protocol};
9use kabu_defi_pools::state_readers::UniswapV3EvmStateReader;
10use kabu_defi_pools::{MaverickPool, PancakeV3Pool, UniswapV2Pool, UniswapV3Pool};
11use kabu_evm_db::{AlloyDB, KabuDB};
12use kabu_types_blockchain::GethStateUpdateVec;
13use kabu_types_market::MarketState;
14use kabu_types_market::SwapDirection;
15use kabu_types_market::{Market, Pool, PoolId, PoolProtocol, PoolWrapper, get_protocol_by_factory};
16use std::collections::BTreeMap;
17use std::sync::Arc;
18use tracing::{debug, error};
19
20pub async fn get_affected_pools_from_code<P, N>(
21    client: P,
22    market: Arc<RwLock<Market>>,
23    state_update: &GethStateUpdateVec,
24    evm_env: &EvmEnv,
25) -> eyre::Result<BTreeMap<PoolWrapper, Vec<SwapDirection>>>
26where
27    N: Network,
28    P: Provider<N> + Send + Sync + Clone + 'static,
29{
30    let mut market_state = MarketState::new(KabuDB::new());
31
32    market_state.state_db.apply_geth_state_update(state_update, true, false);
33
34    let mut ret: BTreeMap<PoolWrapper, Vec<SwapDirection>> = BTreeMap::new();
35
36    for state_update_record in state_update.iter() {
37        for (address, state_update_entry) in state_update_record.iter() {
38            if let Some(code) = &state_update_entry.code {
39                if UniswapV2Protocol::is_code(code) {
40                    match market.read().await.get_pool(&PoolId::Address(*address)) {
41                        None => {
42                            debug!(?address, "Loading UniswapV2 class pool");
43
44                            let ext_db = AlloyDB::new(client.clone(), BlockNumberOrTag::Latest.into());
45
46                            let Some(ext_db) = ext_db else {
47                                error!("Cannot create AlloyDB");
48                                continue;
49                            };
50
51                            let state_db = market_state.state_db.clone().with_ext_db(ext_db);
52
53                            match UniswapV3EvmStateReader::factory(&state_db, evm_env, *address) {
54                                Ok(_factory_address) => match UniswapV2Pool::fetch_pool_data_evm(&state_db, evm_env, *address) {
55                                    Ok(pool) => {
56                                        let pool = PoolWrapper::new(Arc::new(pool));
57                                        let protocol = pool.get_protocol();
58                                        let swap_directions = pool.get_swap_directions();
59
60                                        debug!(%address, %protocol, ?swap_directions, "UniswapV2 pool loaded");
61                                        ret.insert(pool, swap_directions);
62                                    }
63                                    Err(err) => {
64                                        error!(?address, %err, "Error loading UniswapV2 pool");
65                                    }
66                                },
67                                Err(err) => {
68                                    error!(?address, %err, "Error loading UniswapV2 factory for pool")
69                                }
70                            }
71                        }
72                        Some(pool) => {
73                            debug!(?address, protocol = ?pool.get_protocol(), "Pool already exists");
74                        }
75                    }
76                }
77
78                if UniswapV3Protocol::is_code(code) {
79                    match market.read().await.get_pool(&PoolId::Address(*address)) {
80                        None => {
81                            debug!(%address, "Loading UniswapV3 class pool");
82
83                            let ext_db = AlloyDB::new(client.clone(), BlockNumberOrTag::Latest.into());
84
85                            let Some(ext_db) = ext_db else {
86                                error!("Cannot create AlloyDB");
87                                continue;
88                            };
89
90                            let state_db = market_state.state_db.clone().with_ext_db(ext_db);
91
92                            match UniswapV3EvmStateReader::factory(&state_db, evm_env, *address) {
93                                Ok(factory_address) => {
94                                    match get_protocol_by_factory(factory_address) {
95                                        PoolProtocol::PancakeV3 => {
96                                            let pool = PancakeV3Pool::fetch_pool_data_evm(&state_db, evm_env, *address);
97                                            match pool {
98                                                Ok(pool) => {
99                                                    let swap_directions = pool.get_swap_directions();
100                                                    let protocol = pool.get_protocol();
101                                                    debug!(?address, %protocol, ?swap_directions, "PancakeV3 Pool loaded");
102                                                    ret.insert(PoolWrapper::new(Arc::new(pool)), swap_directions);
103                                                }
104                                                Err(err) => {
105                                                    error!(?address, %err, "Error loading PancakeV3 pool");
106                                                }
107                                            }
108                                        }
109                                        PoolProtocol::Maverick => {
110                                            let pool = MaverickPool::fetch_pool_data_evm(&state_db, evm_env, *address);
111                                            match pool {
112                                                Ok(pool) => {
113                                                    let pool = PoolWrapper::new(Arc::new(pool));
114                                                    let swap_directions = pool.get_swap_directions();
115                                                    let protocol = pool.get_protocol();
116                                                    debug!(?address, %protocol, ?swap_directions, "Maverick Pool loaded");
117
118                                                    ret.insert(pool, swap_directions);
119                                                }
120                                                Err(err) => {
121                                                    error!(?address, %err, "Error loading Maverick pool");
122                                                }
123                                            }
124                                        }
125                                        _ => match UniswapV3Pool::fetch_pool_data_evm(&state_db, evm_env, *address) {
126                                            Ok(pool) => {
127                                                let pool = PoolWrapper::new(Arc::new(pool));
128                                                let swap_directions = pool.get_swap_directions();
129                                                let protocol = pool.get_protocol();
130                                                debug!(
131                                                    %address,
132                                                    %protocol,
133                                                    ?swap_directions,
134                                                    "UniswapV3 Pool loaded"
135                                                );
136                                                ret.insert(pool, swap_directions);
137                                            }
138                                            Err(err) => {
139                                                error!(%address, %err, "Error loading UniswapV3 pool");
140                                            }
141                                        },
142                                    };
143                                }
144                                Err(err) => {
145                                    error!(?address, %err, "Error loading UniswapV3 factory for pool")
146                                }
147                            }
148                        }
149                        Some(pool) => {
150                            debug!(?address, protocol = ?pool.get_protocol(), "Pool already exists")
151                        }
152                    }
153                }
154            }
155        }
156    }
157    if !ret.is_empty() { Ok(ret) } else { Err(eyre!("NO_POOLS_LOADED")) }
158}
159
160/// Check if the state update code contains code for a UniswapV2 pair or UniswapV3 pool by looking for method signatures.
161pub fn is_pool_code(state_update: &GethStateUpdateVec) -> bool {
162    for state_update_record in state_update.iter() {
163        for (_address, state_update_entry) in state_update_record.iter() {
164            if let Some(code) = &state_update_entry.code {
165                if UniswapV3Protocol::is_code(code) {
166                    return true;
167                }
168                if UniswapV2Protocol::is_code(code) {
169                    return true;
170                }
171            }
172        }
173    }
174    false
175}
kabu_strategy_backrun/affected_pools_code.rs

kabu_strategy_backrun/
affected_pools_code.rs
