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::{
19 EdgeData, GraphError, GraphManager, GraphQueryFilter, Path, INLINE_EDGES, INLINE_TOKENS,
20};
21use crate::{
22 feed::{
23 events::{EventError, MarketEvent, MarketEventHandler},
24 market_data::MarketDataView,
25 },
26 types::ComponentId,
27};
28
29#[derive(Debug, Clone)]
31pub struct PairEdge<D> {
32 pools: Vec<EdgeData<D>>,
34}
35
36impl<D> PairEdge<D> {
37 pub fn pools(&self) -> &[EdgeData<D>] {
39 &self.pools
40 }
41
42 fn insert(&mut self, component_id: &ComponentId) {
44 if self
45 .pools
46 .iter()
47 .any(|pool| &pool.component_id == component_id)
48 {
49 return;
50 }
51 self.pools
52 .push(EdgeData::new(component_id.clone()));
53 }
54
55 fn remove(&mut self, component_id: &ComponentId) -> bool {
57 self.pools
58 .retain(|pool| &pool.component_id != component_id);
59 !self.pools.is_empty()
60 }
61
62 #[cfg(test)]
64 fn pool_mut(&mut self, component_id: &ComponentId) -> Option<&mut EdgeData<D>> {
65 self.pools
66 .iter_mut()
67 .find(|pool| &pool.component_id == component_id)
68 }
69}
70
71pub type TokenPath = SmallVec<[NodeIndex; INLINE_TOKENS]>;
73
74pub type TokenGraph<D> = stable_graph::StableDiGraph<Address, PairEdge<D>>;
76
77pub struct TopologyGraph<D> {
81 graph: TokenGraph<D>,
85 tokens: FxHashMap<Address, NodeIndex>,
87 pair_index: FxHashMap<(NodeIndex, NodeIndex), EdgeIndex>,
90}
91
92impl<D> TopologyGraph<D> {
93 pub fn pools_between(&self, from: NodeIndex, to: NodeIndex) -> &[EdgeData<D>] {
95 self.pair_index
96 .get(&(from, to))
97 .and_then(|&edge| self.graph.edge_weight(edge))
98 .map_or(&[], PairEdge::pools)
99 }
100
101 pub fn get_token_ix(&self, token: &Address) -> Option<NodeIndex> {
103 self.tokens.get(token).copied()
104 }
105
106 fn get_edge_ix(&self, from: NodeIndex, to: NodeIndex) -> Option<EdgeIndex> {
108 self.pair_index
109 .get(&(from, to))
110 .copied()
111 }
112
113 fn add_component(&mut self, from: NodeIndex, to: NodeIndex, component_id: &ComponentId) {
116 match self.get_edge_ix(from, to) {
117 Some(edge) => {
118 if let Some(pair) = self.graph.edge_weight_mut(edge) {
119 pair.insert(component_id);
120 }
121 }
122 None => {
123 let pair = PairEdge { pools: vec![EdgeData::new(component_id.clone())] };
124 let edge = self.graph.add_edge(from, to, pair);
125 self.pair_index.insert((from, to), edge);
126 }
127 }
128 }
129
130 fn remove_component(&mut self, from: NodeIndex, to: NodeIndex, component_id: &ComponentId) {
133 let Some(edge) = self.get_edge_ix(from, to) else {
134 return;
135 };
136 let still_traded = self
137 .graph
138 .edge_weight_mut(edge)
139 .is_some_and(|pair| pair.remove(component_id));
140 if !still_traded {
141 self.graph.remove_edge(edge);
142 self.pair_index.remove(&(from, to));
143 }
144 }
145
146 pub fn paths_between(
154 &self,
155 from: &Address,
156 to: &Address,
157 filter: &GraphQueryFilter,
158 ) -> Result<Vec<TokenPath>, GraphError> {
159 let from_ix = self
160 .get_token_ix(from)
161 .ok_or_else(|| GraphError::TokenNotFound(from.clone()))?;
162 let to_ix = self
163 .get_token_ix(to)
164 .ok_or_else(|| GraphError::TokenNotFound(to.clone()))?;
165 Ok(self.paths_between_ix(from_ix, to_ix, filter))
166 }
167
168 pub fn paths_between_ix(
174 &self,
175 from: NodeIndex,
176 to: NodeIndex,
177 filter: &GraphQueryFilter,
178 ) -> Vec<TokenPath> {
179 if filter.min_hops == 0 || filter.min_hops > filter.max_hops {
180 return Vec::new();
181 }
182 if from == to {
183 self.circular_token_paths(to, filter)
184 } else {
185 self.bidirectional_search(from, to, filter)
186 }
187 }
188
189 pub fn expand_path(
200 &self,
201 token_path: &[NodeIndex],
202 max_paths: Option<usize>,
203 ) -> Vec<Path<'_, D>> {
204 let legs: SmallVec<[&[EdgeData<D>]; INLINE_EDGES]> = token_path
207 .windows(2)
208 .map(|pair| self.pools_between(pair[0], pair[1]))
209 .collect();
210
211 if legs.is_empty() || legs.iter().any(|leg| leg.is_empty()) {
216 return Vec::new();
217 }
218
219 let combinations: usize = legs
220 .iter()
221 .map(|leg| leg.len())
222 .product();
223 let wanted = max_paths.map_or(combinations, |cap| cap.min(combinations));
224 let mut chosen: SmallVec<[usize; INLINE_EDGES]> = SmallVec::from_elem(0, legs.len());
225
226 let mut out = Vec::with_capacity(wanted);
227 for _ in 0..wanted {
228 let mut path = Path::new();
229 for (leg, (pair, &pick)) in legs
230 .iter()
231 .zip(token_path.windows(2).zip(chosen.iter()))
232 {
233 let pool = leg
234 .get(pick)
235 .expect("odometer holds every leg inside its own pool count");
236 path.add_hop(&self[pair[0]], pool, &self[pair[1]]);
237 }
238 out.push(path);
239
240 for (pick, leg) in chosen.iter_mut().zip(legs.iter()).rev() {
241 *pick += 1;
242 if *pick < leg.len() {
243 break;
244 }
245 *pick = 0;
246 }
247 }
248
249 out
250 }
251
252 fn add_token(&mut self, address: Address) -> NodeIndex {
254 if let Some(index) = self.get_token_ix(&address) {
255 return index;
256 }
257 let index = self.graph.add_node(address.clone());
258 self.tokens.insert(address, index);
259 index
260 }
261
262 fn bidirectional_search(
281 &self,
282 from: NodeIndex,
283 to: NodeIndex,
284 filter: &GraphQueryFilter,
285 ) -> Vec<TokenPath> {
286 let connector_tokens = filter.connector_tokens.as_ref();
287 let endpoints = (from, to);
288 let head_levels =
289 self.walk_levels(from, filter.max_hops.div_ceil(2), endpoints, connector_tokens);
290 let tail_levels = self.walk_levels(to, filter.max_hops / 2, endpoints, connector_tokens);
291
292 let mut midpoint_index: Vec<Option<FxHashMap<NodeIndex, Vec<TokenPath>>>> =
294 vec![None; tail_levels.len()];
295 let mut token_paths = Vec::new();
296
297 for length in filter.min_hops..=filter.max_hops {
298 let head_hops = length.div_ceil(2);
299 let tail_hops = length - head_hops;
300 let (Some(heads), Some(tails)) =
301 (head_levels.get(head_hops), tail_levels.get(tail_hops))
302 else {
303 continue;
304 };
305 if heads.is_empty() || tails.is_empty() {
306 continue;
307 }
308
309 let tails_by_midpoint = midpoint_index[tail_hops].get_or_insert_with(|| {
310 let mut index: FxHashMap<NodeIndex, Vec<TokenPath>> = FxHashMap::default();
311 for tail in tails {
312 let Some(&midpoint) = tail.last() else {
313 continue;
314 };
315 index
316 .entry(midpoint)
317 .or_default()
318 .push(tail.iter().rev().copied().collect());
319 }
320 index
321 });
322
323 for head in heads {
324 let Some(&midpoint) = head.last() else {
325 continue;
326 };
327 let Some(candidates) = tails_by_midpoint.get(&midpoint) else {
328 continue;
329 };
330 for tail in candidates {
331 let collides = head
334 .iter()
335 .any(|token| *token != midpoint && tail.contains(token));
336 if collides {
337 continue;
338 }
339
340 let mut joined = TokenPath::from_slice(head);
341 joined.extend_from_slice(&tail[1..]);
342 token_paths.push(joined);
343 }
344 }
345 }
346
347 token_paths
348 }
349
350 fn walk_levels(
359 &self,
360 start: NodeIndex,
361 hops: usize,
362 endpoints: (NodeIndex, NodeIndex),
363 connector_tokens: Option<&FxHashSet<Address>>,
364 ) -> Vec<Vec<TokenPath>> {
365 let (from, to) = endpoints;
366 let mut levels = vec![vec![TokenPath::from_slice(&[start])]];
367
368 for hop in 0..hops {
369 let mut next = Vec::new();
370 for sequence in &levels[hop] {
371 let Some(&last) = sequence.last() else {
372 continue;
373 };
374 for neighbor in self.neighbors(last) {
375 if sequence.contains(&neighbor) {
376 continue;
377 }
378 if neighbor != from && neighbor != to {
379 if let Some(allowed) = connector_tokens {
380 if !allowed.contains(&self[neighbor]) {
381 continue;
382 }
383 }
384 }
385
386 let mut extended = TokenPath::from_slice(sequence);
387 extended.push(neighbor);
388 next.push(extended);
389 }
390 }
391 levels.push(next);
392 }
393
394 levels
395 }
396
397 fn circular_token_paths(&self, target: NodeIndex, filter: &GraphQueryFilter) -> Vec<TokenPath> {
403 let mut token_paths = Vec::new();
404 let mut frontier = vec![TokenPath::from_slice(&[target])];
405
406 for hops in 1..=filter.max_hops {
407 let mut next = Vec::new();
408 for sequence in &frontier {
409 let Some(&last) = sequence.last() else {
410 continue;
411 };
412 for neighbor in self.neighbors(last) {
413 if neighbor == target {
414 if hops >= filter.min_hops {
415 let mut closed = TokenPath::from_slice(sequence);
416 closed.push(neighbor);
417 token_paths.push(closed);
418 }
419 continue;
423 }
424 if sequence.contains(&neighbor) {
425 continue;
426 }
427 if let Some(allowed) = filter.connector_tokens.as_ref() {
428 if !allowed.contains(&self[neighbor]) {
429 continue;
430 }
431 }
432
433 let mut extended = TokenPath::from_slice(sequence);
434 extended.push(neighbor);
435 next.push(extended);
436 }
437 }
438 frontier = next;
439 }
440
441 token_paths
442 }
443}
444
445impl<D> Deref for TopologyGraph<D> {
446 type Target = TokenGraph<D>;
447
448 fn deref(&self) -> &Self::Target {
449 &self.graph
450 }
451}
452
453impl<D> Default for TopologyGraph<D> {
454 fn default() -> Self {
455 Self {
456 graph: TokenGraph::default(),
457 tokens: FxHashMap::default(),
458 pair_index: FxHashMap::default(),
459 }
460 }
461}
462
463pub struct TopologyGraphManager<D: Clone> {
467 graph: TopologyGraph<D>,
468 component_pairs: FxHashMap<ComponentId, Vec<(NodeIndex, NodeIndex)>>,
471}
472
473impl<D: Clone> TopologyGraphManager<D> {
474 pub fn new() -> Self {
476 Self { graph: TopologyGraph::default(), component_pairs: FxHashMap::default() }
477 }
478
479 fn add_component_edges(&mut self, component_id: &ComponentId, nodes: &[NodeIndex]) {
481 let pairs: Vec<(NodeIndex, NodeIndex)> = nodes
482 .iter()
483 .enumerate()
484 .flat_map(|(i, &from)| {
485 nodes
486 .iter()
487 .skip(i + 1)
488 .flat_map(move |&to| [(from, to), (to, from)])
489 })
490 .collect();
491
492 for &(from, to) in &pairs {
493 self.graph
494 .add_component(from, to, component_id);
495 }
496 self.component_pairs
497 .insert(component_id.clone(), pairs);
498 }
499
500 fn add_components(
507 &mut self,
508 components: &FxHashMap<ComponentId, Vec<Address>>,
509 ) -> Result<(), GraphError> {
510 let mut invalid = Vec::new();
511 let mut skipped = 0usize;
512
513 let mut sorted: Vec<_> = components.iter().collect();
515 sorted.sort_by_key(|(id, _)| *id);
516
517 for (component_id, tokens) in sorted {
518 if self
519 .component_pairs
520 .contains_key(component_id)
521 {
522 trace!(component_id = %component_id, "skipping already-tracked component");
523 skipped += 1;
524 continue;
525 }
526 if tokens.len() < 2 {
527 invalid.push(component_id.clone());
528 continue;
529 }
530
531 let mut sorted_tokens: Vec<&Address> = tokens.iter().collect();
532 sorted_tokens.sort();
533 let nodes: Vec<NodeIndex> = sorted_tokens
534 .iter()
535 .map(|token| self.graph.add_token((*token).clone()))
536 .collect();
537 self.add_component_edges(component_id, &nodes);
538 }
539
540 if skipped > 0 {
541 debug!(skipped_duplicates = skipped, "skipped duplicate components during add");
542 }
543 if !invalid.is_empty() {
544 return Err(GraphError::InvalidComponents(invalid));
545 }
546 Ok(())
547 }
548
549 fn remove_components(&mut self, components: &[ComponentId]) -> Result<(), GraphError> {
556 let mut missing = Vec::new();
557
558 for component_id in components {
559 let Some(pairs) = self
560 .component_pairs
561 .remove(component_id)
562 else {
563 missing.push(component_id.clone());
564 continue;
565 };
566
567 for (from, to) in pairs {
568 self.graph
569 .remove_component(from, to, component_id);
570 }
571 }
572
573 if !missing.is_empty() {
574 return Err(GraphError::ComponentsNotFound(missing));
575 }
576 Ok(())
577 }
578
579 #[cfg(test)]
581 pub(crate) fn set_pool_weight(
582 &mut self,
583 component_id: &ComponentId,
584 token_in: &Address,
585 token_out: &Address,
586 data: D,
587 bidirectional: bool,
588 ) -> Result<(), GraphError> {
589 let from = self
590 .graph
591 .get_token_ix(token_in)
592 .ok_or_else(|| GraphError::TokenNotFound(token_in.clone()))?;
593 let to = self
594 .graph
595 .get_token_ix(token_out)
596 .ok_or_else(|| GraphError::TokenNotFound(token_out.clone()))?;
597
598 let mut directions = vec![(from, to)];
599 if bidirectional {
600 directions.push((to, from));
601 }
602
603 let mut updated = false;
604 for (source, target) in directions {
605 let Some(edge) = self.graph.get_edge_ix(source, target) else {
606 continue;
607 };
608 if let Some(pool) = self
609 .graph
610 .graph
611 .edge_weight_mut(edge)
612 .and_then(|pair| pair.pool_mut(component_id))
613 {
614 pool.data = Some(data.clone());
615 updated = true;
616 }
617 }
618
619 if !updated {
620 return Err(GraphError::MissingComponentBetweenTokens(
621 token_in.clone(),
622 token_out.clone(),
623 component_id.clone(),
624 ));
625 }
626 Ok(())
627 }
628}
629
630impl<D: Clone + super::EdgeWeightFromSimAndDerived> super::EdgeWeightUpdaterWithDerived
631 for TopologyGraphManager<D>
632{
633 fn update_edge_weights_with_derived(
638 &mut self,
639 market: MarketDataView<'_>,
640 derived: &crate::derived::DerivedData,
641 ) -> usize {
642 let tokens = market.token_registry_ref();
643 let mut updated = 0usize;
644
645 for edge in self
646 .graph
647 .edge_indices()
648 .collect::<Vec<_>>()
649 {
650 let Some((source, target)) = self.graph.edge_endpoints(edge) else {
651 continue;
652 };
653 let (Some(token_in), Some(token_out)) =
656 (tokens.get(&self.graph[source]), tokens.get(&self.graph[target]))
657 else {
658 continue;
659 };
660
661 let Some(pair) = self.graph.graph.edge_weight_mut(edge) else {
662 continue;
663 };
664 for pool in &mut pair.pools {
665 pool.data = market
666 .get_simulation_state(&pool.component_id)
667 .and_then(|state| {
668 D::from_sim_and_derived(
669 state,
670 &pool.component_id,
671 token_in,
672 token_out,
673 derived,
674 )
675 });
676 if pool.data.is_some() {
677 updated += 1;
678 }
679 }
680 }
681
682 updated
683 }
684}
685
686impl<D: Clone> Default for TopologyGraphManager<D> {
687 fn default() -> Self {
688 Self::new()
689 }
690}
691
692impl<D: Clone + Send + Sync> GraphManager<TopologyGraph<D>> for TopologyGraphManager<D> {
693 fn initialize_graph(&mut self, component_topology: &FxHashMap<ComponentId, Vec<Address>>) {
694 self.graph = TopologyGraph::default();
695 self.component_pairs.clear();
696
697 let mut tokens: Vec<Address> = component_topology
700 .values()
701 .flat_map(|addresses| addresses.iter())
702 .cloned()
703 .collect::<FxHashSet<_>>()
704 .into_iter()
705 .collect();
706 tokens.sort();
707
708 for token in tokens {
709 self.graph.add_token(token);
710 }
711
712 if let Err(e) = self.add_components(component_topology) {
716 debug!(error = %e, "components skipped while building the graph");
717 }
718 }
719
720 fn graph(&self) -> &TopologyGraph<D> {
721 &self.graph
722 }
723}
724
725#[async_trait]
726impl<D: Clone + Send> MarketEventHandler for TopologyGraphManager<D> {
727 async fn handle_event(&mut self, event: &MarketEvent) -> Result<(), EventError> {
728 match event {
729 MarketEvent::MarketUpdated { added_components, removed_components, .. } => {
730 let mut errors = Vec::new();
731 if let Err(e) = self.add_components(added_components) {
732 errors.push(e);
733 }
734 if let Err(e) = self.remove_components(removed_components) {
735 errors.push(e);
736 }
737 if errors.is_empty() {
738 Ok(())
739 } else {
740 Err(EventError::GraphErrors(errors))
741 }
742 }
743 }
744 }
745}
746
747#[cfg(test)]
748mod tests {
749 use rstest::rstest;
750 use rustc_hash::{FxHashMap, FxHashSet};
751
752 use super::*;
753 use crate::{
754 algorithm::{
755 most_liquid::DepthAndPrice,
756 test_utils::fixtures::{addrs, diamond_graph, linear_graph, parallel_graph},
757 },
758 graph::{EdgeWeightUpdaterWithDerived, GraphManager},
759 };
760
761 fn addr(byte: u8) -> Address {
762 Address::from([byte; 20])
763 }
764
765 #[tokio::test]
772 async fn test_pair_edge_lives_from_first_pool_to_last() {
773 use crate::feed::events::{MarketEvent, MarketEventHandler};
774
775 let (a, b) = (addr(0x0A), addr(0x0B));
776 let mut manager = TopologyGraphManager::<()>::new();
777 manager.initialize_graph(&FxHashMap::from_iter([(
778 "first".to_string(),
779 vec![a.clone(), b.clone()],
780 )]));
781
782 let (from, to) = (
783 manager
784 .graph()
785 .get_token_ix(&a)
786 .unwrap(),
787 manager
788 .graph()
789 .get_token_ix(&b)
790 .unwrap(),
791 );
792 assert_eq!(manager.graph().edge_count(), 2, "one edge each way");
793 assert_eq!(
794 manager
795 .graph()
796 .pools_between(from, to)
797 .len(),
798 1
799 );
800
801 let added = |id: &str| MarketEvent::MarketUpdated {
802 added_components: FxHashMap::from_iter([(id.to_string(), vec![a.clone(), b.clone()])]),
803 removed_components: vec![],
804 updated_components: vec![],
805 };
806 let removed = |id: &str| MarketEvent::MarketUpdated {
807 added_components: FxHashMap::default(),
808 removed_components: vec![id.to_string()],
809 updated_components: vec![],
810 };
811
812 manager
814 .handle_event(&added("second"))
815 .await
816 .unwrap();
817 assert_eq!(manager.graph().edge_count(), 2, "a second pool is not a second edge");
818 assert_eq!(
819 manager
820 .graph()
821 .pools_between(from, to)
822 .len(),
823 2
824 );
825
826 manager
828 .handle_event(&removed("second"))
829 .await
830 .unwrap();
831 assert_eq!(manager.graph().edge_count(), 2, "the pair still trades");
832 assert_eq!(
833 manager
834 .graph()
835 .pools_between(from, to)
836 .len(),
837 1
838 );
839
840 manager
842 .handle_event(&removed("first"))
843 .await
844 .unwrap();
845 assert_eq!(manager.graph().edge_count(), 0);
846 assert!(manager
847 .graph()
848 .pools_between(from, to)
849 .is_empty());
850
851 manager
853 .handle_event(&added("third"))
854 .await
855 .unwrap();
856 assert_eq!(manager.graph().edge_count(), 2);
857 assert_eq!(
858 manager
859 .graph()
860 .pools_between(from, to)
861 .len(),
862 1
863 );
864 assert_eq!(manager.graph().pools_between(from, to)[0].component_id, "third");
865
866 let (c, d) = (addr(0x0C), addr(0x0D));
868 manager
869 .handle_event(&MarketEvent::MarketUpdated {
870 added_components: FxHashMap::from_iter([(
871 "fourth".to_string(),
872 vec![c.clone(), d.clone()],
873 )]),
874 removed_components: vec![],
875 updated_components: vec![],
876 })
877 .await
878 .unwrap();
879 assert_eq!(manager.graph().node_count(), 4);
880 assert_eq!(manager.graph().edge_count(), 4);
881 }
882
883 #[test]
884 fn test_edge_weight_cleared_on_spot_price_miss() {
885 use num_bigint::BigUint;
888 use num_traits::One;
889 use tycho_simulation::tycho_core::simulation::protocol_sim::Price;
890
891 use crate::{
892 algorithm::test_utils::{market_read, setup_market_weighted, token, MockProtocolSim},
893 derived::{types::TokenGasPrices, DerivedData},
894 };
895
896 let token_a = token(0x01, "A");
897 let token_b = token(0x02, "B");
898 let (market, mut manager) = setup_market_weighted(vec![(
899 "component1",
900 &token_a,
901 &token_b,
902 MockProtocolSim::new(2.0),
903 )]);
904
905 assert!(
906 manager
907 .graph()
908 .edge_indices()
909 .all(|e| manager
910 .graph()
911 .edge_weight(e)
912 .unwrap()
913 .pools()
914 .iter()
915 .all(|pool| pool.data.is_some())),
916 "edges should have weight data after setup"
917 );
918
919 let mut token_prices = TokenGasPrices::default();
920 for addr in [&token_a.address, &token_b.address] {
921 token_prices.insert(
922 addr.clone(),
923 Price { numerator: BigUint::one(), denominator: BigUint::one() },
924 );
925 }
926 let mut derived = DerivedData::new();
927 derived.set_spot_prices(Default::default(), vec![], 10, true);
928 derived.set_component_depths(Default::default(), vec![], 10, true);
929 derived.set_token_prices(token_prices, vec![], 10, true);
930
931 manager.update_edge_weights_with_derived(market_read(&market), &derived);
932
933 assert!(
934 manager
935 .graph()
936 .edge_indices()
937 .all(|e| manager
938 .graph()
939 .edge_weight(e)
940 .unwrap()
941 .pools()
942 .iter()
943 .all(|pool| pool.data.is_none())),
944 "stale edge weights must be cleared when spot price is unavailable"
945 );
946 }
947
948 fn routes<'a>(
953 graph: &'a TopologyGraph<DepthAndPrice>,
954 from: &Address,
955 to: &Address,
956 min_hops: usize,
957 max_hops: usize,
958 connector_tokens: Option<FxHashSet<Address>>,
959 ) -> Vec<Path<'a, DepthAndPrice>> {
960 let (Some(from), Some(to)) = (graph.get_token_ix(from), graph.get_token_ix(to)) else {
961 return Vec::new();
962 };
963 let filter = GraphQueryFilter { min_hops, max_hops, connector_tokens };
964 graph
965 .paths_between_ix(from, to, &filter)
966 .iter()
967 .flat_map(|token_path| graph.expand_path(token_path, None))
968 .collect()
969 }
970
971 fn all_ids(paths: Vec<Path<'_, DepthAndPrice>>) -> FxHashSet<Vec<&str>> {
972 paths
973 .iter()
974 .map(|p| {
975 p.iter()
976 .map(|(_, e, _)| e.component_id.as_str())
977 .collect()
978 })
979 .collect()
980 }
981
982 #[test]
983 fn test_find_paths_linear_forward_and_reverse() {
984 let (a, b, c, d) = addrs();
985 let m = linear_graph();
986 let g = m.graph();
987
988 let p = routes(g, &a, &b, 1, 1, None);
990 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab"]]));
991
992 let p = routes(g, &a, &c, 1, 2, None);
993 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab", "bc"]]));
994
995 let p = routes(g, &a, &d, 1, 3, None);
996 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab", "bc", "cd"]]));
997
998 let p = routes(g, &d, &a, 1, 3, None);
1000 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["cd", "bc", "ab"]]));
1001 }
1002
1003 #[test]
1004 fn test_find_paths_respects_hop_bounds() {
1005 let (a, _, c, d) = addrs();
1006 let m = linear_graph();
1007 let g = m.graph();
1008
1009 assert!(routes(g, &a, &d, 1, 2, None).is_empty());
1011
1012 assert!(routes(g, &a, &c, 3, 3, None).is_empty());
1014 }
1015
1016 #[test]
1017 fn test_find_paths_parallel_components() {
1018 let (a, b, c, _) = addrs();
1019 let m = parallel_graph();
1020 let g = m.graph();
1021
1022 let p = routes(g, &a, &b, 1, 1, None);
1024 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab1"], vec!["ab2"], vec!["ab3"]]));
1025
1026 let p = routes(g, &a, &c, 1, 2, None);
1028 assert_eq!(
1029 all_ids(p),
1030 FxHashSet::from_iter([
1031 vec!["ab1", "bc1"],
1032 vec!["ab1", "bc2"],
1033 vec!["ab2", "bc1"],
1034 vec!["ab2", "bc2"],
1035 vec!["ab3", "bc1"],
1036 vec!["ab3", "bc2"],
1037 ])
1038 );
1039 }
1040
1041 #[test]
1042 fn test_find_paths_diamond_multiple_routes() {
1043 let (a, _, _, d) = addrs();
1044 let m = diamond_graph();
1045 let g = m.graph();
1046
1047 let p = routes(g, &a, &d, 1, 2, None);
1049 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab", "bd"], vec!["ac", "cd"]]));
1050 }
1051
1052 #[test]
1053 fn test_find_paths_no_intermediate_cycles() {
1054 let (a, b, _, _) = addrs();
1055 let m = linear_graph();
1056 let g = m.graph();
1057
1058 let p = routes(g, &a, &b, 1, 3, None);
1063 assert_eq!(all_ids(p), FxHashSet::from_iter([vec!["ab"]]));
1064 }
1065
1066 #[test]
1067 fn test_find_paths_cyclic_same_source_dest() {
1068 let (a, _, _, _) = addrs();
1069 let m = parallel_graph();
1071 let g = m.graph();
1072
1073 let p = routes(g, &a, &a, 2, 2, None);
1076 assert_eq!(
1077 all_ids(p),
1078 FxHashSet::from_iter([
1079 vec!["ab1", "ab1"],
1080 vec!["ab1", "ab2"],
1081 vec!["ab1", "ab3"],
1082 vec!["ab2", "ab1"],
1083 vec!["ab2", "ab2"],
1084 vec!["ab2", "ab3"],
1085 vec!["ab3", "ab1"],
1086 vec!["ab3", "ab2"],
1087 vec!["ab3", "ab3"],
1088 ])
1089 );
1090 }
1091
1092 #[tokio::test]
1095 async fn test_add_components_reports_the_ones_with_too_few_tokens() {
1096 let (a, b) = (addr(0x0A), addr(0x0B));
1097 let mut manager = TopologyGraphManager::<()>::new();
1098 manager.initialize_graph(&FxHashMap::default());
1099
1100 let result = manager.add_components(&FxHashMap::from_iter([
1101 ("solo".to_string(), vec![a.clone()]),
1102 ("pair".to_string(), vec![a.clone(), b.clone()]),
1103 ]));
1104
1105 match result {
1106 Err(GraphError::InvalidComponents(invalid)) => {
1107 assert_eq!(invalid, vec!["solo".to_string()]);
1108 }
1109 other => panic!("expected InvalidComponents, got {other:?}"),
1110 }
1111 assert_eq!(manager.graph().edge_count(), 2, "the valid pair was still added");
1112 }
1113
1114 #[tokio::test]
1116 async fn test_remove_components_reports_the_ones_it_does_not_hold() {
1117 let (a, b) = (addr(0x0A), addr(0x0B));
1118 let mut manager = TopologyGraphManager::<()>::new();
1119 manager.initialize_graph(&FxHashMap::from_iter([(
1120 "pair".to_string(),
1121 vec![a.clone(), b.clone()],
1122 )]));
1123
1124 let result = manager.remove_components(&["pair".to_string(), "ghost".to_string()]);
1125
1126 match result {
1127 Err(GraphError::ComponentsNotFound(missing)) => {
1128 assert_eq!(missing, vec!["ghost".to_string()]);
1129 }
1130 other => panic!("expected ComponentsNotFound, got {other:?}"),
1131 }
1132 assert_eq!(manager.graph().edge_count(), 0, "the component it did hold was removed");
1133 }
1134
1135 #[test]
1137 fn test_set_pool_weight_rejects_a_pair_the_component_does_not_trade() {
1138 let (a, b, c, _) = addrs();
1139 let mut manager = linear_graph();
1140
1141 let unknown_token = manager.set_pool_weight(
1142 &"ab".to_string(),
1143 &a,
1144 &addr(0x99),
1145 DepthAndPrice::new(1.0, 1.0),
1146 false,
1147 );
1148 assert!(matches!(unknown_token, Err(GraphError::TokenNotFound(_))));
1149
1150 let wrong_pair =
1151 manager.set_pool_weight(&"ab".to_string(), &b, &c, DepthAndPrice::new(1.0, 1.0), false);
1152 assert!(matches!(wrong_pair, Err(GraphError::MissingComponentBetweenTokens(..))));
1153 }
1154
1155 #[rstest]
1157 #[case::zero_min(0, 3)]
1158 #[case::min_above_max(3, 1)]
1159 fn test_paths_between_rejects_impossible_hop_bounds(
1160 #[case] min_hops: usize,
1161 #[case] max_hops: usize,
1162 ) {
1163 let (a, b, _, _) = addrs();
1164 let m = linear_graph();
1165 let g = m.graph();
1166 let (from, to) = (g.get_token_ix(&a).unwrap(), g.get_token_ix(&b).unwrap());
1167
1168 let filter = GraphQueryFilter { min_hops, max_hops, connector_tokens: None };
1169
1170 assert!(g
1171 .paths_between_ix(from, to, &filter)
1172 .is_empty());
1173 assert!(
1174 g.paths_between_ix(from, from, &filter)
1175 .is_empty(),
1176 "the cyclic search is bound by the same rule"
1177 );
1178 }
1179
1180 #[test]
1184 fn test_cyclic_routes_never_pass_through_the_start_token() {
1185 let (a, _, _, _) = addrs();
1186 let m = diamond_graph();
1187 let g = m.graph();
1188 let start = g.get_token_ix(&a).unwrap();
1189
1190 let cycles = g.paths_between_ix(
1191 start,
1192 start,
1193 &GraphQueryFilter { min_hops: 1, max_hops: 4, connector_tokens: None },
1194 );
1195
1196 assert!(!cycles.is_empty(), "the diamond closes cycles through B and through C");
1197 for cycle in &cycles {
1198 assert_eq!(cycle.first(), Some(&start), "a cycle starts on its token");
1199 assert_eq!(cycle.last(), Some(&start), "a cycle ends on its token");
1200 assert!(
1201 !cycle[1..cycle.len() - 1].contains(&start),
1202 "the start token must appear only at the two ends, got {cycle:?}"
1203 );
1204 }
1205 }
1206
1207 fn stacked_graph() -> TopologyGraphManager<DepthAndPrice> {
1219 let (s, x, y, t) = (addr(0x51), addr(0x58), addr(0x59), addr(0x54));
1220 let mut topology = FxHashMap::default();
1221 for (id, from, to) in [
1222 ("sx1", &s, &x),
1223 ("sx2", &s, &x),
1224 ("xy1", &x, &y),
1225 ("xy2", &x, &y),
1226 ("xy3", &x, &y),
1227 ("yt1", &y, &t),
1228 ("yt2", &y, &t),
1229 ("sy1", &s, &y),
1230 ("st1", &s, &t),
1231 ] {
1232 topology.insert(id.to_string(), vec![from.clone(), to.clone()]);
1233 }
1234
1235 let mut manager = TopologyGraphManager::<DepthAndPrice>::new();
1236 manager.initialize_graph(&topology);
1237 manager
1238 }
1239
1240 #[test]
1241 fn test_find_paths_expands_every_pool_combination() {
1242 let (s, x, y, t) = (addr(0x51), addr(0x58), addr(0x59), addr(0x54));
1243 let manager = stacked_graph();
1244 let graph = manager.graph();
1245
1246 assert_eq!(routes(graph, &s, &t, 1, 1, None).len(), 1);
1248 assert_eq!(routes(graph, &s, &t, 2, 2, None).len(), 2);
1250 assert_eq!(routes(graph, &s, &t, 3, 3, None).len(), 12);
1252
1253 let paths = routes(graph, &s, &t, 1, 3, None);
1254 assert_eq!(paths.len(), 15, "every combination, and none of them twice");
1255
1256 let long: FxHashSet<Vec<&str>> = paths
1259 .iter()
1260 .filter(|path| path.len() == 3)
1261 .map(|path| {
1262 path.edge_iter()
1263 .iter()
1264 .map(|edge| edge.component_id.as_str())
1265 .collect()
1266 })
1267 .collect();
1268 assert_eq!(long.len(), 12);
1269
1270 for path in &paths {
1271 assert_eq!(path.tokens.first().copied(), Some(&s));
1272 assert_eq!(path.tokens.last().copied(), Some(&t));
1273 assert_eq!(path.tokens.len(), path.len() + 1);
1274 }
1275
1276 let allowed = FxHashSet::from_iter([x]);
1278 let filtered = routes(graph, &s, &t, 1, 3, Some(allowed.clone()));
1279 assert_eq!(filtered.len(), 1);
1280 assert_eq!(filtered[0].edge_iter()[0].component_id, "st1");
1281
1282 let allowed = FxHashSet::from_iter([y]);
1283 assert_eq!(
1284 routes(graph, &s, &t, 1, 3, Some(allowed.clone())).len(),
1285 3,
1286 "st1, plus sy1 over each of the two Y-T pools"
1287 );
1288 }
1289
1290 #[test]
1291 fn test_find_paths_bfs_ordering() {
1292 let (a, b, c, d) = addrs();
1297 let e = addr(0x0E);
1298 let mut m = TopologyGraphManager::<DepthAndPrice>::new();
1299 let mut t = FxHashMap::default();
1300 t.insert("ae".into(), vec![a.clone(), e.clone()]);
1301 t.insert("ab".into(), vec![a.clone(), b.clone()]);
1302 t.insert("be".into(), vec![b, e.clone()]);
1303 t.insert("ac".into(), vec![a.clone(), c.clone()]);
1304 t.insert("cd".into(), vec![c, d.clone()]);
1305 t.insert("de".into(), vec![d, e.clone()]);
1306 m.initialize_graph(&t);
1307 let g = m.graph();
1308
1309 let p = routes(g, &a, &e, 1, 3, None);
1310
1311 assert_eq!(p.len(), 3, "Expected 3 paths total");
1313 assert_eq!(p[0].len(), 1, "First path should be 1-hop");
1314 assert_eq!(p[1].len(), 2, "Second path should be 2-hop");
1315 assert_eq!(p[2].len(), 3, "Third path should be 3-hop");
1316 }
1317
1318 #[test]
1319 fn test_connector_tokens_blocks_disallowed_intermediate() {
1320 let (a, b, c, d) = addrs();
1322 let m = diamond_graph();
1323 let g = m.graph();
1324 let allowed: FxHashSet<Address> = FxHashSet::from_iter([c.clone()]);
1325 let paths = routes(g, &a, &d, 1, 2, Some(allowed.clone()));
1326 let intermediates: FxHashSet<&Address> = paths
1327 .iter()
1328 .flat_map(|p| p.iter().map(|(node, _, _)| node))
1329 .filter(|addr| *addr != &a && *addr != &d)
1330 .collect();
1331 assert!(!intermediates.contains(&b), "B should be blocked");
1333 assert!(intermediates.contains(&c), "C should be allowed");
1334 }
1335
1336 #[test]
1337 fn test_connector_tokens_allows_endpoints_even_if_not_listed() {
1338 let (a, b, _, _) = addrs();
1340 let m = linear_graph();
1341 let g = m.graph();
1342 let allowed: FxHashSet<Address> = FxHashSet::default();
1345 let paths = routes(g, &a, &b, 1, 1, Some(allowed.clone()));
1346 assert!(!paths.is_empty(), "1-hop direct route should survive empty allowlist");
1347 }
1348
1349 #[test]
1350 fn test_connector_tokens_none_is_unrestricted() {
1351 let (a, b, c, d) = addrs();
1353 let m = diamond_graph();
1354 let g = m.graph();
1355 let paths = routes(g, &a, &d, 1, 2, None);
1356 let intermediates: FxHashSet<&Address> = paths
1357 .iter()
1358 .flat_map(|p| p.iter().map(|(node, _, _)| node))
1359 .filter(|addr| *addr != &a && *addr != &d)
1360 .collect();
1361 assert!(intermediates.contains(&b), "B should appear with no restriction");
1362 assert!(intermediates.contains(&c), "C should appear with no restriction");
1363 }
1364}