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

eth_prices/router/
mod.rs
