Skip to main content

fynd_core/graph/
token_graph.rs

1//! A market graph with one edge per token pair. The pools that trade the pair are the edge weight.
2//!
3//! [`super::petgraph`] holds the same market with one edge per pool. Use that one to walk or relax
4//! pools. Use this one to find routes as sequences of tokens.
5
6use 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/// The pools that trade one directed token pair.
28#[derive(Debug, Clone)]
29pub struct PairEdge<D> {
30    /// In the order they were added, so a search reads them the same way on every run.
31    pools: Vec<EdgeData<D>>,
32}
33
34impl<D> PairEdge<D> {
35    /// The pools serving this pair.
36    pub fn pools(&self) -> &[EdgeData<D>] {
37        &self.pools
38    }
39
40    /// Adds a pool, or does nothing if this component already serves the pair.
41    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    /// Drops a pool. Returns whether the pair has any left.
54    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    /// The pool this component runs on this pair.
61    #[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
69/// A route as a sequence of tokens, before the pools serving each leg are chosen.
70pub type TokenPath = SmallVec<[NodeIndex; INLINE_TOKENS]>;
71
72/// The pools of one leg a solve may use.
73///
74/// A pair the request leaves alone stays a slice into the graph. One it excludes a pool from
75/// cannot, because what is left is not contiguous, so those pools are collected.
76enum Leg<'a, D> {
77    /// Every pool the pair holds.
78    All(&'a [EdgeData<D>]),
79    /// The pools the request allows.
80    Allowed(Vec<&'a EdgeData<D>>),
81}
82
83impl<'a, D> Leg<'a, D> {
84    /// How many pools this leg may use.
85    fn len(&self) -> usize {
86        match self {
87            Self::All(pools) => pools.len(),
88            Self::Allowed(pools) => pools.len(),
89        }
90    }
91
92    /// Whether the leg has no pool a solve may use.
93    fn is_empty(&self) -> bool {
94        self.len() == 0
95    }
96
97    /// The `index`th pool this leg may use.
98    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
106/// Tokens as nodes, one edge per directed token pair.
107pub type TokenGraph<D> = stable_graph::StableDiGraph<Address, PairEdge<D>>;
108
109/// A [`TokenGraph`] with a lookup from a token pair to its edge.
110///
111/// Derefs to the graph, so the usual `petgraph` reads work on it directly.
112pub struct TopologyGraph<D> {
113    /// Taken mutably to change edge weights, which does not change which tokens trade and so
114    /// leaves `pair_index` correct. Adding or removing an edge must go through `add_component` or
115    /// `remove_component`, which keep it up to date.
116    graph: TokenGraph<D>,
117    /// The node holding each token.
118    tokens: FxHashMap<Address, NodeIndex>,
119    /// The edge for each directed token pair. Only changes when an edge is added or removed, which
120    /// only happens when a pair gains its first pool or loses its last.
121    pair_index: FxHashMap<(NodeIndex, NodeIndex), EdgeIndex>,
122}
123
124impl<D> TopologyGraph<D> {
125    /// The pools trading `from` for `to`, empty if the pair is not connected.
126    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    /// The pools trading `from` for `to` that this solve may use.
134    ///
135    /// Collects only when the request excludes one of them.
136    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    /// Whether the pair trades `from` for `to` through a pool this solve may use.
159    ///
160    /// `false` for a pair the graph does not hold, and for one whose every pool the request
161    /// excludes. Either way the pair does not connect its two tokens for this search.
162    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    /// The node holding `token`, or `None` if the market has no such token.
178    pub fn get_token_ix(&self, token: &Address) -> Option<NodeIndex> {
179        self.tokens.get(token).copied()
180    }
181
182    /// The edge between two tokens, or `None` if they do not trade.
183    fn get_edge_ix(&self, from: NodeIndex, to: NodeIndex) -> Option<EdgeIndex> {
184        self.pair_index
185            .get(&(from, to))
186            .copied()
187    }
188
189    /// Records that `component_id` trades `from` for `to`, adding the pair's edge if it is the
190    /// first pool to do so.
191    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    /// Records that `component_id` no longer trades `from` for `to`, removing the pair's edge if
207    /// it was the last pool doing so.
208    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    /// Every route between two tokens, as token sequences.
223    ///
224    /// See [`TopologyGraph::paths_between_ix`], which this resolves the addresses for.
225    ///
226    /// # Errors
227    ///
228    /// [`GraphError::TokenNotFound`] naming whichever of the two the market does not hold.
229    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    /// Every token path from `from` to `to` within the filter's hop bounds.
245    ///
246    /// Empty when there is no route. Check [`TopologyGraph::expand_path`] for expanded pool paths.
247    ///
248    /// Routes come back shortest first. Within a length they are in no particular order.
249    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    /// Writes out one route per combination of pools along `token_path`.
267    ///
268    /// A token sequence stands for as many routes as the product of the pools on each of its legs.
269    /// They are enumerated by counting: the rightmost leg advances first, and carries into the leg
270    /// to its left when it wraps.
271    ///
272    /// That product is unbounded — four legs of twenty pools is 160,000 routes — so `max_paths`
273    /// caps how many are written. The ones past the cap are dropped in counting order, which no
274    /// ranking has seen yet, so a cap trades routes the caller might have wanted for a bound on
275    /// what one sequence can allocate. `None` writes them all.
276    pub fn expand_path(
277        &self,
278        token_path: &[NodeIndex],
279        max_paths: Option<usize>,
280        exclusions: &RouteExclusions,
281    ) -> Vec<Path<'_, D>> {
282        // Inline, like everything else in the search: this runs once per token sequence found and
283        // a route has at most `max_hops` legs. Each leg holds only the pools the request allows,
284        // so an excluded one is not counted into the product either.
285        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        // A sequence of one token names no leg, and a leg with no pool names a pair the graph
291        // does not connect -- the walk disagreeing with the graph rather than a routing outcome.
292        // Either way there is no route to write, and an empty product would otherwise write one
293        // route with no hops.
294        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    /// Adds a token as a node and returns its index, or returns the index it already has.
332    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    /// Every route from `from` to `to`, found by meeting in the middle.
342    ///
343    /// Each length between `min_hops` and `max_hops` is searched on its own. A route of `n` hops is
344    /// split in two: a walk of `ceil(n/2)` hops out of `from` and a walk of `n - ceil(n/2)` hops
345    /// out of `to`. The two are joined wherever they end on the same token, so the deepest
346    /// level is reached by matching rather than by walking. With a branching factor of `b` that
347    /// costs about `b^(n/2)` instead of `b^n`.
348    ///
349    /// Both halves are walked breadth-first by [`TopologyGraph::walk_tokens`]. The tail starts at
350    /// the destination and walks outwards, so it comes back pointing the wrong way and is reversed
351    /// before joining. Only the order needs fixing -- pools are chosen per leg during expansion,
352    /// from the token pair itself, so no edge is ever carried in the wrong direction.
353    ///
354    /// Searching one length at a time keeps shorter routes ahead of longer ones in the result.
355    ///
356    /// Consecutive lengths ask for the same half-walks -- lengths 3 and 4 both take a 2-hop head,
357    /// lengths 4 and 5 both take a 2-hop tail -- so each side is walked once to its deepest and
358    /// every length reads the level it needs.
359    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        // Grouping a tail level by its midpoint is worth doing once per depth, not once per length.
371        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                    // Each half avoids revisits on its own, but together they can name a token
410                    // twice. The midpoint is the only one they are allowed to share.
411                    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    /// The token sequences leaving `start`, one level per hop count: index `k` holds every
429    /// sequence of exactly `k` edges, up to `hops`. Level 0 is `start` on its own.
430    ///
431    /// Walking breadth-first produces every level on the way to the deepest, so they are all kept
432    /// rather than the deepest alone.
433    ///
434    /// `endpoints` is the route's own `(from, to)`. The rule never applies to those two; every
435    /// other token a walk passes through is an intermediate and must be allowed.
436    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    /// Token sequences that begin and end on the same token.
475    ///
476    /// Such a route closes a cycle, so it has no midpoint to split on -- both halves would have to
477    /// start and end on that token, which the no-revisit rule forbids. Searched from the one end
478    /// instead, with the closing hop exempt from that rule.
479    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                        // A closed cycle is a finished route. Walking on from it would put the
501                        // start token in the middle of a longer one, which is not a route the
502                        // executor can take.
503                        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
542/// Builds a [`TopologyGraph`] from the market and keeps it up to date as components come and go.
543///
544/// One per worker.
545pub struct TopologyGraphManager<D: Clone> {
546    graph: TopologyGraph<D>,
547    /// The token pairs each component trades, so removing a component does not mean searching
548    /// every edge for it.
549    component_pairs: FxHashMap<ComponentId, Vec<(NodeIndex, NodeIndex)>>,
550}
551
552impl<D: Clone> TopologyGraphManager<D> {
553    /// Creates an empty manager.
554    pub fn new() -> Self {
555        Self { graph: TopologyGraph::default(), component_pairs: FxHashMap::default() }
556    }
557
558    /// Adds an edge each way between every pair of the component's tokens.
559    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    /// Adds components to the topology graph.
580    ///
581    /// # Errors
582    ///
583    /// [`GraphError::InvalidComponents`] naming the ones with fewer than two tokens. The rest were
584    /// added.
585    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        // Sorted, so nodes and edges land in the same order whatever the map's iteration order.
593        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    /// Removes components.
629    ///
630    /// # Errors
631    ///
632    /// [`GraphError::ComponentsNotFound`] naming the ones the graph does not hold. The rest were
633    /// removed.
634    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    /// Sets one pool's weight, for tests that need a weight without running the derived data.
659    #[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    /// Recomputes every pool's weight from this block's simulation states and derived data.
713    ///
714    /// Returns how many pools ended up with a weight. A pool whose derived data is missing has its
715    /// weight cleared, so no one reads last block's.
716    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            // Both borrow the token registry, not the graph, so the graph is free to be taken
733            // mutably below without either being cloned.
734            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        // Sorted, so the same topology gives the same node indices in every process. Hash iteration
777        // order is seeded per process and would otherwise vary run to run.
778        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        // The same path an incremental update takes, so a graph built here and a graph grown by
792        // events end up identical. A component with fewer than two tokens forms no edge; there is
793        // no caller to report that to at startup, so it is logged.
794        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    /// A pair has one edge however many pools trade it. The edge appears with the first pool and
845    /// goes with the last.
846    ///
847    /// Adding or removing a pool in between must leave both the graph and `pair_index` alone. The
848    /// last step checks that removing the edge also drops its index entry, since a leftover entry
849    /// would point at an edge that no longer exists.
850    #[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        // A second pool on the same pair rides on the edge that is already there.
892        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        // Taking it away again leaves the pair trading, so the edge stays.
906        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        // Taking the last one disconnects the tokens.
920        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        // And the index must have let go with it, or a later pair would read a dead edge.
931        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        // A pool between tokens the graph has never seen brings its nodes with it.
946        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        // Regression: when spot price computation fails, stale edge weights must be cleared so the
965        // component is excluded from path scoring rather than routed with an outdated price.
966        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    /// Every route between two tokens, as one entry per combination of pools.
1028    ///
1029    /// The graph works in nodes; these cases were written in addresses, so this resolves them and
1030    /// runs both halves of what the algorithms do -- search, then expand.
1031    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        // Forward: A->B (1 hop), A->C (2 hops), A->D (3 hops)
1070        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        // Reverse: D->A (bidirectional components)
1080        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        // A->D needs 3 hops, max_hops=2 finds nothing
1091        assert!(routes(g, &a, &d, 1, 2, None).is_empty());
1092
1093        // A->C is 2 hops, min_hops=3 finds nothing
1094        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        // A->B: 3 parallel components = 3 paths
1104        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        // A->C: 3 A->B components × 2 B->C components = 6 paths
1108        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        // A->D: two 2-hop paths
1129        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        // A->B with max_hops=3: only the direct 1-hop path is valid.
1140        // Revisit paths like A->B->C->B or A->B->B->B are pruned because
1141        // they create intermediate cycles unsupported by Tycho execution
1142        // (only first == last cycles are allowed, i.e. from == to).
1143        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        // Use parallel_graph with 3 A<->B components to verify all combinations
1151        let m = parallel_graph();
1152        let g = m.graph();
1153
1154        // A->A (cyclic path) with 2 hops: should find all 9 combinations (3 components × 3
1155        // components) Note: min_hops=2 because cyclic paths require at least 2 hops
1156        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    /// A component needs two tokens to trade. One with fewer is named in the error, and the
1174    /// components that were fine are still added.
1175    #[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    /// Removing a component the graph never held is reported, and the ones it did hold still go.
1196    #[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    /// A weight can only be set on a pair a component actually trades.
1217    #[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    /// Hop bounds that name no length return nothing rather than searching.
1237    #[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    /// A closed cycle is a finished route, so the walk must not carry on from it. Carrying on
1264    /// produces routes that pass through the start token in the middle, which the executor cannot
1265    /// take: `A -> B -> A -> C -> A` visits A three times.
1266    #[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    /// Every route between two tokens under the given hop bounds and exclusions.
1294    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    /// A route may not pass through a token the request excludes.
1315    #[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    /// A pair whose every pool is excluded no longer connects its two tokens.
1327    #[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    /// A pair keeps the pools that are left, so the expansion drops the excluded one alone.
1339    #[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    /// S and T, with parallel pools stacked on every leg of the long way round.
1356    ///
1357    /// ```text
1358    ///   S ==[sx1,sx2]== X ==[xy1,xy2,xy3]== Y ==[yt1,yt2]== T
1359    ///   S --[sy1]------------------------- Y
1360    ///   S --[st1]--------------------------------------------- T
1361    /// ```
1362    ///
1363    /// The other fixtures carry parallel pools on at most one leg, so they never multiply. Here the
1364    /// three-hop route is 2 x 3 x 2, which is what the expansion has to reproduce from a single
1365    /// token sequence.
1366    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        // 1 hop: st1.
1395        assert_eq!(routes(graph, &s, &t, 1, 1, None).len(), 1);
1396        // 2 hops: sy1 x {yt1, yt2}.
1397        assert_eq!(routes(graph, &s, &t, 2, 2, None).len(), 2);
1398        // 3 hops: {sx1, sx2} x {xy1, xy2, xy3} x {yt1, yt2}.
1399        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        // The twelve long routes must name twelve distinct pool triples, not one triple twelve
1405        // times -- a counting slip in the expansion would still produce the right total.
1406        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        // Barring Y as an intermediate leaves only the direct pool.
1425        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        // Build a graph with 1-hop, 2-hop, and 3-hop paths to E:
1441        //   A --[ae]--> E                          (1-hop)
1442        //   A --[ab]--> B --[be]--> E              (2-hop)
1443        //   A --[ac]--> C --[cd]--> D --[de]--> E  (3-hop)
1444        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        // BFS guarantees paths are ordered by hop count
1460        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        // Diamond: A->B->D, A->C->D. Only C in allowlist → only A->C->D survives.
1469        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        // B must not appear; C must appear
1480        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        // Allowlist contains neither token_in nor token_out, but a 1-hop route should still work.
1487        let (a, b, _, _) = addrs();
1488        let m = linear_graph();
1489        let g = m.graph();
1490        // Empty allowlist, so no token may serve as an intermediate. A 1-hop route reaches the
1491        // destination directly and has none.
1492        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        // None allowlist → both paths in diamond graph returned
1500        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}