Skip to main content

fynd_core/algorithm/
bellman_ford.rs

1//! Bellman-Ford algorithm with SPFA optimization for simulation-driven routing.
2//!
3//! Runs actual pool simulations (`get_amount_out()`) during edge relaxation to find
4//! optimal A-to-B routes that account for slippage, fees, and pool mechanics at the
5//! given trade size.
6//!
7//! Key features:
8//! - **Gas-aware relaxation**: When token prices and gas price are available, relaxation compares
9//!   net amounts (gross output minus cumulative gas cost in token terms) instead of gross output
10//!   alone. Falls back to gross comparison when data is unavailable.
11//! - **Subgraph extraction**: BFS prunes the graph to nodes reachable within `max_hops`
12//! - **SPFA (Shortest Path Faster Algorithm) queuing**: Only re-relaxes edges from nodes whose
13//!   amount improved
14//! - **Forbid revisits**: Skips edges that would revisit a token or pool already in the path
15//!
16//! # Known limitation: SPFA order-dependence
17//!
18//! SPFA reads and writes the same `amount[]` array within a round (unlike
19//! textbook Bellman-Ford which snapshots between rounds). Processing node B
20//! before node C can update intermediate amounts that C then builds on,
21//! producing different routes depending on iteration order. Active nodes are
22//! sorted by `NodeIndex` for determinism, but the chosen ordering is not
23//! guaranteed to find the globally optimal route. A proper fix would be to
24//! snapshot amounts between rounds or use a priority-based processing order.
25
26use std::{
27    collections::{HashMap, HashSet, VecDeque},
28    time::{Duration, Instant},
29};
30
31use num_bigint::{BigInt, BigUint};
32use num_traits::{ToPrimitive, Zero};
33use petgraph::{graph::NodeIndex, prelude::EdgeRef};
34use tracing::{debug, instrument, trace, warn};
35use tycho_simulation::{
36    tycho_common::models::Address,
37    tycho_core::{models::token::Token, simulation::protocol_sim::Price},
38};
39
40use super::{Algorithm, AlgorithmConfig, AlgorithmError, NoPathReason};
41use crate::{
42    derived::{
43        computation::ComputationRequirements,
44        types::{SpotPrices, TokenGasPrices},
45        SharedDerivedDataRef,
46    },
47    feed::market_data::{MarketData, StateLabel},
48    graph::{petgraph::StableDiGraph, PetgraphStableDiGraphManager},
49    types::{ComponentId, Order, Route, RouteResult, Swap},
50};
51
52/// BFS subgraph: adjacency list, token node set, and component ID set.
53type Subgraph =
54    (HashMap<NodeIndex, Vec<(NodeIndex, ComponentId)>>, HashSet<NodeIndex>, HashSet<ComponentId>);
55
56/// Bellman-Ford algorithm with SPFA optimisation for simulation-driven DEX routing.
57///
58/// Finds optimal A→B routes by running actual pool simulations during edge relaxation,
59/// accounting for slippage, fees, and pool mechanics at the requested trade size.
60/// Gas costs are subtracted when price data is available.
61pub struct BellmanFordAlgorithm {
62    max_hops: usize,
63    timeout: Duration,
64    gas_aware: bool,
65    connector_tokens: Option<HashSet<Address>>,
66}
67
68impl BellmanFordAlgorithm {
69    pub(crate) fn with_config(config: AlgorithmConfig) -> Result<Self, AlgorithmError> {
70        Ok(Self {
71            max_hops: config.max_hops(),
72            timeout: config.timeout(),
73            gas_aware: config.gas_aware(),
74            connector_tokens: config.connector_tokens().cloned(),
75        })
76    }
77
78    /// Computes gas-adjusted net amount: gross_amount - gas_cost_in_token.
79    ///
80    /// If `token_price` is None (no conversion rate available), returns the gross amount
81    /// unchanged (falls back to gross comparison for this node).
82    fn gas_adjusted_amount(
83        gross: &BigUint,
84        cumul_gas: &BigUint,
85        gas_price_wei: &BigUint,
86        token_price: Option<&Price>,
87    ) -> BigInt {
88        match token_price {
89            Some(price) if !price.denominator.is_zero() => {
90                let gas_cost = cumul_gas * gas_price_wei * &price.numerator / &price.denominator;
91                BigInt::from(gross.clone()) - BigInt::from(gas_cost)
92            }
93            _ => BigInt::from(gross.clone()),
94        }
95    }
96
97    /// Computes the cumulative spot price product when extending a path by one edge.
98    ///
99    /// Returns `parent_spot * spot_price(component, token_u, token_v)`.
100    /// Returns 0.0 if the spot price is unavailable (disables the fallback for this path).
101    fn compute_edge_spot_product(
102        parent_spot: f64,
103        component_id: &ComponentId,
104        u_addr: Option<&Address>,
105        v_addr: Option<&Address>,
106        spot_prices: Option<&SpotPrices>,
107    ) -> f64 {
108        if parent_spot == 0.0 {
109            return 0.0;
110        }
111        let (Some(u), Some(v), Some(prices)) = (u_addr, v_addr, spot_prices) else {
112            return 0.0;
113        };
114        let key = (component_id.clone(), u.clone(), v.clone());
115        match prices.get(&key) {
116            Some(&spot) if spot > 0.0 => parent_spot * spot,
117            _ => 0.0,
118        }
119    }
120
121    /// Resolves the gas-to-token conversion rate for gas cost calculation.
122    ///
123    /// 1. Primary: use `token_prices[v_addr]` from derived data (direct lookup).
124    /// 2. Fallback: if `token_prices[token_in]` exists and `spot_product > 0`, estimate the rate as
125    ///    `token_prices[token_in] * spot_product` (converted to a Price).
126    /// 3. Last resort: returns None (gas adjustment skipped for this comparison).
127    fn resolve_token_price(
128        v_addr: Option<&Address>,
129        token_prices: Option<&TokenGasPrices>,
130        spot_product: f64,
131        token_in_addr: Option<&Address>,
132    ) -> Option<Price> {
133        let prices = token_prices?;
134        let addr = v_addr?;
135
136        // Primary: direct lookup
137        if let Some(price) = prices.get(addr) {
138            return Some(price.clone());
139        }
140
141        // Fallback: token_in price * cumulative spot product
142        if spot_product > 0.0 {
143            if let Some(in_price) = token_in_addr.and_then(|a| prices.get(a)) {
144                let in_rate_f64 = in_price.numerator.to_f64()? / in_price.denominator.to_f64()?;
145                let estimated_rate = in_rate_f64 * spot_product;
146                let denom = BigUint::from(10u64).pow(18);
147                let numer_f64 = estimated_rate * 1e18;
148                if numer_f64.is_finite() && numer_f64 > 0.0 {
149                    return Some(Price {
150                        numerator: BigUint::from(numer_f64 as u128),
151                        denominator: denom,
152                    });
153                }
154            }
155        }
156
157        None
158    }
159
160    /// Checks whether the target node or pool conflicts with the existing path to `from`.
161    /// Walks the predecessor chain once, checking both conditions simultaneously.
162    fn path_has_conflict(
163        from: NodeIndex,
164        target_node: NodeIndex,
165        target_pool: &ComponentId,
166        predecessor: &[Option<(NodeIndex, ComponentId)>],
167    ) -> bool {
168        let mut current = from;
169        loop {
170            if current == target_node {
171                return true;
172            }
173            match &predecessor[current.index()] {
174                Some((prev, cid)) => {
175                    if cid == target_pool {
176                        return true;
177                    }
178                    current = *prev;
179                }
180                None => return false,
181            }
182        }
183    }
184
185    /// Reconstructs the path from token_out back to token_in by walking the predecessor
186    /// array.
187    fn reconstruct_path(
188        token_out: NodeIndex,
189        token_in: NodeIndex,
190        predecessor: &[Option<(NodeIndex, ComponentId)>],
191    ) -> Result<Vec<(NodeIndex, NodeIndex, ComponentId)>, AlgorithmError> {
192        let mut path = Vec::new();
193        let mut current = token_out;
194        let mut visited = HashSet::new();
195
196        while current != token_in {
197            if !visited.insert(current) {
198                return Err(AlgorithmError::Other("cycle in predecessor chain".to_string()));
199            }
200
201            let idx = current.index();
202            match &predecessor
203                .get(idx)
204                .and_then(|p| p.as_ref())
205            {
206                Some((prev_node, component_id)) => {
207                    path.push((*prev_node, current, component_id.clone()));
208                    current = *prev_node;
209                }
210                None => {
211                    return Err(AlgorithmError::Other(format!(
212                        "broken predecessor chain at node {idx}"
213                    )));
214                }
215            }
216        }
217
218        path.reverse();
219        Ok(path)
220    }
221
222    /// Extracts the subgraph reachable from `token_in_node` within `max_hops` via BFS.
223    ///
224    /// Returns `(adjacency_list, token_nodes, component_ids)` or `NoPath` if the
225    /// subgraph is empty (no outgoing edges from the source).
226    fn get_subgraph(
227        graph: &StableDiGraph<()>,
228        token_in_node: NodeIndex,
229        max_hops: usize,
230        order: &Order,
231    ) -> Result<Subgraph, AlgorithmError> {
232        let mut adj: HashMap<NodeIndex, Vec<(NodeIndex, ComponentId)>> = HashMap::new();
233        let mut token_nodes: HashSet<NodeIndex> = HashSet::new();
234        let mut component_ids: HashSet<ComponentId> = HashSet::new();
235        let mut visited = HashSet::new();
236        let mut queue = VecDeque::new();
237
238        visited.insert(token_in_node);
239        token_nodes.insert(token_in_node);
240        queue.push_back((token_in_node, 0usize));
241
242        while let Some((node, depth)) = queue.pop_front() {
243            if depth >= max_hops {
244                continue;
245            }
246            for edge in graph.edges(node) {
247                let target = edge.target();
248                let cid = edge.weight().component_id.clone();
249
250                adj.entry(node)
251                    .or_default()
252                    .push((target, cid.clone()));
253                component_ids.insert(cid);
254                token_nodes.insert(target);
255
256                if visited.insert(target) {
257                    queue.push_back((target, depth + 1));
258                }
259            }
260        }
261
262        if adj.is_empty() {
263            return Err(AlgorithmError::NoPath {
264                from: order.token_in().clone(),
265                to: order.token_out().clone(),
266                reason: NoPathReason::NoGraphPath,
267            });
268        }
269
270        Ok((adj, token_nodes, component_ids))
271    }
272
273    /// Computes net_amount_out by subtracting gas costs from the output amount.
274    ///
275    /// Uses the same resolution strategy as relaxation: direct token price lookup
276    /// first, then cumulative spot price product fallback for tokens not in the price
277    /// table.
278    #[allow(clippy::too_many_arguments)]
279    fn compute_net_amount_out(
280        amount_out: &BigUint,
281        route: &Route,
282        gas_price: &BigUint,
283        token_prices: Option<&TokenGasPrices>,
284        spot_product: &[f64],
285        node_address: &HashMap<NodeIndex, Address>,
286        token_in_node: NodeIndex,
287    ) -> BigInt {
288        let Some(last_swap) = route.swaps().last() else {
289            return BigInt::from(amount_out.clone());
290        };
291
292        let total_gas = route.total_gas();
293
294        if gas_price.is_zero() {
295            warn!("missing gas price, returning gross amount_out");
296            return BigInt::from(amount_out.clone());
297        }
298
299        let gas_cost_wei = &total_gas * gas_price;
300
301        // Find the output token's node to get its spot_product for the fallback
302        let out_addr = last_swap.token_out();
303        let out_node_spot = node_address
304            .iter()
305            .find(|(_, addr)| *addr == out_addr)
306            .and_then(|(node, _)| spot_product.get(node.index()).copied())
307            .unwrap_or(0.0);
308
309        let output_price = Self::resolve_token_price(
310            Some(out_addr),
311            token_prices,
312            out_node_spot,
313            node_address.get(&token_in_node),
314        );
315
316        match output_price {
317            Some(price) if !price.denominator.is_zero() => {
318                let gas_cost = &gas_cost_wei * &price.numerator / &price.denominator;
319                BigInt::from(amount_out.clone()) - BigInt::from(gas_cost)
320            }
321            _ => {
322                warn!("no gas price for output token, returning gross amount_out");
323                BigInt::from(amount_out.clone())
324            }
325        }
326    }
327}
328
329impl Algorithm for BellmanFordAlgorithm {
330    type GraphType = StableDiGraph<()>;
331    type GraphManager = PetgraphStableDiGraphManager<()>;
332
333    fn name(&self) -> &str {
334        "bellman_ford"
335    }
336
337    #[instrument(level = "debug", skip_all, fields(order_id = %order.id()))]
338    async fn find_best_route(
339        &self,
340        graph: &Self::GraphType,
341        market: MarketData,
342        label: Option<StateLabel>,
343        derived: Option<SharedDerivedDataRef>,
344        order: &Order,
345    ) -> Result<RouteResult, AlgorithmError> {
346        let start = Instant::now();
347
348        if !order.is_sell() {
349            return Err(AlgorithmError::ExactOutNotSupported);
350        }
351
352        let (token_prices, spot_prices) = if let Some(ref derived) = derived {
353            let guard = derived.read().await;
354            (guard.token_prices().cloned(), guard.spot_prices().cloned())
355        } else {
356            (None, None)
357        };
358
359        let token_in_node = graph
360            .node_indices()
361            .find(|&n| &graph[n] == order.token_in())
362            .ok_or(AlgorithmError::NoPath {
363                from: order.token_in().clone(),
364                to: order.token_out().clone(),
365                reason: NoPathReason::SourceTokenNotInGraph,
366            })?;
367        let token_out_node = graph
368            .node_indices()
369            .find(|&n| &graph[n] == order.token_out())
370            .ok_or(AlgorithmError::NoPath {
371                from: order.token_in().clone(),
372                to: order.token_out().clone(),
373                reason: NoPathReason::DestinationTokenNotInGraph,
374            })?;
375
376        // BFS from token_in up to max_hops, building adjacency list and component set.
377        let (adj, token_nodes, component_ids) =
378            Self::get_subgraph(graph, token_in_node, self.max_hops, order)?;
379
380        // Acquire read lock only for market data extraction, then release.
381        let (token_map, market_subset) = {
382            let market = match label.as_ref() {
383                Some(l) => market
384                    .read_labeled(l)
385                    .await
386                    .map_err(|e| AlgorithmError::Other(e.to_string()))?,
387                None => market.read().await,
388            };
389
390            let token_map: HashMap<NodeIndex, Token> = token_nodes
391                .iter()
392                .filter_map(|&node| {
393                    market
394                        .get_token(&graph[node])
395                        .cloned()
396                        .map(|t| (node, t))
397                })
398                .collect();
399
400            let market_subset = market.extract_subset_with_overlay(&component_ids);
401
402            (token_map, market_subset)
403        };
404
405        debug!(
406            edges = adj
407                .values()
408                .map(Vec::len)
409                .sum::<usize>(),
410            tokens = token_map.len(),
411            "subgraph extracted"
412        );
413
414        // SPFA relaxation with forbid-revisits.
415        // amount[node] = best gross output amount reachable at node.
416        // edge_gas[node] = gas estimate for the edge that last improved amount[node].
417        // cumul_gas[node] = total gas units along the best path to this node.
418        let max_idx = graph
419            .node_indices()
420            .map(|n| n.index())
421            .max()
422            .unwrap_or(0) +
423            1;
424
425        let mut amount: Vec<BigUint> = vec![BigUint::ZERO; max_idx];
426        let mut predecessor: Vec<Option<(NodeIndex, ComponentId)>> = vec![None; max_idx];
427        let mut edge_gas: Vec<BigUint> = vec![BigUint::ZERO; max_idx];
428        let mut cumul_gas: Vec<BigUint> = vec![BigUint::ZERO; max_idx];
429
430        amount[token_in_node.index()] = order.amount().clone();
431
432        // Gas-aware relaxation: pre-compute gas price and build address map for lookups.
433        let gas_price_wei = market_subset
434            .gas_price()
435            .map(|gp| gp.effective_gas_price().clone());
436
437        // Build node->address map for token price lookups during relaxation.
438        let node_address: HashMap<NodeIndex, Address> = token_map
439            .iter()
440            .map(|(&node, token)| (node, token.address.clone()))
441            .collect();
442
443        // Track cumulative spot price product from token_in for fallback gas estimation.
444        // spot_product[v] = product of spot prices along the path from token_in to v.
445        let mut spot_product: Vec<f64> = vec![0.0; max_idx];
446        spot_product[token_in_node.index()] = 1.0;
447
448        let gas_aware = self.gas_aware && gas_price_wei.is_some() && token_prices.is_some();
449        if !gas_aware && self.gas_aware {
450            debug!("gas-aware comparison disabled (missing gas_price or token_prices)");
451        } else if !self.gas_aware {
452            debug!("gas-aware comparison disabled by config");
453        }
454
455        let mut active_nodes: Vec<NodeIndex> = vec![token_in_node];
456
457        for round in 0..self.max_hops {
458            if start.elapsed() >= self.timeout {
459                debug!(round, "timeout during relaxation");
460                break;
461            }
462            if active_nodes.is_empty() {
463                debug!(round, "no active nodes, stopping early");
464                break;
465            }
466
467            let mut next_active: HashSet<NodeIndex> = HashSet::new();
468
469            for &u in &active_nodes {
470                let u_idx = u.index();
471                if amount[u_idx].is_zero() {
472                    continue;
473                }
474
475                let Some(token_u) = token_map.get(&u) else { continue };
476                let Some(edges) = adj.get(&u) else { continue };
477
478                for (v, component_id) in edges {
479                    let v_idx = v.index();
480
481                    // Single predecessor walk: skip if target token or pool already in path
482                    if Self::path_has_conflict(u, *v, component_id, &predecessor) {
483                        continue;
484                    }
485
486                    // Skip disallowed connector tokens. Endpoints (token_in / token_out) are
487                    // always permitted regardless of the allowlist.
488                    if let (Some(tokens), Some(v_addr)) =
489                        (&self.connector_tokens, node_address.get(v))
490                    {
491                        if v_addr != order.token_in() &&
492                            v_addr != order.token_out() &&
493                            !tokens.contains(v_addr)
494                        {
495                            continue;
496                        }
497                    }
498
499                    let Some(token_v) = token_map.get(v) else { continue };
500                    let Some(sim) = market_subset.get_simulation_state(component_id) else {
501                        continue;
502                    };
503
504                    let result = match sim.get_amount_out(amount[u_idx].clone(), token_u, token_v) {
505                        Ok(r) => r,
506                        Err(e) => {
507                            trace!(
508                                component_id,
509                                error = %e,
510                                "simulation failed, skipping edge"
511                            );
512                            continue;
513                        }
514                    };
515
516                    let candidate_cumul_gas = &cumul_gas[u_idx] + &result.gas;
517
518                    // Compute spot price product for the candidate path (used for
519                    // gas-aware comparison and for final net amount calculation).
520                    let candidate_spot = Self::compute_edge_spot_product(
521                        spot_product[u_idx],
522                        component_id,
523                        node_address.get(&u),
524                        node_address.get(v),
525                        spot_prices.as_ref(),
526                    );
527
528                    // Gas-aware comparison: compare net amounts (gross - gas cost in token terms)
529                    let is_better = if gas_aware {
530                        let v_price = Self::resolve_token_price(
531                            node_address.get(v),
532                            token_prices.as_ref(),
533                            candidate_spot,
534                            node_address.get(&token_in_node),
535                        );
536
537                        let net_candidate = Self::gas_adjusted_amount(
538                            &result.amount,
539                            &candidate_cumul_gas,
540                            gas_price_wei.as_ref().unwrap(),
541                            v_price.as_ref(),
542                        );
543                        let net_existing = Self::gas_adjusted_amount(
544                            &amount[v_idx],
545                            &cumul_gas[v_idx],
546                            gas_price_wei.as_ref().unwrap(),
547                            v_price.as_ref(),
548                        );
549                        net_candidate > net_existing
550                    } else {
551                        result.amount > amount[v_idx]
552                    };
553
554                    if is_better {
555                        spot_product[v_idx] = candidate_spot;
556                        amount[v_idx] = result.amount;
557                        predecessor[v_idx] = Some((u, component_id.clone()));
558                        edge_gas[v_idx] = result.gas;
559                        cumul_gas[v_idx] = candidate_cumul_gas;
560                        next_active.insert(*v);
561                    }
562                }
563            }
564
565            active_nodes = next_active.into_iter().collect();
566            // Deterministic order: HashSet iteration is random per process.
567            // This pins SPFA to a fixed propagation order. The chosen order
568            // may not yield the optimal route (see module docs), but the
569            // previous random order was statistically no better.
570            active_nodes.sort_unstable();
571        }
572
573        // Check if destination was reached
574        let out_idx = token_out_node.index();
575        if amount[out_idx].is_zero() {
576            return Err(AlgorithmError::NoPath {
577                from: order.token_in().clone(),
578                to: order.token_out().clone(),
579                reason: NoPathReason::NoGraphPath,
580            });
581        }
582
583        // Reconstruct path and build route directly from stored distances/gas
584        // (no re-simulation needed since forbid-revisits guarantees relaxation
585        // amounts match sequential execution).
586        let path_edges = Self::reconstruct_path(token_out_node, token_in_node, &predecessor)?;
587
588        let mut swaps = Vec::with_capacity(path_edges.len());
589        let mut tokens: HashMap<Address, Token> = HashMap::new();
590        for (from_node, to_node, component_id) in &path_edges {
591            let token_in = token_map
592                .get(from_node)
593                .ok_or_else(|| AlgorithmError::DataNotFound {
594                    kind: "token",
595                    id: Some(format!("{:?}", from_node)),
596                })?;
597            let token_out = token_map
598                .get(to_node)
599                .ok_or_else(|| AlgorithmError::DataNotFound {
600                    kind: "token",
601                    id: Some(format!("{:?}", to_node)),
602                })?;
603            let component = market_subset
604                .get_component(component_id)
605                .ok_or_else(|| AlgorithmError::DataNotFound {
606                    kind: "component",
607                    id: Some(component_id.clone()),
608                })?;
609            let sim_state = market_subset
610                .get_simulation_state(component_id)
611                .ok_or_else(|| AlgorithmError::DataNotFound {
612                    kind: "simulation state",
613                    id: Some(component_id.clone()),
614                })?;
615
616            swaps.push(Swap::new(
617                component_id.clone(),
618                component.protocol_system.clone(),
619                token_in.address.clone(),
620                token_out.address.clone(),
621                amount[from_node.index()].clone(),
622                amount[to_node.index()].clone(),
623                edge_gas[to_node.index()].clone(),
624                component.clone(),
625                sim_state.clone_box(),
626            ));
627            tokens
628                .entry(token_in.address.clone())
629                .or_insert_with(|| token_in.clone());
630            tokens
631                .entry(token_out.address.clone())
632                .or_insert_with(|| token_out.clone());
633        }
634
635        let route = Route::new(swaps, tokens);
636        let final_amount_out = amount[out_idx].clone();
637
638        let gas_price = gas_price_wei.unwrap_or_default();
639
640        let net_amount_out = Self::compute_net_amount_out(
641            &final_amount_out,
642            &route,
643            &gas_price,
644            token_prices.as_ref(),
645            &spot_product,
646            &node_address,
647            token_in_node,
648        );
649
650        let result = RouteResult::new(route, net_amount_out, gas_price);
651
652        let solve_time_ms = start.elapsed().as_millis() as u64;
653        debug!(
654            solve_time_ms,
655            hops = result.route().swaps().len(),
656            amount_in = %order.amount(),
657            amount_out = %final_amount_out,
658            net_amount_out = %result.net_amount_out(),
659            "bellman_ford route found"
660        );
661
662        Ok(result)
663    }
664
665    fn computation_requirements(&self) -> ComputationRequirements {
666        // Static requirements for independent computations; cannot conflict.
667        // The trait returns ComputationRequirements (not Result), so expect is
668        // the appropriate pattern for this infallible case.
669        ComputationRequirements::none()
670            .allow_stale("token_prices")
671            .expect("token_prices requirement conflicts (bug)")
672            .allow_stale("spot_prices")
673            .expect("spot_prices requirement conflicts (bug)")
674    }
675
676    fn timeout(&self) -> Duration {
677        self.timeout
678    }
679}
680
681#[cfg(test)]
682mod tests {
683    use std::sync::Arc;
684
685    use num_bigint::BigInt;
686    use tokio::sync::RwLock;
687    use tycho_simulation::{
688        tycho_common::{models::Address, simulation::protocol_sim::ProtocolSim},
689        tycho_ethereum::gas::{BlockGasPrice, GasPrice},
690    };
691
692    use super::*;
693    use crate::{
694        algorithm::test_utils::{component, order, token, MockProtocolSim},
695        derived::{types::TokenGasPrices, DerivedData},
696        feed::market_data::{MarketData, MarketState},
697        graph::GraphManager,
698        types::quote::OrderSide,
699    };
700
701    // ==================== Test Utilities ====================
702
703    /// Sets up market and graph with `()` edge weights for BellmanFord tests.
704    fn setup_market_bf(
705        pools: Vec<(&str, &Token, &Token, MockProtocolSim)>,
706    ) -> (MarketData, PetgraphStableDiGraphManager<()>) {
707        let mut market = MarketState::new();
708
709        market.update_gas_price(BlockGasPrice {
710            block_number: 1,
711            block_hash: Default::default(),
712            block_timestamp: 0,
713            pricing: GasPrice::Legacy { gas_price: BigUint::from(100u64) },
714        });
715        market.update_last_updated(crate::types::BlockInfo::new(1, "0x00".into(), 0));
716
717        for (pool_id, token_in, token_out, state) in pools {
718            let tokens = vec![token_in.clone(), token_out.clone()];
719            let comp = component(pool_id, &tokens);
720            market.upsert_components(std::iter::once(comp));
721            market.update_states([(pool_id.to_string(), Box::new(state) as Box<dyn ProtocolSim>)]);
722            market.upsert_tokens(tokens);
723        }
724
725        let mut graph_manager = PetgraphStableDiGraphManager::default();
726        graph_manager.initialize_graph(&market.component_topology());
727
728        (MarketData::new(Arc::new(RwLock::new(market))), graph_manager)
729    }
730
731    fn setup_derived_with_token_prices(
732        token_addresses: &[Address],
733    ) -> crate::derived::SharedDerivedDataRef {
734        use tycho_simulation::tycho_core::simulation::protocol_sim::Price;
735
736        let mut token_prices: TokenGasPrices = HashMap::new();
737        for address in token_addresses {
738            token_prices.insert(
739                address.clone(),
740                Price { numerator: BigUint::from(1u64), denominator: BigUint::from(1u64) },
741            );
742        }
743
744        let mut derived_data = DerivedData::new();
745        derived_data.set_token_prices(token_prices, vec![], 1, true);
746        Arc::new(RwLock::new(derived_data))
747    }
748
749    fn bf_algorithm(max_hops: usize, timeout_ms: u64) -> BellmanFordAlgorithm {
750        BellmanFordAlgorithm::with_config(
751            AlgorithmConfig::new(1, max_hops, Duration::from_millis(timeout_ms), None).unwrap(),
752        )
753        .unwrap()
754    }
755
756    // ==================== Unit Tests ====================
757
758    #[tokio::test]
759    async fn test_linear_path_found() {
760        let token_a = token(0x01, "A");
761        let token_b = token(0x02, "B");
762        let token_c = token(0x03, "C");
763        let token_d = token(0x04, "D");
764
765        let (market, manager) = setup_market_bf(vec![
766            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
767            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
768            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(4.0)),
769        ]);
770
771        let algo = bf_algorithm(4, 1000);
772        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
773
774        let result = algo
775            .find_best_route(manager.graph(), market, None, None, &ord)
776            .await
777            .unwrap();
778
779        assert_eq!(result.route().swaps().len(), 3);
780        // A->B: 100*2=200, B->C: 200*3=600, C->D: 600*4=2400
781        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(200u64));
782        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(600u64));
783        assert_eq!(result.route().swaps()[2].amount_out(), &BigUint::from(2400u64));
784    }
785
786    #[tokio::test]
787    async fn test_picks_better_of_two_paths() {
788        // Diamond graph: A->B->D (2*3=6x) vs A->C->D (4*1=4x)
789        let token_a = token(0x01, "A");
790        let token_b = token(0x02, "B");
791        let token_c = token(0x03, "C");
792        let token_d = token(0x04, "D");
793
794        let (market, manager) = setup_market_bf(vec![
795            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
796            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(3.0)),
797            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(4.0)),
798            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(1.0)),
799        ]);
800
801        let algo = bf_algorithm(3, 1000);
802        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
803
804        let result = algo
805            .find_best_route(manager.graph(), market, None, None, &ord)
806            .await
807            .unwrap();
808
809        // A->B->D: 100*2*3=600 is better than A->C->D: 100*4*1=400
810        assert_eq!(result.route().swaps().len(), 2);
811        assert_eq!(result.route().swaps()[0].component_id(), "pool_ab");
812        assert_eq!(result.route().swaps()[1].component_id(), "pool_bd");
813        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(600u64));
814    }
815
816    #[tokio::test]
817    async fn test_parallel_pools() {
818        // Two pools between A and B with different multipliers
819        let token_a = token(0x01, "A");
820        let token_b = token(0x02, "B");
821
822        let (market, manager) = setup_market_bf(vec![
823            ("pool1", &token_a, &token_b, MockProtocolSim::new(2.0)),
824            ("pool2", &token_a, &token_b, MockProtocolSim::new(5.0)),
825        ]);
826
827        let algo = bf_algorithm(2, 1000);
828        let ord = order(&token_a, &token_b, 100, OrderSide::Sell);
829
830        let result = algo
831            .find_best_route(manager.graph(), market, None, None, &ord)
832            .await
833            .unwrap();
834
835        assert_eq!(result.route().swaps().len(), 1);
836        assert_eq!(result.route().swaps()[0].component_id(), "pool2");
837        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(500u64));
838    }
839
840    #[tokio::test]
841    async fn test_no_path_returns_error() {
842        let token_a = token(0x01, "A");
843        let token_b = token(0x02, "B");
844        let token_c = token(0x03, "C");
845
846        // A-B connected, C disconnected
847        let (market, manager) =
848            setup_market_bf(vec![("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0))]);
849
850        // Add token_c to market without connecting it
851        {
852            let mut m = market.write().await;
853            m.upsert_tokens(vec![token_c.clone()]);
854        }
855
856        let algo = bf_algorithm(3, 1000);
857        let ord = order(&token_a, &token_c, 100, OrderSide::Sell);
858
859        let result = algo
860            .find_best_route(manager.graph(), market, None, None, &ord)
861            .await;
862        assert!(matches!(result, Err(AlgorithmError::NoPath { .. })));
863    }
864
865    #[tokio::test]
866    async fn test_source_not_in_graph() {
867        let token_a = token(0x01, "A");
868        let token_b = token(0x02, "B");
869        let token_x = token(0x99, "X");
870
871        let (market, manager) =
872            setup_market_bf(vec![("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0))]);
873
874        let algo = bf_algorithm(3, 1000);
875        let ord = order(&token_x, &token_b, 100, OrderSide::Sell);
876
877        let result = algo
878            .find_best_route(manager.graph(), market, None, None, &ord)
879            .await;
880        assert!(matches!(
881            result,
882            Err(AlgorithmError::NoPath { reason: NoPathReason::SourceTokenNotInGraph, .. })
883        ));
884    }
885
886    #[tokio::test]
887    async fn test_destination_not_in_graph() {
888        let token_a = token(0x01, "A");
889        let token_b = token(0x02, "B");
890        let token_x = token(0x99, "X");
891
892        let (market, manager) =
893            setup_market_bf(vec![("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0))]);
894
895        let algo = bf_algorithm(3, 1000);
896        let ord = order(&token_a, &token_x, 100, OrderSide::Sell);
897
898        let result = algo
899            .find_best_route(manager.graph(), market, None, None, &ord)
900            .await;
901        assert!(matches!(
902            result,
903            Err(AlgorithmError::NoPath { reason: NoPathReason::DestinationTokenNotInGraph, .. })
904        ));
905    }
906
907    #[tokio::test]
908    async fn test_respects_max_hops() {
909        // Path A->B->C->D exists but requires 3 hops; max_hops=2
910        let token_a = token(0x01, "A");
911        let token_b = token(0x02, "B");
912        let token_c = token(0x03, "C");
913        let token_d = token(0x04, "D");
914
915        let (market, manager) = setup_market_bf(vec![
916            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
917            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
918            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(4.0)),
919        ]);
920
921        let algo = bf_algorithm(2, 1000);
922        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
923
924        let result = algo
925            .find_best_route(manager.graph(), market, None, None, &ord)
926            .await;
927        assert!(
928            matches!(result, Err(AlgorithmError::NoPath { .. })),
929            "Should not find 3-hop path with max_hops=2"
930        );
931    }
932
933    #[tokio::test]
934    async fn test_source_token_revisit_blocked() {
935        // Forbid-revisits prevents paths like A->B->A->B->C. The algorithm
936        // should find the direct A->B->C path instead.
937        let token_a = token(0x01, "A");
938        let token_b = token(0x02, "B");
939        let token_c = token(0x03, "C");
940
941        let (market, manager) = setup_market_bf(vec![
942            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
943            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
944        ]);
945
946        let algo = bf_algorithm(4, 1000);
947        let ord = order(&token_a, &token_c, 100, OrderSide::Sell);
948
949        let result = algo
950            .find_best_route(manager.graph(), market, None, None, &ord)
951            .await
952            .unwrap();
953
954        // Should find exactly the 2-hop path A->B->C = 100*2*3 = 600
955        assert_eq!(result.route().swaps().len(), 2);
956        assert_eq!(result.route().swaps()[0].component_id(), "pool_ab");
957        assert_eq!(result.route().swaps()[1].component_id(), "pool_bc");
958        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(600u64));
959    }
960
961    #[tokio::test]
962    async fn test_hub_token_revisit_blocked() {
963        // Forbid-revisits blocks A->B->C->B->D (B visited twice).
964        // The algorithm should find the direct A->B->D = 400 instead.
965        let token_a = token(0x01, "A");
966        let token_c = token(0x02, "C");
967        let token_b = token(0x03, "B");
968        let token_d = token(0x04, "D");
969
970        let (market, manager) = setup_market_bf(vec![
971            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
972            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
973            ("pool_cb", &token_c, &token_b, MockProtocolSim::new(100.0)),
974            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(2.0)),
975        ]);
976
977        let algo = bf_algorithm(4, 1000);
978        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
979
980        let result = algo
981            .find_best_route(manager.graph(), market, None, None, &ord)
982            .await
983            .unwrap();
984
985        // Should find A->B->D = 100*2*2 = 400 (the direct 2-hop path)
986        // The 4-hop revisit path A->B->C->B->D is blocked
987        assert_eq!(result.route().swaps().len(), 2, "should use direct 2-hop path");
988        assert_eq!(result.route().swaps()[0].component_id(), "pool_ab");
989        assert_eq!(result.route().swaps()[1].component_id(), "pool_bd");
990        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(400u64));
991    }
992
993    #[tokio::test]
994    async fn test_route_amounts_are_sequential() {
995        // Verify that swap amount_in[i+1] == amount_out[i] in the built route
996        let token_a = token(0x01, "A");
997        let token_b = token(0x02, "B");
998        let token_c = token(0x03, "C");
999
1000        let (market, manager) = setup_market_bf(vec![
1001            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1002            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
1003        ]);
1004
1005        let algo = bf_algorithm(3, 1000);
1006        let ord = order(&token_a, &token_c, 100, OrderSide::Sell);
1007
1008        let result = algo
1009            .find_best_route(manager.graph(), market, None, None, &ord)
1010            .await
1011            .unwrap();
1012
1013        assert_eq!(result.route().swaps().len(), 2);
1014        // amount_in of second swap == amount_out of first swap
1015        assert_eq!(result.route().swaps()[1].amount_in(), result.route().swaps()[0].amount_out());
1016    }
1017
1018    #[tokio::test]
1019    async fn test_gas_deduction() {
1020        let token_a = token(0x01, "A");
1021        let token_b = token(0x02, "B");
1022
1023        let (market, manager) = setup_market_bf(vec![(
1024            "pool1",
1025            &token_a,
1026            &token_b,
1027            MockProtocolSim::new(2.0).with_gas(10),
1028        )]);
1029
1030        let algo = bf_algorithm(2, 1000);
1031        let ord = order(&token_a, &token_b, 1000, OrderSide::Sell);
1032
1033        let derived = setup_derived_with_token_prices(std::slice::from_ref(&token_b.address));
1034
1035        let result = algo
1036            .find_best_route(manager.graph(), market, None, Some(derived), &ord)
1037            .await
1038            .unwrap();
1039
1040        // Output: 1000 * 2 = 2000
1041        // Gas: 10 gas units * 100 gas_price = 1000 wei * 1/1 price = 1000
1042        // Net: 2000 - 1000 = 1000
1043        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(2000u64));
1044        assert_eq!(result.net_amount_out(), &BigInt::from(1000));
1045    }
1046
1047    #[tokio::test]
1048    async fn test_timeout_respected() {
1049        let token_a = token(0x01, "A");
1050        let token_b = token(0x02, "B");
1051        let token_c = token(0x03, "C");
1052
1053        let (market, manager) = setup_market_bf(vec![
1054            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1055            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
1056        ]);
1057
1058        // 0ms timeout
1059        let algo = bf_algorithm(3, 0);
1060        let ord = order(&token_a, &token_c, 100, OrderSide::Sell);
1061
1062        let result = algo
1063            .find_best_route(manager.graph(), market, None, None, &ord)
1064            .await;
1065
1066        // With 0ms timeout, we expect either:
1067        // - A partial result (if some layers completed before timeout check)
1068        // - Timeout error
1069        // - NoPath (if timeout prevented completing enough layers to reach dest)
1070        match result {
1071            Ok(r) => {
1072                assert!(!r.route().swaps().is_empty());
1073            }
1074            Err(AlgorithmError::Timeout { .. }) | Err(AlgorithmError::NoPath { .. }) => {
1075                // Both are acceptable for 0ms timeout
1076            }
1077            Err(e) => panic!("Unexpected error: {:?}", e),
1078        }
1079    }
1080
1081    // ==================== Integration-style Tests ====================
1082
1083    #[tokio::test]
1084    async fn test_with_fees() {
1085        let token_a = token(0x01, "A");
1086        let token_b = token(0x02, "B");
1087
1088        // Pool with 10% fee
1089        let (market, manager) = setup_market_bf(vec![(
1090            "pool1",
1091            &token_a,
1092            &token_b,
1093            MockProtocolSim::new(2.0).with_fee(0.1),
1094        )]);
1095
1096        let algo = bf_algorithm(2, 1000);
1097        let ord = order(&token_a, &token_b, 1000, OrderSide::Sell);
1098
1099        let result = algo
1100            .find_best_route(manager.graph(), market, None, None, &ord)
1101            .await
1102            .unwrap();
1103
1104        // 1000 * 2 * (1-0.1) = 1800
1105        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(1800u64));
1106    }
1107
1108    #[tokio::test]
1109    async fn test_large_trade_slippage() {
1110        let token_a = token(0x01, "A");
1111        let token_b = token(0x02, "B");
1112
1113        // Pool with limited liquidity (500 tokens)
1114        let (market, manager) = setup_market_bf(vec![(
1115            "pool1",
1116            &token_a,
1117            &token_b,
1118            MockProtocolSim::new(2.0).with_liquidity(500),
1119        )]);
1120
1121        let algo = bf_algorithm(2, 1000);
1122        let ord = order(&token_a, &token_b, 1000, OrderSide::Sell);
1123
1124        // Should fail due to insufficient liquidity
1125        let result = algo
1126            .find_best_route(manager.graph(), market, None, None, &ord)
1127            .await;
1128        assert!(
1129            matches!(result, Err(AlgorithmError::NoPath { .. })),
1130            "Should fail when trade exceeds pool liquidity"
1131        );
1132    }
1133
1134    #[tokio::test]
1135    async fn test_disconnected_tokens_return_no_path() {
1136        // A-B connected, D-E disconnected. Routing A->E should fail.
1137        let token_a = token(0x01, "A");
1138        let token_b = token(0x02, "B");
1139        let token_d = token(0x04, "D");
1140        let token_e = token(0x05, "E");
1141
1142        let (market, manager) = setup_market_bf(vec![
1143            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1144            ("pool_de", &token_d, &token_e, MockProtocolSim::new(4.0)),
1145        ]);
1146
1147        let algo = bf_algorithm(3, 1000);
1148        let ord = order(&token_a, &token_e, 100, OrderSide::Sell);
1149
1150        let result = algo
1151            .find_best_route(manager.graph(), market, None, None, &ord)
1152            .await;
1153        assert!(
1154            matches!(result, Err(AlgorithmError::NoPath { .. })),
1155            "should not find path to disconnected component"
1156        );
1157    }
1158
1159    #[tokio::test]
1160    async fn test_spfa_skips_failed_simulations() {
1161        // Pool that will fail simulation (liquidity=0 would cause error for any amount)
1162        let token_a = token(0x01, "A");
1163        let token_b = token(0x02, "B");
1164        let token_c = token(0x03, "C");
1165
1166        let (market, manager) = setup_market_bf(vec![
1167            // Direct path with failing pool
1168            ("pool_ab_bad", &token_a, &token_b, MockProtocolSim::new(2.0).with_liquidity(0)),
1169            // Alternative path that works
1170            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(2.0)),
1171            ("pool_cb", &token_c, &token_b, MockProtocolSim::new(3.0)),
1172        ]);
1173
1174        let algo = bf_algorithm(3, 1000);
1175        let ord = order(&token_a, &token_b, 100, OrderSide::Sell);
1176
1177        let result = algo
1178            .find_best_route(manager.graph(), market, None, None, &ord)
1179            .await;
1180
1181        // Should find A->C->B despite A->B failing
1182        // Note: MockProtocolSim with liquidity=0 will fail for amount > 0
1183        // The direct A->B edge should be skipped and the 2-hop path used
1184        match result {
1185            Ok(r) => {
1186                // Found alternative path
1187                assert!(!r.route().swaps().is_empty());
1188            }
1189            Err(AlgorithmError::NoPath { .. }) => {
1190                // Also acceptable if liquidity=0 blocks all paths through B
1191                // (since the failing pool might also block the reverse B->A edge)
1192            }
1193            Err(e) => panic!("Unexpected error: {:?}", e),
1194        }
1195    }
1196
1197    #[tokio::test]
1198    async fn test_resimulation_produces_correct_amounts() {
1199        // Verifies that re-simulation produces the same correct sequential amounts
1200        let token_a = token(0x01, "A");
1201        let token_b = token(0x02, "B");
1202        let token_c = token(0x03, "C");
1203
1204        let (market, manager) = setup_market_bf(vec![
1205            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1206            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
1207        ]);
1208
1209        let algo = bf_algorithm(3, 1000);
1210        let ord = order(&token_a, &token_c, 100, OrderSide::Sell);
1211
1212        let result = algo
1213            .find_best_route(manager.graph(), market, None, None, &ord)
1214            .await
1215            .unwrap();
1216
1217        // Verify the final amounts are from re-simulation, not relaxation
1218        // A->B: 100*2=200, B->C: 200*3=600
1219        assert_eq!(result.route().swaps()[0].amount_in(), &BigUint::from(100u64));
1220        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(200u64));
1221        assert_eq!(result.route().swaps()[1].amount_in(), &BigUint::from(200u64));
1222        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(600u64));
1223    }
1224
1225    // ==================== Trait getter tests ====================
1226
1227    #[test]
1228    fn algorithm_name() {
1229        let algo = bf_algorithm(4, 200);
1230        assert_eq!(algo.name(), "bellman_ford");
1231    }
1232
1233    #[test]
1234    fn algorithm_timeout() {
1235        let algo = bf_algorithm(4, 200);
1236        assert_eq!(algo.timeout(), Duration::from_millis(200));
1237    }
1238
1239    // ==================== Forbid-revisit helper tests ====================
1240
1241    #[tokio::test]
1242    async fn test_gas_aware_relaxation_picks_cheaper_path() {
1243        // Diamond graph: A -> B -> D vs A -> C -> D
1244        // Path via B: higher gross output (3x * 2x = 6x) but extreme gas (100M per hop)
1245        // Path via C: lower gross output (2x * 2x = 4x) but cheap gas (100 per hop)
1246        //
1247        // With gas_price=100, token_prices[D]=1:1 for WETH conversion:
1248        // Path B gas cost: (100M + 100M) * 100 * 1 = 20B
1249        // Path C gas cost: (100 + 100) * 100 * 1 = 20K
1250        //
1251        // For an input of 1B:
1252        // Path B: gross = 6B, net = 6B - 20B = -14B
1253        // Path C: gross = 4B, net = 4B - 20K ≈ 4B
1254        //
1255        // Without gas awareness: Path B wins (6B > 4B)
1256        // With gas awareness: Path C wins (4B net > -14B net)
1257        let token_a = token(0x01, "A");
1258        let token_b = token(0x02, "B");
1259        let token_c = token(0x03, "C");
1260        let token_d = token(0x04, "D");
1261
1262        let high_gas: u64 = 100_000_000;
1263        let low_gas: u64 = 100;
1264
1265        let (market, manager) = setup_market_bf(vec![
1266            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(3.0).with_gas(high_gas)),
1267            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(2.0).with_gas(high_gas)),
1268            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(2.0).with_gas(low_gas)),
1269            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(2.0).with_gas(low_gas)),
1270        ]);
1271
1272        let algo = bf_algorithm(3, 1000);
1273        let ord = order(&token_a, &token_d, 1_000_000_000, OrderSide::Sell);
1274
1275        // With gas-aware relaxation (derived data with token prices + gas price in market)
1276        let derived = setup_derived_with_token_prices(&[
1277            token_a.address.clone(),
1278            token_b.address.clone(),
1279            token_c.address.clone(),
1280            token_d.address.clone(),
1281        ]);
1282
1283        let result = algo
1284            .find_best_route(manager.graph(), market, None, Some(derived), &ord)
1285            .await
1286            .unwrap();
1287
1288        // Gas-aware relaxation should pick the cheaper path A -> C -> D
1289        assert_eq!(result.route().swaps().len(), 2);
1290        assert_eq!(result.route().swaps()[0].component_id(), "pool_ac");
1291        assert_eq!(result.route().swaps()[1].component_id(), "pool_cd");
1292    }
1293
1294    #[tokio::test]
1295    async fn test_gas_aware_falls_back_to_gross_without_derived() {
1296        // Same diamond graph as above, but without derived data.
1297        // Should fall back to gross comparison and pick Path B (higher gross).
1298        let token_a = token(0x01, "A");
1299        let token_b = token(0x02, "B");
1300        let token_c = token(0x03, "C");
1301        let token_d = token(0x04, "D");
1302
1303        let high_gas: u64 = 100_000_000;
1304        let low_gas: u64 = 100;
1305
1306        let (market, manager) = setup_market_bf(vec![
1307            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(3.0).with_gas(high_gas)),
1308            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(2.0).with_gas(high_gas)),
1309            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(2.0).with_gas(low_gas)),
1310            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(2.0).with_gas(low_gas)),
1311        ]);
1312
1313        let algo = bf_algorithm(3, 1000);
1314        let ord = order(&token_a, &token_d, 1_000_000_000, OrderSide::Sell);
1315
1316        // No derived data: should fall back to gross comparison
1317        let result = algo
1318            .find_best_route(manager.graph(), market, None, None, &ord)
1319            .await
1320            .unwrap();
1321
1322        // Without gas awareness, picks the higher-gross path A -> B -> D
1323        assert_eq!(result.route().swaps().len(), 2);
1324        assert_eq!(result.route().swaps()[0].component_id(), "pool_ab");
1325        assert_eq!(result.route().swaps()[1].component_id(), "pool_bd");
1326    }
1327
1328    // ==================== Connector token tests ====================
1329
1330    /// Build a BellmanFord algorithm whose config includes a specific connector token allowlist.
1331    fn bf_algorithm_with_connectors(
1332        max_hops: usize,
1333        timeout_ms: u64,
1334        connector_tokens: HashSet<Address>,
1335    ) -> BellmanFordAlgorithm {
1336        BellmanFordAlgorithm::with_config(
1337            AlgorithmConfig::new(1, max_hops, Duration::from_millis(timeout_ms), None)
1338                .unwrap()
1339                .with_connector_tokens(connector_tokens),
1340        )
1341        .unwrap()
1342    }
1343
1344    #[tokio::test]
1345    async fn test_connector_tokens_blocks_disallowed_intermediate() {
1346        //      A
1347        //    /   \
1348        //   B     C   ← only C is in the allowlist
1349        //    \   /
1350        //      D
1351        // A->B->D is pruned; only A->C->D survives.
1352        let token_a = token(0x01, "A");
1353        let token_b = token(0x02, "B");
1354        let token_c = token(0x03, "C");
1355        let token_d = token(0x04, "D");
1356
1357        let (market, manager) = setup_market_bf(vec![
1358            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1359            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(2.0)),
1360            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(3.0)),
1361            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(3.0)),
1362        ]);
1363
1364        let connectors: HashSet<Address> = [token_c.address.clone()].into();
1365        let algo = bf_algorithm_with_connectors(3, 1000, connectors);
1366        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
1367
1368        let result = algo
1369            .find_best_route(manager.graph(), market, None, None, &ord)
1370            .await
1371            .unwrap();
1372
1373        // Only A->C->D is reachable; B was pruned.
1374        assert_eq!(result.route().swaps().len(), 2);
1375        assert_eq!(result.route().swaps()[0].component_id(), "pool_ac");
1376        assert_eq!(result.route().swaps()[1].component_id(), "pool_cd");
1377    }
1378
1379    #[tokio::test]
1380    async fn test_connector_tokens_allows_endpoints_even_if_not_listed() {
1381        // token_in (A) and token_out (B) must be reachable even when connector list is empty.
1382        let token_a = token(0x01, "A");
1383        let token_b = token(0x02, "B");
1384
1385        let (market, manager) =
1386            setup_market_bf(vec![("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0))]);
1387
1388        // Empty allowlist — no intermediate tokens allowed, but direct hop A->B should work.
1389        let algo = bf_algorithm_with_connectors(1, 1000, HashSet::new());
1390        let ord = order(&token_a, &token_b, 100, OrderSide::Sell);
1391
1392        let result = algo
1393            .find_best_route(manager.graph(), market, None, None, &ord)
1394            .await
1395            .unwrap();
1396
1397        assert_eq!(result.route().swaps().len(), 1);
1398        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(200u64));
1399    }
1400
1401    #[tokio::test]
1402    async fn test_connector_tokens_none_is_unrestricted() {
1403        // No connector_tokens set: both A->B->D and A->C->D are evaluated.
1404        let token_a = token(0x01, "A");
1405        let token_b = token(0x02, "B");
1406        let token_c = token(0x03, "C");
1407        let token_d = token(0x04, "D");
1408
1409        let (market, manager) = setup_market_bf(vec![
1410            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1411            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(3.0)),
1412            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(1.0)),
1413            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(1.0)),
1414        ]);
1415
1416        let algo = bf_algorithm(3, 1000);
1417        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
1418
1419        let result = algo
1420            .find_best_route(manager.graph(), market, None, None, &ord)
1421            .await
1422            .unwrap();
1423
1424        // Best path is A->B->D = 100*2*3 = 600
1425        assert_eq!(result.route().swaps()[0].component_id(), "pool_ab");
1426        assert_eq!(result.route().swaps()[1].component_id(), "pool_bd");
1427        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(600u64));
1428    }
1429
1430    #[test]
1431    fn test_path_has_conflict_detects_node_and_pool() {
1432        // Path: 0 -[pool_a]-> 1 -[pool_b]-> 2
1433        let mut pred: Vec<Option<(NodeIndex, ComponentId)>> = vec![None; 4];
1434        pred[1] = Some((NodeIndex::new(0), "pool_a".into()));
1435        pred[2] = Some((NodeIndex::new(1), "pool_b".into()));
1436
1437        // Node conflicts: node 0 is in path, node 3 is not
1438        assert!(BellmanFordAlgorithm::path_has_conflict(
1439            NodeIndex::new(2),
1440            NodeIndex::new(0),
1441            &"any".into(),
1442            &pred
1443        ));
1444        assert!(!BellmanFordAlgorithm::path_has_conflict(
1445            NodeIndex::new(2),
1446            NodeIndex::new(3),
1447            &"any".into(),
1448            &pred
1449        ));
1450        // Self-check: node 2 is itself in the "path from 2"
1451        assert!(BellmanFordAlgorithm::path_has_conflict(
1452            NodeIndex::new(2),
1453            NodeIndex::new(2),
1454            &"any".into(),
1455            &pred
1456        ));
1457
1458        // Pool conflicts: pool_a and pool_b are used, pool_c is not
1459        assert!(BellmanFordAlgorithm::path_has_conflict(
1460            NodeIndex::new(2),
1461            NodeIndex::new(3),
1462            &"pool_a".into(),
1463            &pred
1464        ));
1465        assert!(BellmanFordAlgorithm::path_has_conflict(
1466            NodeIndex::new(2),
1467            NodeIndex::new(3),
1468            &"pool_b".into(),
1469            &pred
1470        ));
1471        assert!(!BellmanFordAlgorithm::path_has_conflict(
1472            NodeIndex::new(2),
1473            NodeIndex::new(3),
1474            &"pool_c".into(),
1475            &pred
1476        ));
1477    }
1478}