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        for (from_node, to_node, component_id) in &path_edges {
590            let token_in = token_map
591                .get(from_node)
592                .ok_or_else(|| AlgorithmError::DataNotFound {
593                    kind: "token",
594                    id: Some(format!("{:?}", from_node)),
595                })?;
596            let token_out = token_map
597                .get(to_node)
598                .ok_or_else(|| AlgorithmError::DataNotFound {
599                    kind: "token",
600                    id: Some(format!("{:?}", to_node)),
601                })?;
602            let component = market_subset
603                .get_component(component_id)
604                .ok_or_else(|| AlgorithmError::DataNotFound {
605                    kind: "component",
606                    id: Some(component_id.clone()),
607                })?;
608            let sim_state = market_subset
609                .get_simulation_state(component_id)
610                .ok_or_else(|| AlgorithmError::DataNotFound {
611                    kind: "simulation state",
612                    id: Some(component_id.clone()),
613                })?;
614
615            swaps.push(Swap::new(
616                component_id.clone(),
617                component.protocol_system.clone(),
618                token_in.address.clone(),
619                token_out.address.clone(),
620                amount[from_node.index()].clone(),
621                amount[to_node.index()].clone(),
622                edge_gas[to_node.index()].clone(),
623                component.clone(),
624                sim_state.clone_box(),
625            ));
626        }
627
628        let route = Route::new(swaps);
629        let final_amount_out = amount[out_idx].clone();
630
631        let gas_price = gas_price_wei.unwrap_or_default();
632
633        let net_amount_out = Self::compute_net_amount_out(
634            &final_amount_out,
635            &route,
636            &gas_price,
637            token_prices.as_ref(),
638            &spot_product,
639            &node_address,
640            token_in_node,
641        );
642
643        let result = RouteResult::new(route, net_amount_out, gas_price);
644
645        let solve_time_ms = start.elapsed().as_millis() as u64;
646        debug!(
647            solve_time_ms,
648            hops = result.route().swaps().len(),
649            amount_in = %order.amount(),
650            amount_out = %final_amount_out,
651            net_amount_out = %result.net_amount_out(),
652            "bellman_ford route found"
653        );
654
655        Ok(result)
656    }
657
658    fn computation_requirements(&self) -> ComputationRequirements {
659        // Static requirements for independent computations; cannot conflict.
660        // The trait returns ComputationRequirements (not Result), so expect is
661        // the appropriate pattern for this infallible case.
662        ComputationRequirements::none()
663            .allow_stale("token_prices")
664            .expect("token_prices requirement conflicts (bug)")
665            .allow_stale("spot_prices")
666            .expect("spot_prices requirement conflicts (bug)")
667    }
668
669    fn timeout(&self) -> Duration {
670        self.timeout
671    }
672}
673
674#[cfg(test)]
675mod tests {
676    use std::sync::Arc;
677
678    use num_bigint::BigInt;
679    use tokio::sync::RwLock;
680    use tycho_simulation::{
681        tycho_common::{models::Address, simulation::protocol_sim::ProtocolSim},
682        tycho_ethereum::gas::{BlockGasPrice, GasPrice},
683    };
684
685    use super::*;
686    use crate::{
687        algorithm::test_utils::{component, order, token, MockProtocolSim},
688        derived::{types::TokenGasPrices, DerivedData},
689        feed::market_data::{MarketData, MarketState},
690        graph::GraphManager,
691        types::quote::OrderSide,
692    };
693
694    // ==================== Test Utilities ====================
695
696    /// Sets up market and graph with `()` edge weights for BellmanFord tests.
697    fn setup_market_bf(
698        pools: Vec<(&str, &Token, &Token, MockProtocolSim)>,
699    ) -> (MarketData, PetgraphStableDiGraphManager<()>) {
700        let mut market = MarketState::new();
701
702        market.update_gas_price(BlockGasPrice {
703            block_number: 1,
704            block_hash: Default::default(),
705            block_timestamp: 0,
706            pricing: GasPrice::Legacy { gas_price: BigUint::from(100u64) },
707        });
708        market.update_last_updated(crate::types::BlockInfo::new(1, "0x00".into(), 0));
709
710        for (pool_id, token_in, token_out, state) in pools {
711            let tokens = vec![token_in.clone(), token_out.clone()];
712            let comp = component(pool_id, &tokens);
713            market.upsert_components(std::iter::once(comp));
714            market.update_states([(pool_id.to_string(), Box::new(state) as Box<dyn ProtocolSim>)]);
715            market.upsert_tokens(tokens);
716        }
717
718        let mut graph_manager = PetgraphStableDiGraphManager::default();
719        graph_manager.initialize_graph(&market.component_topology());
720
721        (MarketData::new(Arc::new(RwLock::new(market))), graph_manager)
722    }
723
724    fn setup_derived_with_token_prices(
725        token_addresses: &[Address],
726    ) -> crate::derived::SharedDerivedDataRef {
727        use tycho_simulation::tycho_core::simulation::protocol_sim::Price;
728
729        let mut token_prices: TokenGasPrices = HashMap::new();
730        for address in token_addresses {
731            token_prices.insert(
732                address.clone(),
733                Price { numerator: BigUint::from(1u64), denominator: BigUint::from(1u64) },
734            );
735        }
736
737        let mut derived_data = DerivedData::new();
738        derived_data.set_token_prices(token_prices, vec![], 1, true);
739        Arc::new(RwLock::new(derived_data))
740    }
741
742    fn bf_algorithm(max_hops: usize, timeout_ms: u64) -> BellmanFordAlgorithm {
743        BellmanFordAlgorithm::with_config(
744            AlgorithmConfig::new(1, max_hops, Duration::from_millis(timeout_ms), None).unwrap(),
745        )
746        .unwrap()
747    }
748
749    // ==================== Unit Tests ====================
750
751    #[tokio::test]
752    async fn test_linear_path_found() {
753        let token_a = token(0x01, "A");
754        let token_b = token(0x02, "B");
755        let token_c = token(0x03, "C");
756        let token_d = token(0x04, "D");
757
758        let (market, manager) = setup_market_bf(vec![
759            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
760            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
761            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(4.0)),
762        ]);
763
764        let algo = bf_algorithm(4, 1000);
765        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
766
767        let result = algo
768            .find_best_route(manager.graph(), market, None, None, &ord)
769            .await
770            .unwrap();
771
772        assert_eq!(result.route().swaps().len(), 3);
773        // A->B: 100*2=200, B->C: 200*3=600, C->D: 600*4=2400
774        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(200u64));
775        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(600u64));
776        assert_eq!(result.route().swaps()[2].amount_out(), &BigUint::from(2400u64));
777    }
778
779    #[tokio::test]
780    async fn test_picks_better_of_two_paths() {
781        // Diamond graph: A->B->D (2*3=6x) vs A->C->D (4*1=4x)
782        let token_a = token(0x01, "A");
783        let token_b = token(0x02, "B");
784        let token_c = token(0x03, "C");
785        let token_d = token(0x04, "D");
786
787        let (market, manager) = setup_market_bf(vec![
788            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
789            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(3.0)),
790            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(4.0)),
791            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(1.0)),
792        ]);
793
794        let algo = bf_algorithm(3, 1000);
795        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
796
797        let result = algo
798            .find_best_route(manager.graph(), market, None, None, &ord)
799            .await
800            .unwrap();
801
802        // A->B->D: 100*2*3=600 is better than A->C->D: 100*4*1=400
803        assert_eq!(result.route().swaps().len(), 2);
804        assert_eq!(result.route().swaps()[0].component_id(), "pool_ab");
805        assert_eq!(result.route().swaps()[1].component_id(), "pool_bd");
806        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(600u64));
807    }
808
809    #[tokio::test]
810    async fn test_parallel_pools() {
811        // Two pools between A and B with different multipliers
812        let token_a = token(0x01, "A");
813        let token_b = token(0x02, "B");
814
815        let (market, manager) = setup_market_bf(vec![
816            ("pool1", &token_a, &token_b, MockProtocolSim::new(2.0)),
817            ("pool2", &token_a, &token_b, MockProtocolSim::new(5.0)),
818        ]);
819
820        let algo = bf_algorithm(2, 1000);
821        let ord = order(&token_a, &token_b, 100, OrderSide::Sell);
822
823        let result = algo
824            .find_best_route(manager.graph(), market, None, None, &ord)
825            .await
826            .unwrap();
827
828        assert_eq!(result.route().swaps().len(), 1);
829        assert_eq!(result.route().swaps()[0].component_id(), "pool2");
830        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(500u64));
831    }
832
833    #[tokio::test]
834    async fn test_no_path_returns_error() {
835        let token_a = token(0x01, "A");
836        let token_b = token(0x02, "B");
837        let token_c = token(0x03, "C");
838
839        // A-B connected, C disconnected
840        let (market, manager) =
841            setup_market_bf(vec![("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0))]);
842
843        // Add token_c to market without connecting it
844        {
845            let mut m = market.write().await;
846            m.upsert_tokens(vec![token_c.clone()]);
847        }
848
849        let algo = bf_algorithm(3, 1000);
850        let ord = order(&token_a, &token_c, 100, OrderSide::Sell);
851
852        let result = algo
853            .find_best_route(manager.graph(), market, None, None, &ord)
854            .await;
855        assert!(matches!(result, Err(AlgorithmError::NoPath { .. })));
856    }
857
858    #[tokio::test]
859    async fn test_source_not_in_graph() {
860        let token_a = token(0x01, "A");
861        let token_b = token(0x02, "B");
862        let token_x = token(0x99, "X");
863
864        let (market, manager) =
865            setup_market_bf(vec![("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0))]);
866
867        let algo = bf_algorithm(3, 1000);
868        let ord = order(&token_x, &token_b, 100, OrderSide::Sell);
869
870        let result = algo
871            .find_best_route(manager.graph(), market, None, None, &ord)
872            .await;
873        assert!(matches!(
874            result,
875            Err(AlgorithmError::NoPath { reason: NoPathReason::SourceTokenNotInGraph, .. })
876        ));
877    }
878
879    #[tokio::test]
880    async fn test_destination_not_in_graph() {
881        let token_a = token(0x01, "A");
882        let token_b = token(0x02, "B");
883        let token_x = token(0x99, "X");
884
885        let (market, manager) =
886            setup_market_bf(vec![("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0))]);
887
888        let algo = bf_algorithm(3, 1000);
889        let ord = order(&token_a, &token_x, 100, OrderSide::Sell);
890
891        let result = algo
892            .find_best_route(manager.graph(), market, None, None, &ord)
893            .await;
894        assert!(matches!(
895            result,
896            Err(AlgorithmError::NoPath { reason: NoPathReason::DestinationTokenNotInGraph, .. })
897        ));
898    }
899
900    #[tokio::test]
901    async fn test_respects_max_hops() {
902        // Path A->B->C->D exists but requires 3 hops; max_hops=2
903        let token_a = token(0x01, "A");
904        let token_b = token(0x02, "B");
905        let token_c = token(0x03, "C");
906        let token_d = token(0x04, "D");
907
908        let (market, manager) = setup_market_bf(vec![
909            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
910            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
911            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(4.0)),
912        ]);
913
914        let algo = bf_algorithm(2, 1000);
915        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
916
917        let result = algo
918            .find_best_route(manager.graph(), market, None, None, &ord)
919            .await;
920        assert!(
921            matches!(result, Err(AlgorithmError::NoPath { .. })),
922            "Should not find 3-hop path with max_hops=2"
923        );
924    }
925
926    #[tokio::test]
927    async fn test_source_token_revisit_blocked() {
928        // Forbid-revisits prevents paths like A->B->A->B->C. The algorithm
929        // should find the direct A->B->C path instead.
930        let token_a = token(0x01, "A");
931        let token_b = token(0x02, "B");
932        let token_c = token(0x03, "C");
933
934        let (market, manager) = setup_market_bf(vec![
935            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
936            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
937        ]);
938
939        let algo = bf_algorithm(4, 1000);
940        let ord = order(&token_a, &token_c, 100, OrderSide::Sell);
941
942        let result = algo
943            .find_best_route(manager.graph(), market, None, None, &ord)
944            .await
945            .unwrap();
946
947        // Should find exactly the 2-hop path A->B->C = 100*2*3 = 600
948        assert_eq!(result.route().swaps().len(), 2);
949        assert_eq!(result.route().swaps()[0].component_id(), "pool_ab");
950        assert_eq!(result.route().swaps()[1].component_id(), "pool_bc");
951        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(600u64));
952    }
953
954    #[tokio::test]
955    async fn test_hub_token_revisit_blocked() {
956        // Forbid-revisits blocks A->B->C->B->D (B visited twice).
957        // The algorithm should find the direct A->B->D = 400 instead.
958        let token_a = token(0x01, "A");
959        let token_c = token(0x02, "C");
960        let token_b = token(0x03, "B");
961        let token_d = token(0x04, "D");
962
963        let (market, manager) = setup_market_bf(vec![
964            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
965            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
966            ("pool_cb", &token_c, &token_b, MockProtocolSim::new(100.0)),
967            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(2.0)),
968        ]);
969
970        let algo = bf_algorithm(4, 1000);
971        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
972
973        let result = algo
974            .find_best_route(manager.graph(), market, None, None, &ord)
975            .await
976            .unwrap();
977
978        // Should find A->B->D = 100*2*2 = 400 (the direct 2-hop path)
979        // The 4-hop revisit path A->B->C->B->D is blocked
980        assert_eq!(result.route().swaps().len(), 2, "should use direct 2-hop path");
981        assert_eq!(result.route().swaps()[0].component_id(), "pool_ab");
982        assert_eq!(result.route().swaps()[1].component_id(), "pool_bd");
983        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(400u64));
984    }
985
986    #[tokio::test]
987    async fn test_route_amounts_are_sequential() {
988        // Verify that swap amount_in[i+1] == amount_out[i] in the built route
989        let token_a = token(0x01, "A");
990        let token_b = token(0x02, "B");
991        let token_c = token(0x03, "C");
992
993        let (market, manager) = setup_market_bf(vec![
994            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
995            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
996        ]);
997
998        let algo = bf_algorithm(3, 1000);
999        let ord = order(&token_a, &token_c, 100, OrderSide::Sell);
1000
1001        let result = algo
1002            .find_best_route(manager.graph(), market, None, None, &ord)
1003            .await
1004            .unwrap();
1005
1006        assert_eq!(result.route().swaps().len(), 2);
1007        // amount_in of second swap == amount_out of first swap
1008        assert_eq!(result.route().swaps()[1].amount_in(), result.route().swaps()[0].amount_out());
1009    }
1010
1011    #[tokio::test]
1012    async fn test_gas_deduction() {
1013        let token_a = token(0x01, "A");
1014        let token_b = token(0x02, "B");
1015
1016        let (market, manager) = setup_market_bf(vec![(
1017            "pool1",
1018            &token_a,
1019            &token_b,
1020            MockProtocolSim::new(2.0).with_gas(10),
1021        )]);
1022
1023        let algo = bf_algorithm(2, 1000);
1024        let ord = order(&token_a, &token_b, 1000, OrderSide::Sell);
1025
1026        let derived = setup_derived_with_token_prices(std::slice::from_ref(&token_b.address));
1027
1028        let result = algo
1029            .find_best_route(manager.graph(), market, None, Some(derived), &ord)
1030            .await
1031            .unwrap();
1032
1033        // Output: 1000 * 2 = 2000
1034        // Gas: 10 gas units * 100 gas_price = 1000 wei * 1/1 price = 1000
1035        // Net: 2000 - 1000 = 1000
1036        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(2000u64));
1037        assert_eq!(result.net_amount_out(), &BigInt::from(1000));
1038    }
1039
1040    #[tokio::test]
1041    async fn test_timeout_respected() {
1042        let token_a = token(0x01, "A");
1043        let token_b = token(0x02, "B");
1044        let token_c = token(0x03, "C");
1045
1046        let (market, manager) = setup_market_bf(vec![
1047            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1048            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
1049        ]);
1050
1051        // 0ms timeout
1052        let algo = bf_algorithm(3, 0);
1053        let ord = order(&token_a, &token_c, 100, OrderSide::Sell);
1054
1055        let result = algo
1056            .find_best_route(manager.graph(), market, None, None, &ord)
1057            .await;
1058
1059        // With 0ms timeout, we expect either:
1060        // - A partial result (if some layers completed before timeout check)
1061        // - Timeout error
1062        // - NoPath (if timeout prevented completing enough layers to reach dest)
1063        match result {
1064            Ok(r) => {
1065                assert!(!r.route().swaps().is_empty());
1066            }
1067            Err(AlgorithmError::Timeout { .. }) | Err(AlgorithmError::NoPath { .. }) => {
1068                // Both are acceptable for 0ms timeout
1069            }
1070            Err(e) => panic!("Unexpected error: {:?}", e),
1071        }
1072    }
1073
1074    // ==================== Integration-style Tests ====================
1075
1076    #[tokio::test]
1077    async fn test_with_fees() {
1078        let token_a = token(0x01, "A");
1079        let token_b = token(0x02, "B");
1080
1081        // Pool with 10% fee
1082        let (market, manager) = setup_market_bf(vec![(
1083            "pool1",
1084            &token_a,
1085            &token_b,
1086            MockProtocolSim::new(2.0).with_fee(0.1),
1087        )]);
1088
1089        let algo = bf_algorithm(2, 1000);
1090        let ord = order(&token_a, &token_b, 1000, OrderSide::Sell);
1091
1092        let result = algo
1093            .find_best_route(manager.graph(), market, None, None, &ord)
1094            .await
1095            .unwrap();
1096
1097        // 1000 * 2 * (1-0.1) = 1800
1098        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(1800u64));
1099    }
1100
1101    #[tokio::test]
1102    async fn test_large_trade_slippage() {
1103        let token_a = token(0x01, "A");
1104        let token_b = token(0x02, "B");
1105
1106        // Pool with limited liquidity (500 tokens)
1107        let (market, manager) = setup_market_bf(vec![(
1108            "pool1",
1109            &token_a,
1110            &token_b,
1111            MockProtocolSim::new(2.0).with_liquidity(500),
1112        )]);
1113
1114        let algo = bf_algorithm(2, 1000);
1115        let ord = order(&token_a, &token_b, 1000, OrderSide::Sell);
1116
1117        // Should fail due to insufficient liquidity
1118        let result = algo
1119            .find_best_route(manager.graph(), market, None, None, &ord)
1120            .await;
1121        assert!(
1122            matches!(result, Err(AlgorithmError::NoPath { .. })),
1123            "Should fail when trade exceeds pool liquidity"
1124        );
1125    }
1126
1127    #[tokio::test]
1128    async fn test_disconnected_tokens_return_no_path() {
1129        // A-B connected, D-E disconnected. Routing A->E should fail.
1130        let token_a = token(0x01, "A");
1131        let token_b = token(0x02, "B");
1132        let token_d = token(0x04, "D");
1133        let token_e = token(0x05, "E");
1134
1135        let (market, manager) = setup_market_bf(vec![
1136            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1137            ("pool_de", &token_d, &token_e, MockProtocolSim::new(4.0)),
1138        ]);
1139
1140        let algo = bf_algorithm(3, 1000);
1141        let ord = order(&token_a, &token_e, 100, OrderSide::Sell);
1142
1143        let result = algo
1144            .find_best_route(manager.graph(), market, None, None, &ord)
1145            .await;
1146        assert!(
1147            matches!(result, Err(AlgorithmError::NoPath { .. })),
1148            "should not find path to disconnected component"
1149        );
1150    }
1151
1152    #[tokio::test]
1153    async fn test_spfa_skips_failed_simulations() {
1154        // Pool that will fail simulation (liquidity=0 would cause error for any amount)
1155        let token_a = token(0x01, "A");
1156        let token_b = token(0x02, "B");
1157        let token_c = token(0x03, "C");
1158
1159        let (market, manager) = setup_market_bf(vec![
1160            // Direct path with failing pool
1161            ("pool_ab_bad", &token_a, &token_b, MockProtocolSim::new(2.0).with_liquidity(0)),
1162            // Alternative path that works
1163            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(2.0)),
1164            ("pool_cb", &token_c, &token_b, MockProtocolSim::new(3.0)),
1165        ]);
1166
1167        let algo = bf_algorithm(3, 1000);
1168        let ord = order(&token_a, &token_b, 100, OrderSide::Sell);
1169
1170        let result = algo
1171            .find_best_route(manager.graph(), market, None, None, &ord)
1172            .await;
1173
1174        // Should find A->C->B despite A->B failing
1175        // Note: MockProtocolSim with liquidity=0 will fail for amount > 0
1176        // The direct A->B edge should be skipped and the 2-hop path used
1177        match result {
1178            Ok(r) => {
1179                // Found alternative path
1180                assert!(!r.route().swaps().is_empty());
1181            }
1182            Err(AlgorithmError::NoPath { .. }) => {
1183                // Also acceptable if liquidity=0 blocks all paths through B
1184                // (since the failing pool might also block the reverse B->A edge)
1185            }
1186            Err(e) => panic!("Unexpected error: {:?}", e),
1187        }
1188    }
1189
1190    #[tokio::test]
1191    async fn test_resimulation_produces_correct_amounts() {
1192        // Verifies that re-simulation produces the same correct sequential amounts
1193        let token_a = token(0x01, "A");
1194        let token_b = token(0x02, "B");
1195        let token_c = token(0x03, "C");
1196
1197        let (market, manager) = setup_market_bf(vec![
1198            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1199            ("pool_bc", &token_b, &token_c, MockProtocolSim::new(3.0)),
1200        ]);
1201
1202        let algo = bf_algorithm(3, 1000);
1203        let ord = order(&token_a, &token_c, 100, OrderSide::Sell);
1204
1205        let result = algo
1206            .find_best_route(manager.graph(), market, None, None, &ord)
1207            .await
1208            .unwrap();
1209
1210        // Verify the final amounts are from re-simulation, not relaxation
1211        // A->B: 100*2=200, B->C: 200*3=600
1212        assert_eq!(result.route().swaps()[0].amount_in(), &BigUint::from(100u64));
1213        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(200u64));
1214        assert_eq!(result.route().swaps()[1].amount_in(), &BigUint::from(200u64));
1215        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(600u64));
1216    }
1217
1218    // ==================== Trait getter tests ====================
1219
1220    #[test]
1221    fn algorithm_name() {
1222        let algo = bf_algorithm(4, 200);
1223        assert_eq!(algo.name(), "bellman_ford");
1224    }
1225
1226    #[test]
1227    fn algorithm_timeout() {
1228        let algo = bf_algorithm(4, 200);
1229        assert_eq!(algo.timeout(), Duration::from_millis(200));
1230    }
1231
1232    // ==================== Forbid-revisit helper tests ====================
1233
1234    #[tokio::test]
1235    async fn test_gas_aware_relaxation_picks_cheaper_path() {
1236        // Diamond graph: A -> B -> D vs A -> C -> D
1237        // Path via B: higher gross output (3x * 2x = 6x) but extreme gas (100M per hop)
1238        // Path via C: lower gross output (2x * 2x = 4x) but cheap gas (100 per hop)
1239        //
1240        // With gas_price=100, token_prices[D]=1:1 for WETH conversion:
1241        // Path B gas cost: (100M + 100M) * 100 * 1 = 20B
1242        // Path C gas cost: (100 + 100) * 100 * 1 = 20K
1243        //
1244        // For an input of 1B:
1245        // Path B: gross = 6B, net = 6B - 20B = -14B
1246        // Path C: gross = 4B, net = 4B - 20K ≈ 4B
1247        //
1248        // Without gas awareness: Path B wins (6B > 4B)
1249        // With gas awareness: Path C wins (4B net > -14B net)
1250        let token_a = token(0x01, "A");
1251        let token_b = token(0x02, "B");
1252        let token_c = token(0x03, "C");
1253        let token_d = token(0x04, "D");
1254
1255        let high_gas: u64 = 100_000_000;
1256        let low_gas: u64 = 100;
1257
1258        let (market, manager) = setup_market_bf(vec![
1259            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(3.0).with_gas(high_gas)),
1260            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(2.0).with_gas(high_gas)),
1261            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(2.0).with_gas(low_gas)),
1262            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(2.0).with_gas(low_gas)),
1263        ]);
1264
1265        let algo = bf_algorithm(3, 1000);
1266        let ord = order(&token_a, &token_d, 1_000_000_000, OrderSide::Sell);
1267
1268        // With gas-aware relaxation (derived data with token prices + gas price in market)
1269        let derived = setup_derived_with_token_prices(&[
1270            token_a.address.clone(),
1271            token_b.address.clone(),
1272            token_c.address.clone(),
1273            token_d.address.clone(),
1274        ]);
1275
1276        let result = algo
1277            .find_best_route(manager.graph(), market, None, Some(derived), &ord)
1278            .await
1279            .unwrap();
1280
1281        // Gas-aware relaxation should pick the cheaper path A -> C -> D
1282        assert_eq!(result.route().swaps().len(), 2);
1283        assert_eq!(result.route().swaps()[0].component_id(), "pool_ac");
1284        assert_eq!(result.route().swaps()[1].component_id(), "pool_cd");
1285    }
1286
1287    #[tokio::test]
1288    async fn test_gas_aware_falls_back_to_gross_without_derived() {
1289        // Same diamond graph as above, but without derived data.
1290        // Should fall back to gross comparison and pick Path B (higher gross).
1291        let token_a = token(0x01, "A");
1292        let token_b = token(0x02, "B");
1293        let token_c = token(0x03, "C");
1294        let token_d = token(0x04, "D");
1295
1296        let high_gas: u64 = 100_000_000;
1297        let low_gas: u64 = 100;
1298
1299        let (market, manager) = setup_market_bf(vec![
1300            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(3.0).with_gas(high_gas)),
1301            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(2.0).with_gas(high_gas)),
1302            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(2.0).with_gas(low_gas)),
1303            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(2.0).with_gas(low_gas)),
1304        ]);
1305
1306        let algo = bf_algorithm(3, 1000);
1307        let ord = order(&token_a, &token_d, 1_000_000_000, OrderSide::Sell);
1308
1309        // No derived data: should fall back to gross comparison
1310        let result = algo
1311            .find_best_route(manager.graph(), market, None, None, &ord)
1312            .await
1313            .unwrap();
1314
1315        // Without gas awareness, picks the higher-gross path A -> B -> D
1316        assert_eq!(result.route().swaps().len(), 2);
1317        assert_eq!(result.route().swaps()[0].component_id(), "pool_ab");
1318        assert_eq!(result.route().swaps()[1].component_id(), "pool_bd");
1319    }
1320
1321    // ==================== Connector token tests ====================
1322
1323    /// Build a BellmanFord algorithm whose config includes a specific connector token allowlist.
1324    fn bf_algorithm_with_connectors(
1325        max_hops: usize,
1326        timeout_ms: u64,
1327        connector_tokens: HashSet<Address>,
1328    ) -> BellmanFordAlgorithm {
1329        BellmanFordAlgorithm::with_config(
1330            AlgorithmConfig::new(1, max_hops, Duration::from_millis(timeout_ms), None)
1331                .unwrap()
1332                .with_connector_tokens(connector_tokens),
1333        )
1334        .unwrap()
1335    }
1336
1337    #[tokio::test]
1338    async fn test_connector_tokens_blocks_disallowed_intermediate() {
1339        //      A
1340        //    /   \
1341        //   B     C   ← only C is in the allowlist
1342        //    \   /
1343        //      D
1344        // A->B->D is pruned; only A->C->D survives.
1345        let token_a = token(0x01, "A");
1346        let token_b = token(0x02, "B");
1347        let token_c = token(0x03, "C");
1348        let token_d = token(0x04, "D");
1349
1350        let (market, manager) = setup_market_bf(vec![
1351            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1352            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(2.0)),
1353            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(3.0)),
1354            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(3.0)),
1355        ]);
1356
1357        let connectors: HashSet<Address> = [token_c.address.clone()].into();
1358        let algo = bf_algorithm_with_connectors(3, 1000, connectors);
1359        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
1360
1361        let result = algo
1362            .find_best_route(manager.graph(), market, None, None, &ord)
1363            .await
1364            .unwrap();
1365
1366        // Only A->C->D is reachable; B was pruned.
1367        assert_eq!(result.route().swaps().len(), 2);
1368        assert_eq!(result.route().swaps()[0].component_id(), "pool_ac");
1369        assert_eq!(result.route().swaps()[1].component_id(), "pool_cd");
1370    }
1371
1372    #[tokio::test]
1373    async fn test_connector_tokens_allows_endpoints_even_if_not_listed() {
1374        // token_in (A) and token_out (B) must be reachable even when connector list is empty.
1375        let token_a = token(0x01, "A");
1376        let token_b = token(0x02, "B");
1377
1378        let (market, manager) =
1379            setup_market_bf(vec![("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0))]);
1380
1381        // Empty allowlist — no intermediate tokens allowed, but direct hop A->B should work.
1382        let algo = bf_algorithm_with_connectors(1, 1000, HashSet::new());
1383        let ord = order(&token_a, &token_b, 100, OrderSide::Sell);
1384
1385        let result = algo
1386            .find_best_route(manager.graph(), market, None, None, &ord)
1387            .await
1388            .unwrap();
1389
1390        assert_eq!(result.route().swaps().len(), 1);
1391        assert_eq!(result.route().swaps()[0].amount_out(), &BigUint::from(200u64));
1392    }
1393
1394    #[tokio::test]
1395    async fn test_connector_tokens_none_is_unrestricted() {
1396        // No connector_tokens set: both A->B->D and A->C->D are evaluated.
1397        let token_a = token(0x01, "A");
1398        let token_b = token(0x02, "B");
1399        let token_c = token(0x03, "C");
1400        let token_d = token(0x04, "D");
1401
1402        let (market, manager) = setup_market_bf(vec![
1403            ("pool_ab", &token_a, &token_b, MockProtocolSim::new(2.0)),
1404            ("pool_bd", &token_b, &token_d, MockProtocolSim::new(3.0)),
1405            ("pool_ac", &token_a, &token_c, MockProtocolSim::new(1.0)),
1406            ("pool_cd", &token_c, &token_d, MockProtocolSim::new(1.0)),
1407        ]);
1408
1409        let algo = bf_algorithm(3, 1000);
1410        let ord = order(&token_a, &token_d, 100, OrderSide::Sell);
1411
1412        let result = algo
1413            .find_best_route(manager.graph(), market, None, None, &ord)
1414            .await
1415            .unwrap();
1416
1417        // Best path is A->B->D = 100*2*3 = 600
1418        assert_eq!(result.route().swaps()[0].component_id(), "pool_ab");
1419        assert_eq!(result.route().swaps()[1].component_id(), "pool_bd");
1420        assert_eq!(result.route().swaps()[1].amount_out(), &BigUint::from(600u64));
1421    }
1422
1423    #[test]
1424    fn test_path_has_conflict_detects_node_and_pool() {
1425        // Path: 0 -[pool_a]-> 1 -[pool_b]-> 2
1426        let mut pred: Vec<Option<(NodeIndex, ComponentId)>> = vec![None; 4];
1427        pred[1] = Some((NodeIndex::new(0), "pool_a".into()));
1428        pred[2] = Some((NodeIndex::new(1), "pool_b".into()));
1429
1430        // Node conflicts: node 0 is in path, node 3 is not
1431        assert!(BellmanFordAlgorithm::path_has_conflict(
1432            NodeIndex::new(2),
1433            NodeIndex::new(0),
1434            &"any".into(),
1435            &pred
1436        ));
1437        assert!(!BellmanFordAlgorithm::path_has_conflict(
1438            NodeIndex::new(2),
1439            NodeIndex::new(3),
1440            &"any".into(),
1441            &pred
1442        ));
1443        // Self-check: node 2 is itself in the "path from 2"
1444        assert!(BellmanFordAlgorithm::path_has_conflict(
1445            NodeIndex::new(2),
1446            NodeIndex::new(2),
1447            &"any".into(),
1448            &pred
1449        ));
1450
1451        // Pool conflicts: pool_a and pool_b are used, pool_c is not
1452        assert!(BellmanFordAlgorithm::path_has_conflict(
1453            NodeIndex::new(2),
1454            NodeIndex::new(3),
1455            &"pool_a".into(),
1456            &pred
1457        ));
1458        assert!(BellmanFordAlgorithm::path_has_conflict(
1459            NodeIndex::new(2),
1460            NodeIndex::new(3),
1461            &"pool_b".into(),
1462            &pred
1463        ));
1464        assert!(!BellmanFordAlgorithm::path_has_conflict(
1465            NodeIndex::new(2),
1466            NodeIndex::new(3),
1467            &"pool_c".into(),
1468            &pred
1469        ));
1470    }
1471}