eth_prices/router/
mod.rs

1use std::collections::HashMap;
2
3use petgraph::{
4    dot::Dot,
5    graph::{NodeIndex, UnGraph},
6};
7use tracing::info;
8
9use crate::{
10    Result,
11    asset::AssetIdentifier,
12    quoter::{AnyQuoter, RateDirection},
13};
14use route::{Route, RouteStep};
15
16pub mod route;
17
18#[derive(Debug, Clone)]
19pub struct Router {
20    pub quoters: Vec<AnyQuoter>,
21    pub graph: UnGraph<String, String>,
22    pub token_map: HashMap<String, NodeIndex<u32>>,
23}
24
25impl Default for Router {
26    fn default() -> Self {
27        Self {
28            quoters: Vec::new(),
29            graph: UnGraph::new_undirected(),
30            token_map: HashMap::new(),
31        }
32    }
33}
34
35impl FromIterator<AnyQuoter> for Router {
36    fn from_iter<T: IntoIterator<Item = AnyQuoter>>(iter: T) -> Self {
37        let mut graph = Self::default();
38        for quoter in iter {
39            graph.add_quoter(quoter);
40        }
41        graph
42    }
43}
44
45impl Router {
46    pub fn get_token_index(&self, token: &AssetIdentifier) -> Option<NodeIndex<u32>> {
47        self.token_map.get(&token.to_string()).copied()
48    }
49
50    pub fn get_token_by_index(&self, index: NodeIndex<u32>) -> Option<AssetIdentifier> {
51        self.token_map
52            .iter()
53            .find(|x| *x.1 == index)
54            .map(|(token, _)| token.clone())
55            .map(AssetIdentifier::try_from)
56            .and_then(|x| x.ok())
57    }
58
59    pub fn add_token(&mut self, token: &AssetIdentifier) -> NodeIndex<u32> {
60        match self.token_map.get(&token.to_string()) {
61            Some(node_index) => *node_index,
62            None => {
63                let slug = token.to_string();
64                let node_index = self.graph.add_node(slug.to_owned());
65                self.token_map.insert(slug, node_index);
66                node_index
67            }
68        }
69    }
70
71    pub fn add_quoter(&mut self, quoter: AnyQuoter) {
72        let slug = quoter.to_string();
73        let (token_in, token_out) = quoter.tokens();
74        self.quoters.push(quoter);
75
76        let token_in_index = self.add_token(&token_in);
77        let token_out_index = self.add_token(&token_out);
78
79        self.graph
80            .extend_with_edges([(token_in_index, token_out_index, slug)]);
81    }
82
83    pub fn to_graphviz(&self) -> String {
84        Dot::new(&self.graph).to_string()
85    }
86
87    /// compute a route given an input and output token
88    pub fn compute(
89        &self,
90        input_token: &AssetIdentifier,
91        output_token: &AssetIdentifier,
92    ) -> Result<Route> {
93        let token_a_index = self
94            .get_token_index(input_token)
95            .ok_or_else(|| crate::error::EthPricesError::AssetNotFound(input_token.to_string()))?;
96        let token_b_index = self
97            .get_token_index(output_token)
98            .ok_or_else(|| crate::error::EthPricesError::AssetNotFound(output_token.to_string()))?;
99
100        info!(
101            target: "router::compute_start",
102            input_token = %input_token,
103            output_token = %output_token,
104        );
105
106        let path = petgraph::algo::astar(
107            &self.graph,
108            token_a_index,
109            |x| x == token_b_index,
110            |_| 0,
111            |_| 0,
112        );
113
114        match path {
115            None => Err(crate::error::EthPricesError::NoRouteFound(
116                input_token.to_string(),
117                output_token.to_string(),
118            )),
119            Some((_cost, node_path)) => {
120                info!(
121                    target: "router::compute_end",
122                    node_path = ?node_path,
123                );
124                let token_route = node_path
125                    .iter()
126                    .map(|x| {
127                        self.get_token_by_index(*x)
128                            .ok_or_else(|| crate::error::EthPricesError::MissingTokenInRoute)
129                    })
130                    .collect::<Result<Vec<AssetIdentifier>>>()?;
131
132                let mut path = Vec::new();
133
134                let mut previous_token = input_token;
135                for next_token in token_route.iter() {
136                    if *previous_token == *next_token {
137                        continue;
138                    };
139
140                    let quoter = self
141                        .quoters
142                        .iter()
143                        .find(|x| {
144                            let (token_in, token_out) = x.tokens();
145
146                            (token_in == *previous_token && token_out == *next_token)
147                                || (token_in == *next_token && token_out == *previous_token)
148                        })
149                        .ok_or_else(|| crate::error::EthPricesError::MissingQuoterInRoute)?;
150
151                    path.push(RouteStep {
152                        quoter: quoter.clone(),
153                        direction: if *previous_token == quoter.tokens().0 {
154                            RateDirection::Forward
155                        } else {
156                            RateDirection::Reverse
157                        },
158                    });
159                    previous_token = next_token;
160                }
161
162                if path.len() != node_path.len() - 1 {
163                    return Err(crate::error::EthPricesError::PathLengthMismatch {
164                        expected: node_path.len() - 1,
165                        actual: path.len(),
166                    });
167                }
168
169                Ok(Route {
170                    path,
171                    input_token: input_token.clone(),
172                    output_token: output_token.clone(),
173                })
174            }
175        }
176    }
177}
eth_prices/router/mod.rs

eth_prices/router/
mod.rs
