1use std::ops::Deref;
7
8use async_trait::async_trait;
9use petgraph::{
10 graph::{EdgeIndex, NodeIndex},
11 stable_graph,
12};
13use rustc_hash::{FxHashMap, FxHashSet};
14use smallvec::SmallVec;
15use tracing::{debug, trace};
16use tycho_simulation::tycho_common::models::Address;
17
18use super::{EdgeData, GraphError, GraphManager, Path, RouteSearch, INLINE_EDGES, INLINE_TOKENS};
19use crate::{
20 feed::{
21 events::{EventError, MarketEvent, MarketEventHandler},
22 market_data::MarketDataView,
23 },
24 types::{ComponentId, RouteExclusions},
25};
26
27#[derive(Debug, Clone)]
29pub struct PairEdge<D> {
30 pools: Vec<EdgeData<D>>,
32}
33
34impl<D> PairEdge<D> {
35 pub fn pools(&self) -> &[EdgeData<D>] {
37 &self.pools
38 }
39
40 fn insert(&mut self, component_id: &ComponentId) {
42 if self
43 .pools
44 .iter()
45 .any(|pool| &pool.component_id == component_id)
46 {
47 return;
48 }
49 self.pools
50 .push(EdgeData::new(component_id.clone()));
51 }
52
53 fn remove(&mut self, component_id: &ComponentId) -> bool {
55 self.pools
56 .retain(|pool| &pool.component_id != component_id);
57 !self.pools.is_empty()
58 }
59
60 #[cfg(any(test, feature = "test-utils"))]
62 fn pool_mut(&mut self, component_id: &ComponentId) -> Option<&mut EdgeData<D>> {
63 self.pools
64 .iter_mut()
65 .find(|pool| &pool.component_id == component_id)
66 }
67}
68
69pub type TokenPath = SmallVec<[NodeIndex; INLINE_TOKENS]>;
71
72enum Leg<'a, D> {
77 All(&'a [EdgeData<D>]),
79 Allowed(Vec<&'a EdgeData<D>>),
81}
82
83impl<'a, D> Leg<'a, D> {
84 fn len(&self) -> usize {
86 match self {
87 Self::All(pools) => pools.len(),
88 Self::Allowed(pools) => pools.len(),
89 }
90 }
91
92 fn is_empty(&self) -> bool {
94 self.len() == 0
95 }
96
97 fn get(&self, index: usize) -> Option<&'a EdgeData<D>> {
99 match self {
100 Self::All(pools) => (*pools).get(index),
101 Self::Allowed(pools) => pools.get(index).copied(),
102 }
103 }
104}
105
106pub type TokenGraph<D> = stable_graph::StableDiGraph<Address, PairEdge<D>>;
108
109pub struct TopologyGraph<D> {
113 graph: TokenGraph<D>,
117 tokens: FxHashMap<Address, NodeIndex>,
119 pair_index: FxHashMap<(NodeIndex, NodeIndex), EdgeIndex>,
122}
123
124impl<D> TopologyGraph<D> {
125 pub fn pools_between(&self, from: NodeIndex, to: NodeIndex) -> &[EdgeData<D>] {
127 self.pair_index
128 .get(&(from, to))
129 .and_then(|&edge| self.graph.edge_weight(edge))
130 .map_or(&[], PairEdge::pools)
131 }
132
133 fn leg_between(
137 &self,
138 from: NodeIndex,
139 to: NodeIndex,
140 exclusions: &RouteExclusions,
141 ) -> Leg<'_, D> {
142 let pools = self.pools_between(from, to);
143 if exclusions.is_empty() ||
144 !pools
145 .iter()
146 .any(|pool| exclusions.excludes_pool(&pool.component_id))
147 {
148 return Leg::All(pools);
149 }
150 Leg::Allowed(
151 pools
152 .iter()
153 .filter(|pool| !exclusions.excludes_pool(&pool.component_id))
154 .collect(),
155 )
156 }
157
158 fn pair_has_allowed_pool(
163 &self,
164 from: NodeIndex,
165 to: NodeIndex,
166 exclusions: &RouteExclusions,
167 ) -> bool {
168 let pools = self.pools_between(from, to);
169 if exclusions.is_empty() {
170 return !pools.is_empty();
171 }
172 pools
173 .iter()
174 .any(|pool| !exclusions.excludes_pool(&pool.component_id))
175 }
176
177 pub fn get_token_ix(&self, token: &Address) -> Option<NodeIndex> {
179 self.tokens.get(token).copied()
180 }
181
182 fn get_edge_ix(&self, from: NodeIndex, to: NodeIndex) -> Option<EdgeIndex> {
184 self.pair_index
185 .get(&(from, to))
186 .copied()
187 }
188
189 fn add_component(&mut self, from: NodeIndex, to: NodeIndex, component_id: &ComponentId) {
192 match self.get_edge_ix(from, to) {
193 Some(edge) => {
194 if let Some(pair) = self.graph.edge_weight_mut(edge) {
195 pair.insert(component_id);
196 }
197 }
198 None => {
199 let pair = PairEdge { pools: vec![EdgeData::new(component_id.clone())] };
200 let edge = self.graph.add_edge(from, to, pair);
201 self.pair_index.insert((from, to), edge);
202 }
203 }
204 }
205
206 fn remove_component(&mut self, from: NodeIndex, to: NodeIndex, component_id: &ComponentId) {
209 let Some(edge) = self.get_edge_ix(from, to) else {
210 return;
211 };
212 let still_traded = self
213 .graph
214 .edge_weight_mut(edge)
215 .is_some_and(|pair| pair.remove(component_id));
216 if !still_traded {
217 self.graph.remove_edge(edge);
218 self.pair_index.remove(&(from, to));
219 }
220 }
221
222 pub fn paths_between(
230 &self,
231 from: &Address,
232 to: &Address,
233 search: RouteSearch<'_>,
234 ) -> Result<Vec<TokenPath>, GraphError> {
235 let from_ix = self
236 .get_token_ix(from)
237 .ok_or_else(|| GraphError::TokenNotFound(from.clone()))?;
238 let to_ix = self
239 .get_token_ix(to)
240 .ok_or_else(|| GraphError::TokenNotFound(to.clone()))?;
241 Ok(self.paths_between_ix(from_ix, to_ix, search))
242 }
243
244 pub fn paths_between_ix(
250 &self,
251 from: NodeIndex,
252 to: NodeIndex,
253 search: RouteSearch<'_>,
254 ) -> Vec<TokenPath> {
255 let filter = search.bounds;
256 if filter.min_hops == 0 || filter.min_hops > filter.max_hops {
257 return Vec::new();
258 }
259 if from == to {
260 self.circular_token_paths(to, search)
261 } else {
262 self.bidirectional_search(from, to, search)
263 }
264 }
265
266 pub fn expand_path(
277 &self,
278 token_path: &[NodeIndex],
279 max_paths: Option<usize>,
280 exclusions: &RouteExclusions,
281 ) -> Vec<Path<'_, D>> {
282 let legs: SmallVec<[Leg<'_, D>; INLINE_EDGES]> = token_path
286 .windows(2)
287 .map(|pair| self.leg_between(pair[0], pair[1], exclusions))
288 .collect();
289
290 if legs.is_empty() || legs.iter().any(|leg| leg.is_empty()) {
295 return Vec::new();
296 }
297
298 let combinations: usize = legs
299 .iter()
300 .map(|leg| leg.len())
301 .product();
302 let wanted = max_paths.map_or(combinations, |cap| cap.min(combinations));
303 let mut chosen: SmallVec<[usize; INLINE_EDGES]> = SmallVec::from_elem(0, legs.len());
304
305 let mut out = Vec::with_capacity(wanted);
306 for _ in 0..wanted {
307 let mut path = Path::new();
308 for (leg, (pair, &pick)) in legs
309 .iter()
310 .zip(token_path.windows(2).zip(chosen.iter()))
311 {
312 let pool = leg
313 .get(pick)
314 .expect("odometer holds every leg inside its own pool count");
315 path.add_hop(&self[pair[0]], pool, &self[pair[1]]);
316 }
317 out.push(path);
318
319 for (pick, leg) in chosen.iter_mut().zip(legs.iter()).rev() {
320 *pick += 1;
321 if *pick < leg.len() {
322 break;
323 }
324 *pick = 0;
325 }
326 }
327
328 out
329 }
330
331 fn add_token(&mut self, address: Address) -> NodeIndex {
333 if let Some(index) = self.get_token_ix(&address) {
334 return index;
335 }
336 let index = self.graph.add_node(address.clone());
337 self.tokens.insert(address, index);
338 index
339 }
340
341 fn bidirectional_search(
360 &self,
361 from: NodeIndex,
362 to: NodeIndex,
363 search: RouteSearch<'_>,
364 ) -> Vec<TokenPath> {
365 let filter = search.bounds;
366 let endpoints = (from, to);
367 let head_levels = self.walk_levels(from, filter.max_hops.div_ceil(2), endpoints, search);
368 let tail_levels = self.walk_levels(to, filter.max_hops / 2, endpoints, search);
369
370 let mut midpoint_index: Vec<Option<FxHashMap<NodeIndex, Vec<TokenPath>>>> =
372 vec![None; tail_levels.len()];
373 let mut token_paths = Vec::new();
374
375 for length in filter.min_hops..=filter.max_hops {
376 let head_hops = length.div_ceil(2);
377 let tail_hops = length - head_hops;
378 let (Some(heads), Some(tails)) =
379 (head_levels.get(head_hops), tail_levels.get(tail_hops))
380 else {
381 continue;
382 };
383 if heads.is_empty() || tails.is_empty() {
384 continue;
385 }
386
387 let tails_by_midpoint = midpoint_index[tail_hops].get_or_insert_with(|| {
388 let mut index: FxHashMap<NodeIndex, Vec<TokenPath>> = FxHashMap::default();
389 for tail in tails {
390 let Some(&midpoint) = tail.last() else {
391 continue;
392 };
393 index
394 .entry(midpoint)
395 .or_default()
396 .push(tail.iter().rev().copied().collect());
397 }
398 index
399 });
400
401 for head in heads {
402 let Some(&midpoint) = head.last() else {
403 continue;
404 };
405 let Some(candidates) = tails_by_midpoint.get(&midpoint) else {
406 continue;
407 };
408 for tail in candidates {
409 let collides = head
412 .iter()
413 .any(|token| *token != midpoint && tail.contains(token));
414 if collides {
415 continue;
416 }
417
418 let mut joined = TokenPath::from_slice(head);
419 joined.extend_from_slice(&tail[1..]);
420 token_paths.push(joined);
421 }
422 }
423 }
424
425 token_paths
426 }
427
428 fn walk_levels(
437 &self,
438 start: NodeIndex,
439 hops: usize,
440 endpoints: (NodeIndex, NodeIndex),
441 search: RouteSearch<'_>,
442 ) -> Vec<Vec<TokenPath>> {
443 let (from, to) = endpoints;
444 let mut levels = vec![vec![TokenPath::from_slice(&[start])]];
445
446 for hop in 0..hops {
447 let mut next = Vec::new();
448 for sequence in &levels[hop] {
449 let Some(&last) = sequence.last() else {
450 continue;
451 };
452 for neighbor in self.neighbors(last) {
453 if sequence.contains(&neighbor) {
454 continue;
455 }
456 if !search.allows_token(&self[neighbor], (&self[from], &self[to])) {
457 continue;
458 }
459 if !self.pair_has_allowed_pool(last, neighbor, search.exclusions) {
460 continue;
461 }
462
463 let mut extended = TokenPath::from_slice(sequence);
464 extended.push(neighbor);
465 next.push(extended);
466 }
467 }
468 levels.push(next);
469 }
470
471 levels
472 }
473
474 fn circular_token_paths(&self, target: NodeIndex, search: RouteSearch<'_>) -> Vec<TokenPath> {
480 let filter = search.bounds;
481 let mut token_paths = Vec::new();
482 let mut frontier = vec![TokenPath::from_slice(&[target])];
483
484 for hops in 1..=filter.max_hops {
485 let mut next = Vec::new();
486 for sequence in &frontier {
487 let Some(&last) = sequence.last() else {
488 continue;
489 };
490 for neighbor in self.neighbors(last) {
491 if !self.pair_has_allowed_pool(last, neighbor, search.exclusions) {
492 continue;
493 }
494 if neighbor == target {
495 if hops >= filter.min_hops {
496 let mut closed = TokenPath::from_slice(sequence);
497 closed.push(neighbor);
498 token_paths.push(closed);
499 }
500 continue;
504 }
505 if sequence.contains(&neighbor) {
506 continue;
507 }
508 if !search.allows_token(&self[neighbor], (&self[target], &self[target])) {
509 continue;
510 }
511
512 let mut extended = TokenPath::from_slice(sequence);
513 extended.push(neighbor);
514 next.push(extended);
515 }
516 }
517 frontier = next;
518 }
519
520 token_paths
521 }
522}
523
524impl<D> Deref for TopologyGraph<D> {
525 type Target = TokenGraph<D>;
526
527 fn deref(&self) -> &Self::Target {
528 &self.graph
529 }
530}
531
532impl<D> Default for TopologyGraph<D> {
533 fn default() -> Self {
534 Self {
535 graph: TokenGraph::default(),
536 tokens: FxHashMap::default(),
537 pair_index: FxHashMap::default(),
538 }
539 }
540}
541
542pub struct TopologyGraphManager<D: Clone> {
546 graph: TopologyGraph<D>,
547 component_pairs: FxHashMap<ComponentId, Vec<(NodeIndex, NodeIndex)>>,
550}
551
552impl<D: Clone> TopologyGraphManager<D> {
553 pub fn new() -> Self {
555 Self { graph: TopologyGraph::default(), component_pairs: FxHashMap::default() }
556 }
557
558 fn add_component_edges(&mut self, component_id: &ComponentId, nodes: &[NodeIndex]) {
560 let pairs: Vec<(NodeIndex, NodeIndex)> = nodes
561 .iter()
562 .enumerate()
563 .flat_map(|(i, &from)| {
564 nodes
565 .iter()
566 .skip(i + 1)
567 .flat_map(move |&to| [(from, to), (to, from)])
568 })
569 .collect();
570
571 for &(from, to) in &pairs {
572 self.graph
573 .add_component(from, to, component_id);
574 }
575 self.component_pairs
576 .insert(component_id.clone(), pairs);
577 }
578
579 fn add_components(
586 &mut self,
587 components: &FxHashMap<ComponentId, Vec<Address>>,
588 ) -> Result<(), GraphError> {
589 let mut invalid = Vec::new();
590 let mut skipped = 0usize;
591
592 let mut sorted: Vec<_> = components.iter().collect();
594 sorted.sort_by_key(|(id, _)| *id);
595
596 for (component_id, tokens) in sorted {
597 if self
598 .component_pairs
599 .contains_key(component_id)
600 {
601 trace!(component_id = %component_id, "skipping already-tracked component");
602 skipped += 1;
603 continue;
604 }
605 if tokens.len() < 2 {
606 invalid.push(component_id.clone());
607 continue;
608 }
609
610 let mut sorted_tokens: Vec<&Address> = tokens.iter().collect();
611 sorted_tokens.sort();
612 let nodes: Vec<NodeIndex> = sorted_tokens
613 .iter()
614 .map(|token| self.graph.add_token((*token).clone()))
615 .collect();
616 self.add_component_edges(component_id, &nodes);
617 }
618
619 if skipped > 0 {
620 debug!(skipped_duplicates = skipped, "skipped duplicate components during add");
621 }
622 if !invalid.is_empty() {
623 return Err(GraphError::InvalidComponents(invalid));
624 }
625 Ok(())
626 }
627
628 fn remove_components(&mut self, components: &[ComponentId]) -> Result<(), GraphError> {
635 let mut missing = Vec::new();
636
637 for component_id in components {
638 let Some(pairs) = self
639 .component_pairs
640 .remove(component_id)
641 else {
642 missing.push(component_id.clone());
643 continue;
644 };
645
646 for (from, to) in pairs {
647 self.graph
648 .remove_component(from, to, component_id);
649 }
650 }
651
652 if !missing.is_empty() {
653 return Err(GraphError::ComponentsNotFound(missing));
654 }
655 Ok(())
656 }
657
658 #[cfg(any(test, feature = "test-utils"))]
660 pub(crate) fn set_pool_weight(
661 &mut self,
662 component_id: &ComponentId,
663 token_in: &Address,
664 token_out: &Address,
665 data: D,
666 bidirectional: bool,
667 ) -> Result<(), GraphError> {
668 let from = self
669 .graph
670 .get_token_ix(token_in)
671 .ok_or_else(|| GraphError::TokenNotFound(token_in.clone()))?;
672 let to = self
673 .graph
674 .get_token_ix(token_out)
675 .ok_or_else(|| GraphError::TokenNotFound(token_out.clone()))?;
676
677 let mut directions = vec![(from, to)];
678 if bidirectional {
679 directions.push((to, from));
680 }
681
682 let mut updated = false;
683 for (source, target) in directions {
684 let Some(edge) = self.graph.get_edge_ix(source, target) else {
685 continue;
686 };
687 if let Some(pool) = self
688 .graph
689 .graph
690 .edge_weight_mut(edge)
691 .and_then(|pair| pair.pool_mut(component_id))
692 {
693 pool.data = Some(data.clone());
694 updated = true;
695 }
696 }
697
698 if !updated {
699 return Err(GraphError::MissingComponentBetweenTokens(
700 token_in.clone(),
701 token_out.clone(),
702 component_id.clone(),
703 ));
704 }
705 Ok(())
706 }
707}
708
709impl<D: Clone + super::EdgeWeightFromSimAndDerived> super::EdgeWeightUpdaterWithDerived
710 for TopologyGraphManager<D>
711{
712 fn update_edge_weights_with_derived(
717 &mut self,
718 market: MarketDataView<'_>,
719 derived: &crate::derived::DerivedData,
720 ) -> usize {
721 let tokens = market.token_registry_ref();
722 let mut updated = 0usize;
723
724 for edge in self
725 .graph
726 .edge_indices()
727 .collect::<Vec<_>>()
728 {
729 let Some((source, target)) = self.graph.edge_endpoints(edge) else {
730 continue;
731 };
732 let (Some(token_in), Some(token_out)) =
735 (tokens.get(&self.graph[source]), tokens.get(&self.graph[target]))
736 else {
737 continue;
738 };
739
740 let Some(pair) = self.graph.graph.edge_weight_mut(edge) else {
741 continue;
742 };
743 for pool in &mut pair.pools {
744 pool.data = market
745 .get_simulation_state(&pool.component_id)
746 .and_then(|state| {
747 D::from_sim_and_derived(
748 state,
749 &pool.component_id,
750 token_in,
751 token_out,
752 derived,
753 )
754 });
755 if pool.data.is_some() {
756 updated += 1;
757 }
758 }
759 }
760
761 updated
762 }
763}
764
765impl<D: Clone> Default for TopologyGraphManager<D> {
766 fn default() -> Self {
767 Self::new()
768 }
769}
770
771impl<D: Clone + Send + Sync> GraphManager<TopologyGraph<D>> for TopologyGraphManager<D> {
772 fn initialize_graph(&mut self, component_topology: &FxHashMap<ComponentId, Vec<Address>>) {
773 self.graph = TopologyGraph::default();
774 self.component_pairs.clear();
775
776 let mut tokens: Vec<Address> = component_topology
779 .values()
780 .flat_map(|addresses| addresses.iter())
781 .cloned()
782 .collect::<FxHashSet<_>>()
783 .into_iter()
784 .collect();
785 tokens.sort();
786
787 for token in tokens {
788 self.graph.add_token(token);
789 }
790
791 if let Err(e) = self.add_components(component_topology) {
795 debug!(error = %e, "components skipped while building the graph");
796 }
797 }
798
799 fn graph(&self) -> &TopologyGraph<D> {
800 &self.graph
801 }
802}
803
804#[async_trait]
805impl<D: Clone + Send> MarketEventHandler for TopologyGraphManager<D> {
806 async fn handle_event(&mut self, event: &MarketEvent) -> Result<(), EventError> {
807 match event {
808 MarketEvent::MarketUpdated { added_components, removed_components, .. } => {
809 let mut errors = Vec::new();
810 if let Err(e) = self.add_components(added_components) {
811 errors.push(e);
812 }
813 if let Err(e) = self.remove_components(removed_components) {
814 errors.push(e);
815 }
816 if errors.is_empty() {
817 Ok(())
818 } else {
819 Err(EventError::GraphErrors(errors))
820 }
821 }
822 }
823 }
824}
825
826#[cfg(test)]
827mod tests {
828 use rstest::rstest;
829 use rustc_hash::{FxHashMap, FxHashSet};
830
831 use super::*;
832 use crate::{
833 algorithm::{
834 most_liquid::DepthAndPrice,
835 test_utils::fixtures::{addrs, diamond_graph, linear_graph, parallel_graph},
836 },
837 graph::{EdgeWeightUpdaterWithDerived, GraphManager, GraphQueryFilter},
838 };
839
840 fn addr(byte: u8) -> Address {
841 Address::from([byte; 20])
842 }
843
844 #[tokio::test]
851 async fn test_pair_edge_lives_from_first_pool_to_last() {
852 use crate::feed::events::{MarketEvent, MarketEventHandler};
853
854 let (a, b) = (addr(0x0A), addr(0x0B));
855 let mut manager = TopologyGraphManager::<()>::new();
856 manager.initialize_graph(&FxHashMap::from_iter([(
857 "first".to_string(),
858 vec![a.clone(), b.clone()],
859 )]));
860
861 let (from, to) = (
862 manager
863 .graph()
864 .get_token_ix(&a)
865 .unwrap(),
866 manager
867 .graph()
868 .get_token_ix(&b)
869 .unwrap(),
870 );
871 assert_eq!(manager.graph().edge_count(), 2, "one edge each way");
872 assert_eq!(
873 manager
874 .graph()
875 .pools_between(from, to)
876 .len(),
877 1
878 );
879
880 let added = |id: &str| MarketEvent::MarketUpdated {
881 added_components: FxHashMap::from_iter([(id.to_string(), vec![a.clone(), b.clone()])]),
882 removed_components: vec![],
883 updated_components: vec![],
884 };
885 let removed = |id: &str| MarketEvent::MarketUpdated {
886 added_components: FxHashMap::default(),
887 removed_components: vec![id.to_string()],
888 updated_components: vec![],
889 };
890
891 manager
893 .handle_event(&added("second"))
894 .await
895 .unwrap();
896 assert_eq!(manager.graph().edge_count(), 2, "a second pool is not a second edge");
897 assert_eq!(
898 manager
899 .graph()
900 .pools_between(from, to)
901 .len(),
902 2
903 );
904
905 manager
907 .handle_event(&removed("second"))
908 .await
909 .unwrap();
910 assert_eq!(manager.graph().edge_count(), 2, "the pair still trades");
911 assert_eq!(
912 manager
913 .graph()
914 .pools_between(from, to)
915 .len(),
916 1
917 );
918
919 manager
921 .handle_event(&removed("first"))
922 .await
923 .unwrap();
924 assert_eq!(manager.graph().edge_count(), 0);
925 assert!(manager
926 .graph()
927 .pools_between(from, to)
928 .is_empty());
929
930 manager
932 .handle_event(&added("third"))
933 .await
934 .unwrap();
935 assert_eq!(manager.graph().edge_count(), 2);
936 assert_eq!(
937 manager
938 .graph()
939 .pools_between(from, to)
940 .len(),
941 1
942 );
943 assert_eq!(manager.graph().pools_between(from, to)[0].component_id, "third");
944
945 let (c, d) = (addr(0x0C), addr(0x0D));
947 manager
948 .handle_event(&MarketEvent::MarketUpdated {
949 added_components: FxHashMap::from_iter([(
950 "fourth".to_string(),
951 vec![c.clone(), d.clone()],
952 )]),
953 removed_components: vec![],
954 updated_components: vec![],
955 })
956 .await
957 .unwrap();
958 assert_eq!(manager.graph().node_count(), 4);
959 assert_eq!(manager.graph().edge_count(), 4);
960 }
961
962 #[test]
963 fn test_edge_weight_cleared_on_spot_price_miss() {
964 use num_bigint::BigUint;
967 use num_traits::One;
968 use tycho_simulation::tycho_core::simulation::protocol_sim::Price;
969
970 use crate::{
971 algorithm::test_utils::{market_read, setup_market_weighted, token, MockProtocolSim},
972 derived::{types::TokenGasPrices, DerivedData},
973 };
974
975 let token_a = token(0x01, "A");
976 let token_b = token(0x02, "B");
977 let (market, mut manager) = setup_market_weighted(vec![(
978 "component1",
979 &token_a,
980 &token_b,
981 MockProtocolSim::new(2.0),
982 )]);
983
984 assert!(
985 manager
986 .graph()
987 .edge_indices()
988 .all(|e| manager
989 .graph()
990 .edge_weight(e)
991 .unwrap()
992 .pools()
993 .iter()
994 .all(|pool| pool.data.is_some())),
995 "edges should have weight data after setup"
996 );
997
998 let mut token_prices = TokenGasPrices::default();
999 for addr in [&token_a.address, &token_b.address] {
1000 token_prices.insert(
1001 addr.clone(),
1002 Price { numerator: BigUint::one(), denominator: BigUint::one() },
1003 );
1004 }
1005 let mut derived = DerivedData::new();
1006 derived.set_spot_prices(Default::default(), vec![], 10, true);
1007 derived.set_component_depths(Default::default(), vec![], 10, true);
1008 derived.set_token_prices(token_prices, vec![], 10, true);
1009
1010 manager.update_edge_weights_with_derived(market_read(&market), &derived);
1011
1012 assert!(
1013 manager
1014 .graph()
1015 .edge_indices()
1016 .all(|e| manager
1017 .graph()
1018 .edge_weight(e)
1019 .unwrap()
1020 .pools()
1021 .iter()
1022 .all(|pool| pool.data.is_none())),
1023 "stale edge weights must be cleared when spot price is unavailable"
1024 );
1025 }
1026
1027 fn routes<'a>(
1032 graph: &'a TopologyGraph<DepthAndPrice>,
1033 from: &Address,
1034 to: &Address,
1035 min_hops: usize,
1036 max_hops: usize,
1037 connector_tokens: Option<FxHashSet<Address>>,
1038 ) -> Vec<Path<'a, DepthAndPrice>> {
1039 let (Some(from), Some(to)) = (graph.get_token_ix(from), graph.get_token_ix(to)) else {
1040 return Vec::new();
1041 };
1042 let filter = GraphQueryFilter { min_hops, max_hops, connector_tokens };
1043 let exclusions = RouteExclusions::default();
1044 let search = RouteSearch { bounds: &filter, exclusions: &exclusions };
1045 graph
1046 .paths_between_ix(from, to, search)
1047 .iter()
1048 .flat_map(|token_path| graph.expand_path(token_path, None, &exclusions))
1049 .collect()
1050 }
1051
1052 fn all_ids(paths: Vec<Path<'_, DepthAndPrice>>) -> FxHashSet<Vec<&str>> {
1053 paths
1054 .iter()
1055 .map(|p| {
1056 p.iter()
1057 .map(|(_, e, _)| e.component_id.as_str())
1058 .collect()
1059 })
1060 .collect()
1061 }
1062
1063 #[test]
1064 fn test_find_paths_linear_forward_and_reverse() {
1065 let (a, b, c, d) = addrs();
1066 let m = linear_graph();
1067 let g = m.graph();
1068
1069 let p = routes(g, &a, &b, 1, 1, None);
1071 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab"]]));
1072
1073 let p = routes(g, &a, &c, 1, 2, None);
1074 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab", "bc"]]));
1075
1076 let p = routes(g, &a, &d, 1, 3, None);
1077 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab", "bc", "cd"]]));
1078
1079 let p = routes(g, &d, &a, 1, 3, None);
1081 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["cd", "bc", "ab"]]));
1082 }
1083
1084 #[test]
1085 fn test_find_paths_respects_hop_bounds() {
1086 let (a, _, c, d) = addrs();
1087 let m = linear_graph();
1088 let g = m.graph();
1089
1090 assert!(routes(g, &a, &d, 1, 2, None).is_empty());
1092
1093 assert!(routes(g, &a, &c, 3, 3, None).is_empty());
1095 }
1096
1097 #[test]
1098 fn test_find_paths_parallel_components() {
1099 let (a, b, c, _) = addrs();
1100 let m = parallel_graph();
1101 let g = m.graph();
1102
1103 let p = routes(g, &a, &b, 1, 1, None);
1105 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab1"], vec!["ab2"], vec!["ab3"]]));
1106
1107 let p = routes(g, &a, &c, 1, 2, None);
1109 assert_eq!(
1110 all_ids(p),
1111 FxHashSet::from_iter([
1112 vec!["ab1", "bc1"],
1113 vec!["ab1", "bc2"],
1114 vec!["ab2", "bc1"],
1115 vec!["ab2", "bc2"],
1116 vec!["ab3", "bc1"],
1117 vec!["ab3", "bc2"],
1118 ])
1119 );
1120 }
1121
1122 #[test]
1123 fn test_find_paths_diamond_multiple_routes() {
1124 let (a, _, _, d) = addrs();
1125 let m = diamond_graph();
1126 let g = m.graph();
1127
1128 let p = routes(g, &a, &d, 1, 2, None);
1130 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab", "bd"], vec!["ac", "cd"]]));
1131 }
1132
1133 #[test]
1134 fn test_find_paths_no_intermediate_cycles() {
1135 let (a, b, _, _) = addrs();
1136 let m = linear_graph();
1137 let g = m.graph();
1138
1139 let p = routes(g, &a, &b, 1, 3, None);
1144 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab"]]));
1145 }
1146
1147 #[test]
1148 fn test_find_paths_cyclic_same_source_dest() {
1149 let (a, _, _, _) = addrs();
1150 let m = parallel_graph();
1152 let g = m.graph();
1153
1154 let p = routes(g, &a, &a, 2, 2, None);
1157 assert_eq!(
1158 all_ids(p),
1159 FxHashSet::from_iter([
1160 vec!["ab1", "ab1"],
1161 vec!["ab1", "ab2"],
1162 vec!["ab1", "ab3"],
1163 vec!["ab2", "ab1"],
1164 vec!["ab2", "ab2"],
1165 vec!["ab2", "ab3"],
1166 vec!["ab3", "ab1"],
1167 vec!["ab3", "ab2"],
1168 vec!["ab3", "ab3"],
1169 ])
1170 );
1171 }
1172
1173 #[tokio::test]
1176 async fn test_add_components_reports_the_ones_with_too_few_tokens() {
1177 let (a, b) = (addr(0x0A), addr(0x0B));
1178 let mut manager = TopologyGraphManager::<()>::new();
1179 manager.initialize_graph(&FxHashMap::default());
1180
1181 let result = manager.add_components(&FxHashMap::from_iter([
1182 ("solo".to_string(), vec![a.clone()]),
1183 ("pair".to_string(), vec![a.clone(), b.clone()]),
1184 ]));
1185
1186 match result {
1187 Err(GraphError::InvalidComponents(invalid)) => {
1188 assert_eq!(invalid, vec!["solo".to_string()]);
1189 }
1190 other => panic!("expected InvalidComponents, got {other:?}"),
1191 }
1192 assert_eq!(manager.graph().edge_count(), 2, "the valid pair was still added");
1193 }
1194
1195 #[tokio::test]
1197 async fn test_remove_components_reports_the_ones_it_does_not_hold() {
1198 let (a, b) = (addr(0x0A), addr(0x0B));
1199 let mut manager = TopologyGraphManager::<()>::new();
1200 manager.initialize_graph(&FxHashMap::from_iter([(
1201 "pair".to_string(),
1202 vec![a.clone(), b.clone()],
1203 )]));
1204
1205 let result = manager.remove_components(&["pair".to_string(), "ghost".to_string()]);
1206
1207 match result {
1208 Err(GraphError::ComponentsNotFound(missing)) => {
1209 assert_eq!(missing, vec!["ghost".to_string()]);
1210 }
1211 other => panic!("expected ComponentsNotFound, got {other:?}"),
1212 }
1213 assert_eq!(manager.graph().edge_count(), 0, "the component it did hold was removed");
1214 }
1215
1216 #[test]
1218 fn test_set_pool_weight_rejects_a_pair_the_component_does_not_trade() {
1219 let (a, b, c, _) = addrs();
1220 let mut manager = linear_graph();
1221
1222 let unknown_token = manager.set_pool_weight(
1223 &"ab".to_string(),
1224 &a,
1225 &addr(0x99),
1226 DepthAndPrice::new(1.0, 1.0),
1227 false,
1228 );
1229 assert!(matches!(unknown_token, Err(GraphError::TokenNotFound(_))));
1230
1231 let wrong_pair =
1232 manager.set_pool_weight(&"ab".to_string(), &b, &c, DepthAndPrice::new(1.0, 1.0), false);
1233 assert!(matches!(wrong_pair, Err(GraphError::MissingComponentBetweenTokens(..))));
1234 }
1235
1236 #[rstest]
1238 #[case::zero_min(0, 3)]
1239 #[case::min_above_max(3, 1)]
1240 fn test_paths_between_rejects_impossible_hop_bounds(
1241 #[case] min_hops: usize,
1242 #[case] max_hops: usize,
1243 ) {
1244 let (a, b, _, _) = addrs();
1245 let m = linear_graph();
1246 let g = m.graph();
1247 let (from, to) = (g.get_token_ix(&a).unwrap(), g.get_token_ix(&b).unwrap());
1248
1249 let filter = GraphQueryFilter { min_hops, max_hops, connector_tokens: None };
1250 let exclusions = RouteExclusions::default();
1251 let search = RouteSearch { bounds: &filter, exclusions: &exclusions };
1252
1253 assert!(g
1254 .paths_between_ix(from, to, search)
1255 .is_empty());
1256 assert!(
1257 g.paths_between_ix(from, from, search)
1258 .is_empty(),
1259 "the cyclic search is bound by the same rule"
1260 );
1261 }
1262
1263 #[test]
1267 fn test_cyclic_routes_never_pass_through_the_start_token() {
1268 let (a, _, _, _) = addrs();
1269 let m = diamond_graph();
1270 let g = m.graph();
1271 let start = g.get_token_ix(&a).unwrap();
1272
1273 let cycles = g.paths_between_ix(
1274 start,
1275 start,
1276 RouteSearch {
1277 bounds: &GraphQueryFilter { min_hops: 1, max_hops: 4, connector_tokens: None },
1278 exclusions: &RouteExclusions::default(),
1279 },
1280 );
1281
1282 assert!(!cycles.is_empty(), "the diamond closes cycles through B and through C");
1283 for cycle in &cycles {
1284 assert_eq!(cycle.first(), Some(&start), "a cycle starts on its token");
1285 assert_eq!(cycle.last(), Some(&start), "a cycle ends on its token");
1286 assert!(
1287 !cycle[1..cycle.len() - 1].contains(&start),
1288 "the start token must appear only at the two ends, got {cycle:?}"
1289 );
1290 }
1291 }
1292
1293 fn routes_excluding<'a>(
1295 graph: &'a TopologyGraph<DepthAndPrice>,
1296 from: &Address,
1297 to: &Address,
1298 hops: (usize, usize),
1299 exclusions: &RouteExclusions,
1300 ) -> FxHashSet<Vec<&'a str>> {
1301 let filter =
1302 GraphQueryFilter { min_hops: hops.0, max_hops: hops.1, connector_tokens: None };
1303 let search = RouteSearch { bounds: &filter, exclusions };
1304 all_ids(
1305 graph
1306 .paths_between(from, to, search)
1307 .unwrap()
1308 .iter()
1309 .flat_map(|token_path| graph.expand_path(token_path, None, exclusions))
1310 .collect(),
1311 )
1312 }
1313
1314 #[test]
1316 fn test_paths_with_excluded_token() {
1317 let (a, b, _, d) = addrs();
1318 let m = diamond_graph();
1319 let exclusions = RouteExclusions::default().with_tokens([b.clone()]);
1320
1321 let routes = routes_excluding(m.graph(), &a, &d, (1, 3), &exclusions);
1322
1323 assert_eq!(routes, FxHashSet::from_iter([vec!["ac", "cd"]]));
1324 }
1325
1326 #[test]
1328 fn test_paths_with_every_pool_of_a_pair_excluded() {
1329 let (a, _, _, d) = addrs();
1330 let m = diamond_graph();
1331 let exclusions = RouteExclusions::default().with_pools(["ab".to_string()]);
1332
1333 let routes = routes_excluding(m.graph(), &a, &d, (1, 3), &exclusions);
1334
1335 assert_eq!(routes, FxHashSet::from_iter([vec!["ac", "cd"]]));
1336 }
1337
1338 #[test]
1340 fn test_expand_path_with_excluded_pool() {
1341 let (a, _, c, _) = addrs();
1342 let m = parallel_graph();
1343 let exclusions =
1344 RouteExclusions::default().with_pools(["ab2".to_string(), "bc1".to_string()]);
1345
1346 let routes = routes_excluding(m.graph(), &a, &c, (2, 2), &exclusions);
1347
1348 assert_eq!(
1349 routes,
1350 FxHashSet::from_iter([vec!["ab1", "bc2"], vec!["ab3", "bc2"]]),
1351 "the two excluded pools are gone; every other pool combination remains"
1352 );
1353 }
1354
1355 fn stacked_graph() -> TopologyGraphManager<DepthAndPrice> {
1367 let (s, x, y, t) = (addr(0x51), addr(0x58), addr(0x59), addr(0x54));
1368 let mut topology = FxHashMap::default();
1369 for (id, from, to) in [
1370 ("sx1", &s, &x),
1371 ("sx2", &s, &x),
1372 ("xy1", &x, &y),
1373 ("xy2", &x, &y),
1374 ("xy3", &x, &y),
1375 ("yt1", &y, &t),
1376 ("yt2", &y, &t),
1377 ("sy1", &s, &y),
1378 ("st1", &s, &t),
1379 ] {
1380 topology.insert(id.to_string(), vec![from.clone(), to.clone()]);
1381 }
1382
1383 let mut manager = TopologyGraphManager::<DepthAndPrice>::new();
1384 manager.initialize_graph(&topology);
1385 manager
1386 }
1387
1388 #[test]
1389 fn test_find_paths_expands_every_pool_combination() {
1390 let (s, x, y, t) = (addr(0x51), addr(0x58), addr(0x59), addr(0x54));
1391 let manager = stacked_graph();
1392 let graph = manager.graph();
1393
1394 assert_eq!(routes(graph, &s, &t, 1, 1, None).len(), 1);
1396 assert_eq!(routes(graph, &s, &t, 2, 2, None).len(), 2);
1398 assert_eq!(routes(graph, &s, &t, 3, 3, None).len(), 12);
1400
1401 let paths = routes(graph, &s, &t, 1, 3, None);
1402 assert_eq!(paths.len(), 15, "every combination, and none of them twice");
1403
1404 let long: FxHashSet<Vec<&str>> = paths
1407 .iter()
1408 .filter(|path| path.len() == 3)
1409 .map(|path| {
1410 path.edge_iter()
1411 .iter()
1412 .map(|edge| edge.component_id.as_str())
1413 .collect()
1414 })
1415 .collect();
1416 assert_eq!(long.len(), 12);
1417
1418 for path in &paths {
1419 assert_eq!(path.tokens.first().copied(), Some(&s));
1420 assert_eq!(path.tokens.last().copied(), Some(&t));
1421 assert_eq!(path.tokens.len(), path.len() + 1);
1422 }
1423
1424 let allowed = FxHashSet::from_iter([x]);
1426 let filtered = routes(graph, &s, &t, 1, 3, Some(allowed.clone()));
1427 assert_eq!(filtered.len(), 1);
1428 assert_eq!(filtered[0].edge_iter()[0].component_id, "st1");
1429
1430 let allowed = FxHashSet::from_iter([y]);
1431 assert_eq!(
1432 routes(graph, &s, &t, 1, 3, Some(allowed.clone())).len(),
1433 3,
1434 "st1, plus sy1 over each of the two Y-T pools"
1435 );
1436 }
1437
1438 #[test]
1439 fn test_find_paths_bfs_ordering() {
1440 let (a, b, c, d) = addrs();
1445 let e = addr(0x0E);
1446 let mut m = TopologyGraphManager::<DepthAndPrice>::new();
1447 let mut t = FxHashMap::default();
1448 t.insert("ae".into(), vec![a.clone(), e.clone()]);
1449 t.insert("ab".into(), vec![a.clone(), b.clone()]);
1450 t.insert("be".into(), vec![b, e.clone()]);
1451 t.insert("ac".into(), vec![a.clone(), c.clone()]);
1452 t.insert("cd".into(), vec![c, d.clone()]);
1453 t.insert("de".into(), vec![d, e.clone()]);
1454 m.initialize_graph(&t);
1455 let g = m.graph();
1456
1457 let p = routes(g, &a, &e, 1, 3, None);
1458
1459 assert_eq!(p.len(), 3, "Expected 3 paths total");
1461 assert_eq!(p[0].len(), 1, "First path should be 1-hop");
1462 assert_eq!(p[1].len(), 2, "Second path should be 2-hop");
1463 assert_eq!(p[2].len(), 3, "Third path should be 3-hop");
1464 }
1465
1466 #[test]
1467 fn test_connector_tokens_blocks_disallowed_intermediate() {
1468 let (a, b, c, d) = addrs();
1470 let m = diamond_graph();
1471 let g = m.graph();
1472 let allowed: FxHashSet<Address> = FxHashSet::from_iter([c.clone()]);
1473 let paths = routes(g, &a, &d, 1, 2, Some(allowed.clone()));
1474 let intermediates: FxHashSet<&Address> = paths
1475 .iter()
1476 .flat_map(|p| p.iter().map(|(node, _, _)| node))
1477 .filter(|addr| *addr != &a && *addr != &d)
1478 .collect();
1479 assert!(!intermediates.contains(&b), "B should be blocked");
1481 assert!(intermediates.contains(&c), "C should be allowed");
1482 }
1483
1484 #[test]
1485 fn test_connector_tokens_allows_endpoints_even_if_not_listed() {
1486 let (a, b, _, _) = addrs();
1488 let m = linear_graph();
1489 let g = m.graph();
1490 let allowed: FxHashSet<Address> = FxHashSet::default();
1493 let paths = routes(g, &a, &b, 1, 1, Some(allowed.clone()));
1494 assert!(!paths.is_empty(), "1-hop direct route should survive empty allowlist");
1495 }
1496
1497 #[test]
1498 fn test_connector_tokens_none_is_unrestricted() {
1499 let (a, b, c, d) = addrs();
1501 let m = diamond_graph();
1502 let g = m.graph();
1503 let paths = routes(g, &a, &d, 1, 2, None);
1504 let intermediates: FxHashSet<&Address> = paths
1505 .iter()
1506 .flat_map(|p| p.iter().map(|(node, _, _)| node))
1507 .filter(|addr| *addr != &a && *addr != &d)
1508 .collect();
1509 assert!(intermediates.contains(&b), "B should appear with no restriction");
1510 assert!(intermediates.contains(&c), "C should appear with no restriction");
1511 }
1512}