pub mod petgraph;
pub mod token_graph;
pub use petgraph::{EdgeData, PetgraphStableDiGraphManager, StableDiGraph};
use rustc_hash::{FxHashMap, FxHashSet};
use smallvec::SmallVec;
use thiserror::Error;
pub use token_graph::{PairEdge, TokenGraph, TokenPath, TopologyGraph, TopologyGraphManager};
use tycho_simulation::{
tycho_common::{models::Address, simulation::protocol_sim::ProtocolSim},
tycho_core::models::token::Token,
};
use crate::{
derived::DerivedData,
feed::market_data::MarketDataView,
types::{ComponentId, RouteExclusions},
};
pub(crate) const INLINE_TOKENS: usize = 5;
pub const INLINE_EDGES: usize = INLINE_TOKENS - 1;
#[derive(Default)]
pub struct Path<'a, D> {
pub tokens: SmallVec<[&'a Address; INLINE_TOKENS]>,
pub edge_data: SmallVec<[&'a EdgeData<D>; INLINE_EDGES]>,
}
impl<D> Clone for Path<'_, D> {
fn clone(&self) -> Self {
Self {
tokens: SmallVec::from_slice(&self.tokens),
edge_data: SmallVec::from_slice(&self.edge_data),
}
}
}
impl<'a, D> Path<'a, D> {
pub fn new() -> Self {
Self { tokens: SmallVec::new(), edge_data: SmallVec::new() }
}
pub fn add_hop(&mut self, from: &'a Address, edge_data: &'a EdgeData<D>, to: &'a Address) {
if self.tokens.is_empty() {
self.tokens.push(from);
}
self.tokens.push(to);
self.edge_data.push(edge_data);
}
pub fn len(&self) -> usize {
self.edge_data.len()
}
pub fn is_empty(&self) -> bool {
self.edge_data.is_empty()
}
pub fn edge_iter(&self) -> &[&'a EdgeData<D>] {
&self.edge_data
}
pub fn iter(&self) -> impl Iterator<Item = (&'a Address, &'a EdgeData<D>, &'a Address)> + '_ {
self.tokens
.windows(2)
.zip(self.edge_data.iter())
.map(|(tokens, edge)| (tokens[0], *edge, tokens[1]))
}
}
#[derive(Error, Debug)]
pub enum GraphError {
#[error("Token not found in graph: {0:?}")]
TokenNotFound(Address),
#[error("Components not found in graph: {0:?}")]
ComponentsNotFound(Vec<ComponentId>),
#[error("Components with less then 2 tokens cannot be added: {0:?}")]
InvalidComponents(Vec<ComponentId>),
#[cfg(any(test, feature = "test-utils"))]
#[error("No edge found between tokens {0:?} and {1:?} for component {2}")]
MissingComponentBetweenTokens(Address, Address, ComponentId),
}
pub trait GraphManager<G>: Send + Sync
where
G: Send + Sync,
{
fn initialize_graph(&mut self, components: &FxHashMap<ComponentId, Vec<Address>>);
fn graph(&self) -> &G;
}
#[derive(Debug, Clone)]
pub struct GraphQueryFilter {
pub min_hops: usize,
pub max_hops: usize,
pub connector_tokens: Option<FxHashSet<Address>>,
}
#[derive(Debug, Clone, Copy)]
pub struct RouteSearch<'a> {
pub bounds: &'a GraphQueryFilter,
pub exclusions: &'a RouteExclusions,
}
impl RouteSearch<'_> {
#[must_use]
pub fn allows_token(self, token: &Address, endpoints: (&Address, &Address)) -> bool {
self.exclusions
.allows_token(token, endpoints) &&
(token == endpoints.0 ||
token == endpoints.1 ||
self.bounds
.connector_tokens
.as_ref()
.is_none_or(|tokens| tokens.contains(token)))
}
}
pub trait EdgeWeightFromSimAndDerived: Sized {
fn from_sim_and_derived(
sim: &dyn ProtocolSim,
component_id: &ComponentId,
token_in: &Token,
token_out: &Token,
derived: &DerivedData,
) -> Option<Self>;
}
impl EdgeWeightFromSimAndDerived for () {
fn from_sim_and_derived(
_sim: &dyn ProtocolSim,
_component_id: &ComponentId,
_token_in: &Token,
_token_out: &Token,
_derived: &DerivedData,
) -> Option<Self> {
Some(())
}
}
pub trait EdgeWeightUpdaterWithDerived {
fn update_edge_weights_with_derived(
&mut self,
market: MarketDataView<'_>,
derived: &DerivedData,
) -> usize;
}