1use num_bigint::{BigInt, BigUint};
9use rustc_hash::FxHashMap;
10use tracing::{instrument, trace};
11use tycho_simulation::{
12 tycho_common::simulation::protocol_sim::ProtocolSim,
13 tycho_core::models::{token::Token, Address},
14};
15
16use super::{most_liquid::DepthAndPrice, NoPathReason};
17use crate::{
18 algorithm::sim_guard::GuardedProtocolSim,
19 derived::types::TokenGasPrices,
20 feed::market_data::{MarketData, MarketDataView, MarketState},
21 graph::{GraphError, Path, RouteSearch, TokenPath, TopologyGraph},
22 types::{ComponentId, Route, RouteResult, Swap},
23 AlgorithmError, StateLabel,
24};
25
26#[instrument(level = "debug", skip(graph, search))]
35pub fn find_paths<'a, D>(
36 graph: &'a TopologyGraph<D>,
37 from: &Address,
38 to: &Address,
39 search: RouteSearch<'_>,
40 max_paths_per_sequence: Option<usize>,
41) -> Result<Vec<Path<'a, D>>, AlgorithmError> {
42 let token_paths = find_token_paths(graph, from, to, search)?;
43
44 let mut paths = Vec::new();
45 for token_path in &token_paths {
46 paths.extend(graph.expand_path(token_path, max_paths_per_sequence, search.exclusions));
47 }
48
49 Ok(paths)
50}
51
52pub fn find_token_paths<D>(
61 graph: &TopologyGraph<D>,
62 from: &Address,
63 to: &Address,
64 search: RouteSearch<'_>,
65) -> Result<Vec<TokenPath>, AlgorithmError> {
66 graph
67 .paths_between(from, to, search)
68 .map_err(|error| AlgorithmError::NoPath {
69 from: from.clone(),
70 to: to.clone(),
71 reason: match error {
72 GraphError::TokenNotFound(ref missing) if missing == from => {
73 NoPathReason::SourceTokenNotInGraph
74 }
75 _ => NoPathReason::DestinationTokenNotInGraph,
76 },
77 })
78}
79
80pub fn try_score_path(path: &Path<DepthAndPrice>) -> Option<f64> {
91 if path.is_empty() {
92 trace!("cannot score empty path");
93 return None;
94 }
95
96 let mut price = 1.0;
97 let mut min_depth = f64::MAX;
98
99 for edge in path.edge_iter() {
100 let Some(data) = edge.data.as_ref() else {
101 trace!(component_id = %edge.component_id, "edge missing weight data, path cannot be scored");
102 return None;
103 };
104
105 price *= data.spot_price;
106 min_depth = min_depth.min(data.depth);
107 }
108
109 Some(price * min_depth)
110}
111
112#[instrument(level = "trace", skip(path, market, token_prices), fields(hop_count = path.len()))]
125pub fn simulate_pool_path<D>(
126 path: &Path<D>,
127 market: &MarketState,
128 token_prices: Option<&TokenGasPrices>,
129 amount_in: BigUint,
130) -> Result<RouteResult, AlgorithmError> {
131 let mut current_amount = amount_in.clone();
132 let mut swaps = Vec::with_capacity(path.len());
133
134 let mut state_overrides: FxHashMap<&ComponentId, Box<dyn ProtocolSim>> = FxHashMap::default();
136 let mut tokens: FxHashMap<Address, Token> = FxHashMap::default();
137
138 for (address_in, edge_data, address_out) in path.iter() {
139 let token_in = get_token(market, address_in)?;
141 let token_out = get_token(market, address_out)?;
142
143 let component_id = &edge_data.component_id;
144 let component = market
145 .get_component(component_id)
146 .ok_or_else(|| AlgorithmError::DataNotFound {
147 kind: "component",
148 id: Some(component_id.clone()),
149 })?;
150 let component_state = market
151 .get_simulation_state(component_id)
152 .ok_or_else(|| AlgorithmError::DataNotFound {
153 kind: "simulation state",
154 id: Some(component_id.clone()),
155 })?;
156
157 let state = state_overrides
158 .get(component_id)
159 .map(Box::as_ref)
160 .unwrap_or(component_state);
161
162 let result = state
164 .get_amount_out_guarded(current_amount.clone(), token_in, token_out)
165 .map_err(|e| AlgorithmError::Other(format!("simulation error: {:?}", e)))?;
166
167 swaps.push(Swap::new(
169 component_id.clone(),
170 component.protocol_system.clone(),
171 token_in.address.clone(),
172 token_out.address.clone(),
173 current_amount.clone(),
174 result.amount.clone(),
175 result.gas,
176 component.clone(),
177 state.clone_box(),
178 ));
179 tokens
180 .entry(token_in.address.clone())
181 .or_insert_with(|| token_in.clone());
182 tokens
183 .entry(token_out.address.clone())
184 .or_insert_with(|| token_out.clone());
185
186 state_overrides.insert(component_id, result.new_state);
187 current_amount = result.amount;
188 }
189
190 let route = Route::new(swaps, tokens)?;
192 let output_amount = route
193 .swaps()
194 .last()
195 .map(|s| s.amount_out().clone())
196 .unwrap_or_else(|| BigUint::ZERO);
197
198 let gas_price = market
199 .gas_price()
200 .ok_or(AlgorithmError::DataNotFound { kind: "gas price", id: None })?
201 .effective_gas_price()
202 .clone();
203
204 let net_amount_out = if let Some(last_swap) = route.swaps().last() {
205 let gas_cost_wei = route.total_gas() * &gas_price;
206
207 match token_prices.and_then(|prices| prices.get(last_swap.token_out())) {
210 Some(price) => {
211 BigInt::from(output_amount) -
212 BigInt::from(gas_cost_wei * &price.numerator / &price.denominator)
213 }
214 None => BigInt::from(output_amount),
215 }
216 } else {
217 BigInt::from(output_amount)
218 };
219
220 Ok(RouteResult::new(route, net_amount_out, gas_price))
221}
222
223pub fn get_token<'a>(
229 market: &'a MarketState,
230 address: &Address,
231) -> Result<&'a Token, AlgorithmError> {
232 market
233 .get_token(address)
234 .ok_or_else(|| AlgorithmError::DataNotFound {
235 kind: "token",
236 id: Some(format!("{:?}", address)),
237 })
238}
239
240pub async fn read_market(
246 market: &MarketData,
247 label: Option<StateLabel>,
248) -> Result<MarketDataView<'_>, AlgorithmError> {
249 match label.as_ref() {
250 Some(l) => market
251 .read_labeled(l)
252 .await
253 .map_err(|e| AlgorithmError::Other(e.to_string())),
254 None => Ok(market.read().await),
255 }
256}
257
258pub fn fetch_gas_price(view: &MarketState) -> Result<BigUint, AlgorithmError> {
264 Ok(view
265 .gas_price()
266 .ok_or(AlgorithmError::DataNotFound { kind: "gas price", id: None })?
267 .effective_gas_price()
268 .clone())
269}
270
271#[cfg(test)]
272mod tests {
273 use num_bigint::BigUint;
274 use rstest::rstest;
275
276 use super::*;
277 use crate::{
278 algorithm::test_utils::{
279 addr,
280 fixtures::{addrs, linear_graph},
281 market_read, setup_market_weighted, token, MockProtocolSim,
282 },
283 graph::{GraphManager, GraphQueryFilter, TopologyGraphManager},
284 types::RouteExclusions,
285 };
286
287 fn hops(min_hops: usize, max_hops: usize) -> GraphQueryFilter {
289 GraphQueryFilter { min_hops, max_hops, connector_tokens: None }
290 }
291
292 #[test]
293 fn test_try_score_path_calculates_correctly() {
294 let (a, b, c, _) = addrs();
295 let mut m = linear_graph();
296
297 m.set_pool_weight(&"ab".to_string(), &a, &b, DepthAndPrice::new(2.0, 1000.0), false)
299 .unwrap();
300 m.set_pool_weight(&"bc".to_string(), &b, &c, DepthAndPrice::new(0.5, 500.0), false)
301 .unwrap();
302
303 let graph = m.graph();
304 let paths = find_paths(
305 graph,
306 &a,
307 &c,
308 RouteSearch { bounds: &hops(2, 2), exclusions: &RouteExclusions::default() },
309 None,
310 )
311 .unwrap();
312 assert_eq!(paths.len(), 1);
313 let path = &paths[0];
314
315 let expected = 2.0 * 0.5 * 500.0;
317 let score = try_score_path(path).unwrap();
318 assert_eq!(score, expected, "expected {expected}, got {score}");
319 }
320
321 #[test]
322 fn test_try_score_path_empty_returns_none() {
323 let path: Path<DepthAndPrice> = Path::new();
324 assert_eq!(try_score_path(&path), None);
325 }
326
327 #[test]
328 fn test_try_score_path_missing_weight_returns_none() {
329 let (a, b, _, _) = addrs();
330 let m = linear_graph();
331 let graph = m.graph();
332 let paths = find_paths(
333 graph,
334 &a,
335 &b,
336 RouteSearch { bounds: &hops(1, 1), exclusions: &RouteExclusions::default() },
337 None,
338 )
339 .unwrap();
340 assert_eq!(paths.len(), 1);
341 assert!(try_score_path(&paths[0]).is_none());
342 }
343
344 #[test]
345 fn test_try_score_path_circular_route() {
346 let (a, b, _, _) = addrs();
348 let mut m = linear_graph();
349
350 m.set_pool_weight(&"ab".to_string(), &a, &b, DepthAndPrice::new(2.0, 1000.0), false)
354 .unwrap();
355 m.set_pool_weight(&"ab".to_string(), &b, &a, DepthAndPrice::new(0.6, 800.0), false)
356 .unwrap();
357
358 let graph = m.graph();
359 let paths = find_paths(
361 graph,
362 &a,
363 &a,
364 RouteSearch { bounds: &hops(2, 2), exclusions: &RouteExclusions::default() },
365 None,
366 )
367 .unwrap();
368
369 assert_eq!(paths.len(), 1);
371
372 let score = try_score_path(&paths[0]).unwrap();
374 let expected = 2.0 * 0.6 * 800.0;
375 assert_eq!(score, expected, "expected {expected}, got {score}");
376 }
377
378 #[rstest]
379 #[case::source_not_in_graph(false, true)]
380 #[case::dest_not_in_graph(true, false)]
381 fn test_find_paths_token_not_in_graph(#[case] from_exists: bool, #[case] to_exists: bool) {
382 let (a, b, _, _) = addrs();
384 let non_existent = addr(0x99);
385 let m = linear_graph();
386 let g = m.graph();
387
388 let from = if from_exists { a } else { non_existent.clone() };
389 let to = if to_exists { b } else { non_existent };
390
391 let result = find_paths(
392 g,
393 &from,
394 &to,
395 RouteSearch { bounds: &hops(1, 3), exclusions: &RouteExclusions::default() },
396 None,
397 );
398
399 assert!(matches!(result, Err(AlgorithmError::NoPath { .. })));
400 }
401
402 #[test]
410 fn test_simulate_path_single_hop() {
411 let token_a = token(0x01, "A");
412 let token_b = token(0x02, "B");
413
414 let (market, manager) = setup_market_weighted(vec![(
415 "component1",
416 &token_a,
417 &token_b,
418 MockProtocolSim::new(2.0),
419 )]);
420
421 let filter = hops(1, 1);
422 let paths = find_paths(
423 manager.graph(),
424 &token_a.address,
425 &token_b.address,
426 RouteSearch { bounds: &filter, exclusions: &RouteExclusions::default() },
427 None,
428 )
429 .unwrap();
430 let path = paths.into_iter().next().unwrap();
431
432 let result = simulate_pool_path(
433 &path,
434 market_read(&market).base_market_state(),
435 None,
436 BigUint::from(100u64),
437 )
438 .unwrap();
439
440 assert_eq!(result.route().swaps().len(), 1);
441 assert_eq!(*result.route().swaps()[0].amount_in(), BigUint::from(100u64));
442 assert_eq!(*result.route().swaps()[0].amount_out(), BigUint::from(200u64)); assert_eq!(result.route().swaps()[0].component_id(), "component1");
444 }
445
446 #[test]
447 fn test_simulate_path_multi_hop_chains_amounts() {
448 let token_a = token(0x01, "A");
449 let token_b = token(0x02, "B");
450 let token_c = token(0x03, "C");
451
452 let (market, manager) = setup_market_weighted(vec![
453 ("component1", &token_a, &token_b, MockProtocolSim::new(2.0)),
454 ("component2", &token_b, &token_c, MockProtocolSim::new(3.0)),
455 ]);
456
457 let filter = hops(2, 2);
458 let paths = find_paths(
459 manager.graph(),
460 &token_a.address,
461 &token_c.address,
462 RouteSearch { bounds: &filter, exclusions: &RouteExclusions::default() },
463 None,
464 )
465 .unwrap();
466 let path = paths.into_iter().next().unwrap();
467
468 let result = simulate_pool_path(
469 &path,
470 market_read(&market).base_market_state(),
471 None,
472 BigUint::from(10u64),
473 )
474 .unwrap();
475
476 assert_eq!(result.route().swaps().len(), 2);
477 assert_eq!(*result.route().swaps()[0].amount_out(), BigUint::from(20u64));
479 assert_eq!(*result.route().swaps()[1].amount_in(), BigUint::from(20u64));
481 assert_eq!(*result.route().swaps()[1].amount_out(), BigUint::from(60u64));
482 }
483
484 #[test]
485 fn test_simulate_path_same_component_twice_uses_updated_state() {
486 let token_a = token(0x01, "A");
489 let token_b = token(0x02, "B");
490
491 let (market, manager) = setup_market_weighted(vec![(
492 "component1",
493 &token_a,
494 &token_b,
495 MockProtocolSim::new(2.0),
496 )]);
497
498 let filter = hops(2, 2);
501 let paths = find_paths(
502 manager.graph(),
503 &token_a.address,
504 &token_a.address,
505 RouteSearch { bounds: &filter, exclusions: &RouteExclusions::default() },
506 None,
507 )
508 .unwrap();
509
510 assert_eq!(paths.len(), 1);
512 let path = paths[0].clone();
513
514 let result = simulate_pool_path(
515 &path,
516 market_read(&market).base_market_state(),
517 None,
518 BigUint::from(10u64),
519 )
520 .unwrap();
521
522 assert_eq!(result.route().swaps().len(), 2);
523 assert_eq!(*result.route().swaps()[0].amount_out(), BigUint::from(20u64));
525 assert_eq!(*result.route().swaps()[1].amount_out(), BigUint::from(6u64));
527 }
528
529 #[test]
530 fn test_simulate_path_missing_token_returns_data_not_found() {
531 let token_a = token(0x01, "A");
532 let token_b = token(0x02, "B");
533 let token_c = token(0x03, "C");
534
535 let (market, _) = setup_market_weighted(vec![(
536 "component1",
537 &token_a,
538 &token_b,
539 MockProtocolSim::new(2.0),
540 )]);
541 let market = market_read(&market);
542
543 let mut topology = market.component_topology();
545 topology.insert(
546 "component2".to_string(),
547 vec![token_b.address.clone(), token_c.address.clone()],
548 );
549 let mut manager = TopologyGraphManager::<DepthAndPrice>::default();
550 manager.initialize_graph(&topology);
551
552 let graph = manager.graph();
553 let filter = hops(2, 2);
554 let paths = find_paths(
555 graph,
556 &token_a.address,
557 &token_c.address,
558 RouteSearch { bounds: &filter, exclusions: &RouteExclusions::default() },
559 None,
560 )
561 .unwrap();
562 let path = paths.into_iter().next().unwrap();
563
564 let result =
565 simulate_pool_path(&path, market.base_market_state(), None, BigUint::from(100u64));
566 assert!(matches!(result, Err(AlgorithmError::DataNotFound { kind: "token", .. })));
567 }
568
569 #[test]
570 fn test_simulate_path_missing_component_returns_data_not_found() {
571 let token_a = token(0x01, "A");
572 let token_b = token(0x02, "B");
573 let (market, manager) = setup_market_weighted(vec![(
574 "component1",
575 &token_a,
576 &token_b,
577 MockProtocolSim::new(2.0),
578 )]);
579
580 let mut market_write = market.try_write().unwrap();
582 market_write.remove_components([&"component1".to_string()]);
583 drop(market_write);
584
585 let graph = manager.graph();
586 let filter = hops(1, 1);
587 let paths = find_paths(
588 graph,
589 &token_a.address,
590 &token_b.address,
591 RouteSearch { bounds: &filter, exclusions: &RouteExclusions::default() },
592 None,
593 )
594 .unwrap();
595 let path = paths.into_iter().next().unwrap();
596
597 let result = simulate_pool_path(
598 &path,
599 market_read(&market).base_market_state(),
600 None,
601 BigUint::from(100u64),
602 );
603 assert!(matches!(result, Err(AlgorithmError::DataNotFound { kind: "component", .. })));
604 }
605}