1use std::collections::HashMap;
11
12use num_bigint::BigUint;
13use tycho_simulation::tycho_common::{models::Address, simulation::protocol_sim::ProtocolSim};
14
15use crate::{
16 algorithm::{sim_guard::GuardedProtocolSim, split_primitives::split_amount},
17 feed::market_data::MarketState,
18 types::{ComponentId, Route, Swap},
19};
20
21#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct RouteReplay {
24 pub amount_out: BigUint,
26 pub gas: BigUint,
28}
29
30#[derive(Debug, thiserror::Error)]
32pub enum ReplayError {
33 #[error("route has no swaps")]
36 EmptyRoute,
37 #[error("no simulation state for component {0}")]
40 MissingState(ComponentId),
41 #[error("token {0} missing from the market state")]
43 MissingToken(Address),
44 #[error("simulation failed on component {component_id}: {error}")]
46 Simulation {
47 component_id: ComponentId,
49 error: String,
51 },
52}
53
54pub fn replay_route(route: &Route, market: &MarketState) -> Result<RouteReplay, ReplayError> {
69 let swaps = route.swaps();
70 let (Some(input_token), Some(output_token)) = (route.input_token(), route.output_token())
71 else {
72 return Err(ReplayError::EmptyRoute);
73 };
74 let total_in: BigUint = swaps
75 .iter()
76 .filter(|swap| *swap.token_in() == input_token)
77 .map(Swap::amount_in)
78 .sum();
79
80 let mut available: HashMap<Address, BigUint> = HashMap::new();
84 available.insert(input_token, total_in);
85 let mut branch_totals: HashMap<Address, BigUint> = HashMap::new();
86 let mut post_swap: HashMap<ComponentId, Box<dyn ProtocolSim>> = HashMap::new();
87 let mut total_gas = BigUint::ZERO;
88
89 for swap in swaps {
90 let token_in = market
91 .get_token(swap.token_in())
92 .ok_or_else(|| ReplayError::MissingToken(swap.token_in().clone()))?;
93 let token_out = market
94 .get_token(swap.token_out())
95 .ok_or_else(|| ReplayError::MissingToken(swap.token_out().clone()))?;
96
97 let branch_total = branch_totals
98 .entry(swap.token_in().clone())
99 .or_insert_with(|| {
100 available
101 .get(swap.token_in())
102 .cloned()
103 .unwrap_or_default()
104 })
105 .clone();
106 let remaining = available
107 .entry(swap.token_in().clone())
108 .or_default();
109 let amount_in = if *swap.split() > 0.0 {
110 let (part, _) = split_amount(&branch_total, *swap.split());
113 part.min(remaining.clone())
114 } else {
115 remaining.clone()
116 };
117 *remaining -= &amount_in;
118
119 let sim = post_swap
120 .get(swap.component_id())
121 .map(|state| state.as_ref())
122 .or_else(|| market.get_simulation_state(swap.component_id()))
123 .ok_or_else(|| ReplayError::MissingState(swap.component_id().to_string()))?;
124 let result = sim
125 .get_amount_out_guarded(amount_in, token_in, token_out)
126 .map_err(|e| ReplayError::Simulation {
127 component_id: swap.component_id().to_string(),
128 error: e.to_string(),
129 })?;
130
131 total_gas += &result.gas;
132 *available
133 .entry(swap.token_out().clone())
134 .or_default() += &result.amount;
135 post_swap.insert(swap.component_id().to_string(), result.new_state);
136 }
137
138 let amount_out = available
139 .remove(&output_token)
140 .unwrap_or_default();
141 Ok(RouteReplay { amount_out, gas: total_gas })
142}
143
144#[cfg(test)]
145mod tests {
146 use rustc_hash::FxHashMap;
147
148 use super::*;
149 use crate::algorithm::test_utils::{component, token, ConstantProductSim, MockProtocolSim};
150
151 fn make_market(
152 pools: Vec<(
153 &str,
154 Vec<tycho_simulation::tycho_common::models::token::Token>,
155 Box<dyn ProtocolSim>,
156 )>,
157 ) -> MarketState {
158 let mut market = MarketState::new();
159 for (pool_id, tokens, sim) in pools {
160 market.upsert_components(std::iter::once(component(pool_id, &tokens)));
161 market.update_states([(pool_id.to_string(), sim)]);
162 market.upsert_tokens(tokens);
163 }
164 market
165 }
166
167 fn route_swap(
170 pool_id: &str,
171 token_in: &tycho_simulation::tycho_common::models::token::Token,
172 token_out: &tycho_simulation::tycho_common::models::token::Token,
173 amount_in: u64,
174 split: f64,
175 ) -> Swap {
176 Swap::new(
177 pool_id.to_string(),
178 "mock".to_string(),
179 token_in.address.clone(),
180 token_out.address.clone(),
181 BigUint::from(amount_in),
182 BigUint::ZERO,
183 BigUint::ZERO,
184 component(pool_id, &[token_in.clone(), token_out.clone()]),
185 Box::new(MockProtocolSim::new(1_000_000.0)),
186 )
187 .with_split(split)
188 }
189
190 fn route(swaps: Vec<Swap>) -> Route {
191 Route::new(swaps, FxHashMap::default()).expect("test route must not be empty")
192 }
193
194 #[test]
195 fn sequential_route_threads_amounts_through_market_state() {
196 let token_a = token(0x0A, "A");
199 let token_b = token(0x0B, "B");
200 let token_c = token(0x0C, "C");
201 let market = make_market(vec![
202 (
203 "pool_ab",
204 vec![token_a.clone(), token_b.clone()],
205 Box::new(MockProtocolSim::new(2.0).with_gas(50_000)),
206 ),
207 (
208 "pool_bc",
209 vec![token_b.clone(), token_c.clone()],
210 Box::new(MockProtocolSim::new(3.0).with_gas(70_000)),
211 ),
212 ]);
213 let route = route(vec![
214 route_swap("pool_ab", &token_a, &token_b, 1_000, 0.0),
215 route_swap("pool_bc", &token_b, &token_c, 2_000, 0.0),
216 ]);
217
218 let replay = replay_route(&route, &market).unwrap();
219 assert_eq!(replay.amount_out, BigUint::from(6_000u64));
220 assert_eq!(replay.gas, BigUint::from(120_000u64));
221 }
222
223 #[test]
224 fn split_route_divides_by_fraction_with_remainder() {
225 let token_a = token(0x0A, "A");
228 let token_b = token(0x0B, "B");
229 let market = make_market(vec![
230 ("pool_1", vec![token_a.clone(), token_b.clone()], Box::new(MockProtocolSim::new(2.0))),
231 ("pool_2", vec![token_a.clone(), token_b.clone()], Box::new(MockProtocolSim::new(3.0))),
232 ]);
233 let route = route(vec![
234 route_swap("pool_1", &token_a, &token_b, 600, 0.6),
235 route_swap("pool_2", &token_a, &token_b, 400, 0.0),
236 ]);
237
238 let replay = replay_route(&route, &market).unwrap();
239 assert_eq!(replay.amount_out, BigUint::from(2_400u64));
240 }
241
242 #[test]
243 fn splits_are_fractions_of_the_collected_total_not_the_remainder() {
244 let token_a = token(0x0A, "A");
247 let token_b = token(0x0B, "B");
248 let market = make_market(vec![
249 ("pool_1", vec![token_a.clone(), token_b.clone()], Box::new(MockProtocolSim::new(1.0))),
250 ("pool_2", vec![token_a.clone(), token_b.clone()], Box::new(MockProtocolSim::new(2.0))),
251 ("pool_3", vec![token_a.clone(), token_b.clone()], Box::new(MockProtocolSim::new(4.0))),
252 ]);
253 let route = route(vec![
254 route_swap("pool_1", &token_a, &token_b, 500, 0.5),
255 route_swap("pool_2", &token_a, &token_b, 300, 0.3),
256 route_swap("pool_3", &token_a, &token_b, 200, 0.0),
257 ]);
258
259 let replay = replay_route(&route, &market).unwrap();
261 assert_eq!(replay.amount_out, BigUint::from(1_900u64));
262 }
263
264 #[test]
265 fn shared_pool_sees_depleted_reserves() {
266 let token_a = token(0x0A, "A");
269 let token_b = token(0x0B, "B");
270 let cp = ConstantProductSim {
271 reserve_0: BigUint::from(10_000u64),
272 reserve_1: BigUint::from(10_000u64),
273 gas: 50_000,
274 };
275 let market = make_market(vec![(
276 "pool",
277 vec![token_a.clone(), token_b.clone()],
278 Box::new(cp.clone()),
279 )]);
280 let route = route(vec![
281 route_swap("pool", &token_a, &token_b, 500, 0.5),
282 route_swap("pool", &token_a, &token_b, 500, 0.0),
283 ]);
284
285 let replay = replay_route(&route, &market).unwrap();
286 let full_swap = cp
287 .get_amount_out(BigUint::from(1_000u64), &token_a, &token_b)
288 .unwrap()
289 .amount;
290 let diff = if replay.amount_out > full_swap {
291 &replay.amount_out - &full_swap
292 } else {
293 &full_swap - &replay.amount_out
294 };
295 assert!(diff <= BigUint::from(2u32), "split through one pool must match one full swap");
296 }
297
298 #[test]
299 fn changed_market_state_changes_the_output() {
300 let token_a = token(0x0A, "A");
302 let token_b = token(0x0B, "B");
303 let market = make_market(vec![(
304 "pool",
305 vec![token_a.clone(), token_b.clone()],
306 Box::new(MockProtocolSim::new(2.5)),
307 )]);
308 let route = route(vec![route_swap("pool", &token_a, &token_b, 1_000, 0.0)]);
310
311 let replay = replay_route(&route, &market).unwrap();
312 assert_eq!(replay.amount_out, BigUint::from(2_500u64));
313 }
314
315 #[test]
316 fn missing_simulation_state_errors() {
317 let token_a = token(0x0A, "A");
318 let token_b = token(0x0B, "B");
319 let mut market = make_market(vec![]);
320 market.upsert_tokens(vec![token_a.clone(), token_b.clone()]);
321 let route = route(vec![route_swap("gone", &token_a, &token_b, 1_000, 0.0)]);
322
323 let err = replay_route(&route, &market).unwrap_err();
324 assert!(matches!(err, ReplayError::MissingState(id) if id == "gone"));
325 }
326
327 #[test]
328 fn missing_token_errors() {
329 let token_a = token(0x0A, "A");
330 let token_b = token(0x0B, "B");
331 let market = make_market(vec![(
332 "pool",
333 vec![token_a.clone(), token_b.clone()],
334 Box::new(MockProtocolSim::new(2.0)),
335 )]);
336 let unknown = token(0x42, "X");
337 let route = route(vec![route_swap("pool", &unknown, &token_b, 1_000, 0.0)]);
338
339 let err = replay_route(&route, &market).unwrap_err();
340 assert!(matches!(err, ReplayError::MissingToken(addr) if addr == unknown.address));
341 }
342}