1#![allow(clippy::type_complexity)]
2
3use alloy_primitives::map::HashMap;
4use alloy_primitives::{Address, U256};
5use eyre::{OptionExt, Result, eyre};
6use std::collections::BTreeMap;
7use std::fmt::Display;
8use std::sync::Arc;
9use tracing::debug;
10
11use crate::{PoolClass, PoolId, PoolWrapper, SwapDirection, SwapPath, SwapPaths, Token, build_swap_path_vec};
12
13#[derive(Default, Clone)]
16pub struct Market {
17 pools: HashMap<PoolId, PoolWrapper>,
19 pools_disabled: HashMap<PoolId, bool>,
21 pools_manager_cells: HashMap<Address, HashMap<U256, PoolId>>,
23 tokens: HashMap<Address, Arc<Token>>,
25 token_symbols: HashMap<String, Address>,
27
28 token_tokens: HashMap<Address, Vec<Address>>,
30 token_token_pools: HashMap<Address, HashMap<Address, Vec<PoolId>>>,
32 token_pools: HashMap<Address, Vec<PoolId>>,
34 swap_paths: SwapPaths,
36}
37
38impl Display for Market {
39 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40 let token_token_len = self.token_tokens.values().map(|inner| inner.len()).sum::<usize>();
41 let token_token_pools_len = self.token_token_pools.values().map(|inner_map| inner_map.len()).sum::<usize>();
42 let token_pool_len = self.token_pools.values().map(|inner| inner.len()).sum::<usize>();
43 let token_pool_len_max = self.token_pools.values().map(|inner| inner.len()).max().unwrap_or_default();
44 let swap_path_len = self.swap_paths.len();
45 let swap_path_len_max = self.swap_paths.len_max();
46
47 write!(
48 f,
49 "Pools: {} Disabled : {} Tokens : {} TT : {} TTP {} TP {}/{} SwapPaths: {}/{}",
50 self.pools.len(),
51 self.pools_disabled.len(),
52 self.tokens.len(),
53 token_token_len,
54 token_token_pools_len,
55 token_pool_len,
56 token_pool_len_max,
57 swap_path_len,
58 swap_path_len_max
59 )
60 }
61}
62
63impl Market {
64 pub fn add_token<T: Into<Arc<Token>>>(&mut self, token: T) {
66 let arc_token: Arc<Token> = token.into();
67 self.token_symbols.insert(arc_token.get_symbol(), arc_token.get_address());
68 self.tokens.insert(arc_token.get_address(), arc_token);
69 }
70
71 #[inline]
73 pub fn is_basic_token(&self, address: &Address) -> bool {
74 self.tokens.get(address).is_some_and(|t| t.is_basic())
75 }
76
77 #[inline]
78 pub fn is_weth_token(&self, address: &Address) -> bool {
79 self.tokens.get(address).is_some_and(|t| t.is_weth())
80 }
81
82 #[inline]
84 pub fn get_token_or_default(&self, address: &Address) -> Arc<Token> {
85 self.tokens.get(address).map_or(Arc::new(Token::new(*address)), |t| t.clone())
86 }
87
88 #[inline]
94 pub fn get_token(&self, address: &Address) -> Option<Arc<Token>> {
95 self.tokens.get(address).cloned()
96 }
97
98 #[inline]
100 pub fn get_token_by_symbol(&self, symbol: &String) -> Option<Arc<Token>> {
101 self.token_symbols.get(symbol).and_then(|address| self.tokens.get(address).cloned())
102 }
103
104 pub fn add_pool<T: Into<PoolWrapper>>(&mut self, pool: T) -> Result<()> {
106 let pool_contract = pool.into();
107 let pool_address = pool_contract.get_pool_id();
108
109 if let Some(pool) = self.pools.get(&pool_address) {
110 return Err(eyre!("Pool already exists {:?}", pool.get_pool_id()));
111 }
112
113 debug!("Adding pool {:?}", pool_address);
114
115 for swap_direction in pool_contract.get_swap_directions().into_iter() {
116 self.token_token_pools.entry(*swap_direction.from()).or_default().entry(*swap_direction.to()).or_default().push(pool_address);
117 self.token_tokens.entry(*swap_direction.from()).or_default().push(*swap_direction.to());
118 self.token_pools.entry(*swap_direction.from()).or_default().push(pool_address);
120 }
121
122 self.pools.insert(pool_address, pool_contract);
123
124 Ok(())
125 }
126
127 pub fn add_paths(&mut self, paths: Vec<SwapPath>) -> Vec<usize> {
129 paths.into_iter().filter_map(|path| self.swap_paths.add(path)).collect()
130 }
131
132 #[inline]
134 pub fn get_pool_paths(&self, pool_address: &PoolId) -> Option<Vec<SwapPath>> {
135 self.swap_paths.get_pool_paths_enabled_vec(pool_address)
136 }
137
138 #[inline]
140 pub fn swap_paths_vec(&self) -> Vec<SwapPath> {
141 self.swap_paths.paths.clone().into_iter().collect::<Vec<_>>()
142 }
143
144 #[inline]
145 pub fn swap_paths_vec_by_idx(&self, swap_path_idx_vec: Vec<usize>) -> Vec<SwapPath> {
146 swap_path_idx_vec.into_iter().filter_map(|idx| self.swap_paths.paths.get(idx).cloned()).collect::<Vec<_>>()
147 }
148
149 #[inline]
150 pub fn pool_swap_paths_idx_vec(&self, pool_id: &PoolId) -> Option<Vec<usize>> {
151 self.swap_paths.pool_paths.get(pool_id).cloned()
152 }
153
154 #[inline]
155 pub fn pool_swap_paths_vec(&self, pool_id: &PoolId) -> Vec<(usize, SwapPath)> {
156 let pool_paths = self.swap_paths.pool_paths.get(pool_id).cloned().unwrap_or_default();
157
158 pool_paths.into_iter().filter_map(|idx| self.swap_paths.paths.get(idx).cloned().map(|a| (idx, a))).collect::<Vec<_>>()
159 }
160
161 #[inline]
163 pub fn get_pool(&self, address: &PoolId) -> Option<&PoolWrapper> {
164 self.pools.get(address).filter(|&pool_wrapper| pool_wrapper.get_class() != PoolClass::Unknown)
165 }
166
167 #[inline]
169 pub fn is_pool(&self, address: &PoolId) -> bool {
170 self.pools.contains_key(address)
171 }
172
173 #[inline]
175 pub fn pools(&self) -> &HashMap<PoolId, PoolWrapper> {
176 &self.pools
177 }
178
179 pub fn swap_paths(&self) -> &SwapPaths {
180 &self.swap_paths
181 }
182
183 pub fn swap_paths_mut(&mut self) -> &mut SwapPaths {
184 &mut self.swap_paths
185 }
186
187 pub fn set_pool_disabled(&mut self, address: &PoolId, token_from: &Address, token_to: &Address, disabled: bool) {
189 self.swap_paths.disable_pool_paths(address, token_from, token_to, disabled);
210 }
211
212 pub fn set_path_disabled(&mut self, swap_path: &SwapPath, disabled: bool) -> bool {
214 self.swap_paths.disable_path(swap_path, disabled)
215 }
216
217 #[inline]
219 pub fn is_pool_disabled(&self, address: &PoolId) -> bool {
220 self.pools_disabled.get(address).is_some_and(|&is_disabled| is_disabled)
221 }
222
223 #[inline]
225 pub fn get_token_token_pools(&self, token_from_address: &Address, token_to_address: &Address) -> Option<&Vec<PoolId>> {
226 self.token_token_pools.get(token_from_address).and_then(|a| a.get(token_to_address))
227 }
228
229 #[inline]
231 pub fn get_token_tokens(&self, token_from_address: &Address) -> Option<&Vec<Address>> {
232 self.token_tokens.get(token_from_address)
233 }
234
235 pub fn get_token_pools(&self, token_address: &Address) -> Option<&Vec<PoolId>> {
237 self.token_pools.get(token_address)
238 }
239
240 pub fn get_token_pools_len(&self, token_address: &Address) -> usize {
242 self.token_pools.get(token_address).map_or(0, |t| t.len())
243 }
244 pub fn build_swap_path_vec(&self, directions: &BTreeMap<PoolWrapper, Vec<SwapDirection>>) -> Result<Vec<SwapPath>> {
246 build_swap_path_vec(self, directions)
247 }
248
249 pub fn swap_path(&self, token_address_vec: Vec<Address>, pool_address_vec: Vec<PoolId>) -> Result<SwapPath> {
251 let mut tokens: Vec<Arc<Token>> = Vec::new();
252 let mut pools: Vec<PoolWrapper> = Vec::new();
253
254 for token_address in token_address_vec.iter() {
255 tokens.push(self.get_token(token_address).ok_or_eyre("TOKEN_NOT_FOUND")?);
256 }
257 for pool_address in pool_address_vec.iter() {
258 pools.push(self.get_pool(pool_address).cloned().ok_or_eyre("TOKEN_NOT_FOUND")?);
259 }
260
261 Ok(SwapPath { tokens, pools, ..Default::default() })
262 }
263
264 pub fn disabled_pools_count(&self) -> usize {
265 self.pools_disabled.len()
266 }
267
268 pub fn add_pool_manager_cell(&mut self, pool_manager_address: Address, pool_id: PoolId, cell: U256) {
269 let pool_manager_entry = self.pools_manager_cells.entry(pool_manager_address).or_default();
270 pool_manager_entry.insert(cell, pool_id);
271 }
272
273 pub fn is_pool_manager(&self, pool_manager_address: &Address) -> bool {
274 self.pools_manager_cells.contains_key(pool_manager_address)
275 }
276
277 pub fn get_pool_id_for_cell(&self, pool_manager_address: &Address, cell: &U256) -> Option<&PoolId> {
278 self.pools_manager_cells.get(pool_manager_address).and_then(|pool_manager_cell| pool_manager_cell.get(cell))
279 }
280}
281
282#[cfg(test)]
283mod tests {
284 use super::*;
285 use crate::mock_pool::MockPool;
286 use alloy_primitives::Address;
287 use eyre::Result;
288 use kabu_defi_address_book::TokenAddressEth;
289
290 #[test]
291 fn test_add_pool() {
292 let mut market = Market::default();
293 let pool_address = Address::random();
294 let token0 = Address::random();
295 let token1 = Address::random();
296 let mock_pool = MockPool::new(token0, token1, pool_address);
297
298 let result = market.add_pool(mock_pool);
299
300 assert!(result.is_ok());
301
302 assert_eq!(market.get_pool(&PoolId::Address(pool_address)).unwrap().pool.get_address(), PoolId::Address(pool_address));
303
304 assert_eq!(*market.get_token_token_pools(&token0, &token1).unwrap().first().unwrap(), PoolId::Address(pool_address));
305 assert_eq!(*market.get_token_token_pools(&token1, &token0).unwrap().first().unwrap(), PoolId::Address(pool_address));
306
307 assert!(market.get_token_tokens(&token0).unwrap().contains(&token1));
308 assert!(market.get_token_tokens(&token1).unwrap().contains(&token0));
309
310 assert!(market.get_token_pools(&token0).unwrap().contains(&PoolId::Address(pool_address)));
311 assert!(market.get_token_pools(&token1).unwrap().contains(&PoolId::Address(pool_address)));
312 }
313
314 #[ignore]
315 #[test]
316 fn test_add_token() {
317 let market = Market::default();
318 let token_address = Address::random();
319
320 assert_eq!(market.get_token(&token_address).unwrap().get_address(), token_address);
321 }
322
323 #[test]
324 fn test_get_token_default() {
325 let market = Market::default();
326 let token_address = Address::random();
327
328 let token = market.get_token_or_default(&token_address);
329
330 assert_eq!(token.get_address(), token_address);
331 }
332
333 #[test]
334 fn test_get_pool() {
335 let mut market = Market::default();
336 let pool_address = Address::random();
337 let mock_pool = MockPool::new(Address::ZERO, Address::ZERO, pool_address);
338 let _ = market.add_pool(mock_pool.clone());
339
340 let pool = market.get_pool(&PoolId::Address(pool_address));
341
342 assert_eq!(pool.unwrap().get_address(), PoolId::Address(pool_address));
343 }
344
345 #[test]
346 fn test_is_pool() {
347 let mut market = Market::default();
348 let pool_address = Address::random();
349 let mock_pool = MockPool::new(Address::ZERO, Address::ZERO, pool_address);
350 let _ = market.add_pool(mock_pool.clone());
351
352 let is_pool = market.is_pool(&PoolId::Address(pool_address));
353
354 assert!(is_pool);
355 }
356
357 #[test]
358 fn test_is_pool_not_found() {
359 let market = Market::default();
360 let pool_address = Address::random();
361
362 let is_pool = market.is_pool(&PoolId::Address(pool_address));
363
364 assert!(!is_pool);
365 }
366
367 #[ignore]
368 #[test]
369 fn test_set_pool_disabled() {
370 let mut market = Market::default();
371 let pool_address = Address::random();
372 let token0 = Address::random();
373 let token1 = Address::random();
374 let mock_pool = MockPool::new(token0, token1, pool_address);
375 let _ = market.add_pool(mock_pool.clone());
376
377 assert!(!market.is_pool_disabled(&PoolId::Address(pool_address)));
378 assert_eq!(market.get_token_token_pools(&token0, &token1).unwrap().len(), 1);
379
380 market.set_pool_disabled(&PoolId::Address(pool_address), &token0, &token1, true);
382 assert!(market.is_pool_disabled(&PoolId::Address(pool_address)));
383 assert_eq!(market.get_token_token_pools(&token0, &token1).unwrap().len(), 1);
384
385 market.set_pool_disabled(&PoolId::Address(pool_address), &token0, &token1, false);
387 assert!(!market.is_pool_disabled(&PoolId::Address(pool_address)));
388 assert_eq!(market.get_token_token_pools(&token0, &token1).unwrap().len(), 1);
389 }
390
391 #[test]
392 fn test_get_token_token_pools() {
393 let mut market = Market::default();
394 let pool_address = Address::random();
395 let token0 = Address::random();
396 let token1 = Address::random();
397 let mock_pool = MockPool::new(token0, token1, pool_address);
398 let _ = market.add_pool(mock_pool);
399
400 let pools = market.get_token_token_pools(&token0, &token1);
401
402 assert_eq!(pools.unwrap().first().unwrap(), &PoolId::Address(pool_address));
403 }
404
405 #[test]
406 fn test_get_token_tokens() {
407 let mut market = Market::default();
408 let pool_address = Address::random();
409 let token0 = Address::random();
410 let token1 = Address::random();
411 let mock_pool = MockPool::new(token0, token1, pool_address);
412 let _ = market.add_pool(mock_pool);
413
414 let tokens = market.get_token_tokens(&token0);
415
416 assert_eq!(tokens.unwrap().first().unwrap(), &token1);
417 }
418
419 #[test]
420 fn test_get_token_pools() {
421 let mut market = Market::default();
422 let pool_address = Address::random();
423 let token0 = Address::random();
424 let token1 = Address::random();
425 let mock_pool = MockPool::new(token0, token1, pool_address);
426 let _ = market.add_pool(mock_pool);
427
428 let pools = market.get_token_pools(&token0).cloned();
429
430 assert_eq!(pools.unwrap().first().unwrap(), &PoolId::Address(pool_address));
431 }
432
433 #[test]
434 fn test_build_swap_path_vec_two_hops() -> Result<()> {
435 let mut market = Market::default();
436
437 let weth_token = Token::new_with_data(TokenAddressEth::WETH, Some("WETH".to_string()), None, Some(18), true, false);
439 market.add_token(weth_token);
440
441 let pool_address1 = Address::random();
443 let token1 = Address::random();
444 let mock_pool1 = PoolWrapper::new(Arc::new(MockPool::new(TokenAddressEth::WETH, token1, pool_address1)));
445 let _ = market.add_pool(mock_pool1.clone());
446
447 let pool_address2 = Address::random();
449 let mock_pool2 = PoolWrapper::new(Arc::new(MockPool::new(TokenAddressEth::WETH, token1, pool_address2)));
450 let _ = market.add_pool(mock_pool2.clone());
451
452 let mut directions = BTreeMap::new();
454 directions.insert(mock_pool2.clone(), mock_pool2.get_swap_directions());
455 let swap_paths = market.build_swap_path_vec(&directions)?;
456
457 assert_eq!(swap_paths.len(), 2);
459 assert_eq!(swap_paths.first().unwrap().pool_count(), 2);
460 assert_eq!(swap_paths.first().unwrap().tokens_count(), 3);
461 assert_eq!(swap_paths.get(1).unwrap().pool_count(), 2);
462 assert_eq!(swap_paths.get(1).unwrap().tokens_count(), 3);
463
464 let (first_path, second_path) =
466 if swap_paths.first().unwrap().pools.first().unwrap().get_address() == PoolId::Address(pool_address1) {
467 (swap_paths.first().unwrap(), swap_paths.get(1).unwrap())
468 } else {
469 (swap_paths.get(1).unwrap(), swap_paths.first().unwrap())
470 };
471
472 let tokens = first_path.tokens.iter().map(|token| token.get_address()).collect::<Vec<Address>>();
474 assert_eq!(tokens.first(), Some(&TokenAddressEth::WETH));
475 assert_eq!(tokens.get(1), Some(&token1));
476 assert_eq!(tokens.get(2), Some(&TokenAddressEth::WETH));
477
478 let pools = first_path.pools.iter().map(|pool| pool.get_address()).collect::<Vec<PoolId>>();
479 assert_eq!(pools.first(), Some(&PoolId::Address(pool_address1)));
480 assert_eq!(pools.get(1), Some(&PoolId::Address(pool_address2)));
481
482 let tokens = second_path.tokens.iter().map(|token| token.get_address()).collect::<Vec<Address>>();
484 assert_eq!(tokens.first(), Some(&TokenAddressEth::WETH));
485 assert_eq!(tokens.get(1), Some(&token1));
486 assert_eq!(tokens.get(2), Some(&TokenAddressEth::WETH));
487
488 let pools = second_path.pools.iter().map(|pool| pool.get_address()).collect::<Vec<PoolId>>();
489 assert_eq!(pools.first(), Some(&PoolId::Address(pool_address2)));
490 assert_eq!(pools.get(1), Some(&PoolId::Address(pool_address1)));
491
492 Ok(())
493 }
494
495 #[test]
496 fn test_build_swap_path_vec_three_hops() -> Result<()> {
497 let mut market = Market::default();
498
499 let weth_token = Token::new_with_data(TokenAddressEth::WETH, Some("WETH".to_string()), None, Some(18), true, false);
501 market.add_token(weth_token);
502
503 let token1 = Address::random();
505 let token2 = Address::random();
506
507 let pool_address1 = Address::random();
509 let mock_pool = PoolWrapper::new(Arc::new(MockPool::new(token1, TokenAddressEth::WETH, pool_address1)));
510 let _ = market.add_pool(mock_pool);
511
512 let pool_address2 = Address::random();
514 let mock_pool2 = PoolWrapper::new(Arc::new(MockPool::new(token1, token2, pool_address2)));
515 let _ = market.add_pool(mock_pool2);
516
517 let pool_address3 = Address::random();
519 let mock_pool3 = PoolWrapper::new(Arc::new(MockPool::new(token2, TokenAddressEth::WETH, pool_address3)));
520 let _ = market.add_pool(mock_pool3.clone());
521
522 let mut directions = BTreeMap::new();
524 directions.insert(mock_pool3.clone(), mock_pool3.get_swap_directions());
525 let swap_paths = market.build_swap_path_vec(&directions)?;
526
527 assert_eq!(swap_paths.len(), 2);
529 assert_eq!(swap_paths.first().unwrap().pool_count(), 3);
530 assert_eq!(swap_paths.first().unwrap().tokens_count(), 4);
531 assert_eq!(swap_paths.get(1).unwrap().pool_count(), 3);
532 assert_eq!(swap_paths.get(1).unwrap().tokens_count(), 4);
533
534 let (first_path, second_path) = if swap_paths.first().unwrap().tokens.get(1).unwrap().get_address() == token1 {
536 (swap_paths.first().unwrap(), swap_paths.get(1).unwrap())
537 } else {
538 (swap_paths.get(1).unwrap(), swap_paths.first().unwrap())
539 };
540
541 let tokens = first_path.tokens.iter().map(|token| token.get_address()).collect::<Vec<Address>>();
543 assert_eq!(tokens.first(), Some(&TokenAddressEth::WETH));
544 assert_eq!(tokens.get(1), Some(&token1));
545 assert_eq!(tokens.get(2), Some(&token2));
546 assert_eq!(tokens.get(3), Some(&TokenAddressEth::WETH));
547
548 let pools = first_path.pools.iter().map(|pool| pool.get_address()).collect::<Vec<PoolId>>();
549 assert_eq!(pools.first(), Some(&PoolId::Address(pool_address1)));
550 assert_eq!(pools.get(1), Some(&PoolId::Address(pool_address2)));
551 assert_eq!(pools.get(2), Some(&PoolId::Address(pool_address3)));
552
553 let tokens = second_path.tokens.iter().map(|token| token.get_address()).collect::<Vec<Address>>();
555 assert_eq!(tokens.first(), Some(&TokenAddressEth::WETH));
556 assert_eq!(tokens.get(1), Some(&token2));
557 assert_eq!(tokens.get(2), Some(&token1));
558 assert_eq!(tokens.get(3), Some(&TokenAddressEth::WETH));
559
560 let pools = second_path.pools.iter().map(|pool| pool.get_address()).collect::<Vec<PoolId>>();
561 assert_eq!(pools.first(), Some(&PoolId::Address(pool_address3)));
562 assert_eq!(pools.get(1), Some(&PoolId::Address(pool_address2)));
563 assert_eq!(pools.get(2), Some(&PoolId::Address(pool_address1)));
564
565 Ok(())
566 }
567}