1use num_bigint::BigUint;
2use rustc_hash::FxHashMap;
3use tycho_simulation::tycho_common::{models::Address, simulation::errors::SimulationError};
4
5use crate::{algorithm::sim_meter, ComponentId};
6
7const INTERPOLATION_GAP_PERCENT: u32 = 10;
10
11#[derive(PartialEq, Eq, Hash)]
14pub struct PoolDirection<'a> {
15 pub component_id: &'a ComponentId,
17 pub address_in: &'a Address,
19 pub address_out: &'a Address,
21}
22
23#[derive(Clone, Copy, PartialEq, Eq)]
30pub enum Refusal {
31 OverLimit,
33 Failed,
35}
36
37impl Refusal {
38 pub fn of(error: &SimulationError) -> Self {
41 match error {
42 SimulationError::InvalidInput(_, _) => Refusal::OverLimit,
43 SimulationError::FatalError(_) | SimulationError::RecoverableError(_) => {
44 Refusal::Failed
45 }
46 }
47 }
48}
49
50#[derive(Clone)]
52pub struct SwapResult {
53 pub amount_out: BigUint,
55 pub gas: BigUint,
57}
58
59#[derive(Default)]
65pub struct SwappedAmounts {
66 amounts_and_results: Vec<(BigUint, Option<SwapResult>)>,
68 failed_at: Option<BigUint>,
71}
72
73impl SwappedAmounts {
74 fn refuses(&self, amount_in: &BigUint) -> bool {
76 self.failed_at
77 .as_ref()
78 .is_some_and(|refused_from| amount_in >= refused_from)
79 }
80
81 fn record(
95 &mut self,
96 insert_at: usize,
97 amount_in: &BigUint,
98 outcome: Result<SwapResult, Refusal>,
99 ) {
100 let refusal = outcome.as_ref().err().copied();
101 self.amounts_and_results
102 .insert(insert_at, (amount_in.clone(), outcome.ok()));
103 if refusal != Some(Refusal::OverLimit) {
104 return;
105 }
106
107 let was_served = |(_, outcome): &(BigUint, Option<SwapResult>)| outcome.is_some();
108 let served_below = self.amounts_and_results[..insert_at]
109 .iter()
110 .any(was_served);
111 let served_above = self.amounts_and_results[insert_at + 1..]
112 .iter()
113 .any(was_served);
114 if !served_below || served_above {
115 return;
116 }
117
118 let lowest_refused = match self.failed_at.take() {
119 Some(already_refused) if already_refused <= *amount_in => already_refused,
120 _ => amount_in.clone(),
121 };
122 self.failed_at = Some(lowest_refused);
123 }
124}
125
126pub struct SwapCache<'a> {
134 by_direction: FxHashMap<PoolDirection<'a>, SwappedAmounts>,
135}
136
137impl Default for SwapCache<'_> {
138 fn default() -> Self {
139 Self::new()
140 }
141}
142
143impl<'a> SwapCache<'a> {
144 #[must_use]
146 pub fn new() -> Self {
147 Self { by_direction: FxHashMap::default() }
148 }
149
150 pub fn swap(
157 &mut self,
158 direction: PoolDirection<'a>,
159 amount_in: &BigUint,
160 label: &'static str,
161 simulate: impl FnOnce() -> Result<SwapResult, Refusal>,
162 may_interpolate: bool,
163 ) -> Option<SwapResult> {
164 let component_id = direction.component_id;
165 let amounts_swapped = self
166 .by_direction
167 .entry(direction)
168 .or_default();
169
170 let insert_at = match amounts_swapped
171 .amounts_and_results
172 .binary_search_by(|(amount, _)| amount.cmp(amount_in))
173 {
174 Ok(asked_before) => {
175 sim_meter::record_cache_hit(component_id, label);
176 return amounts_swapped.amounts_and_results[asked_before]
177 .1
178 .clone();
179 }
180 Err(insert_at) => insert_at,
181 };
182
183 if amounts_swapped.refuses(amount_in) {
186 sim_meter::record_refusal_without_calling(component_id, label);
187 return None;
188 }
189
190 if may_interpolate {
191 if let Some(read_across) = Self::interpolate(amounts_swapped, insert_at, amount_in) {
192 sim_meter::record_interpolation(component_id, label);
193 return Some(read_across);
194 }
195 }
196
197 let outcome = simulate();
200 amounts_swapped.record(insert_at, amount_in, outcome.clone());
201 outcome.ok()
202 }
203
204 fn interpolate(
226 amounts_swapped: &SwappedAmounts,
227 insert_at: usize,
228 amount_in: &BigUint,
229 ) -> Option<SwapResult> {
230 let (lower_amount, lower) = amounts_swapped
231 .amounts_and_results
232 .get(insert_at.checked_sub(1)?)?;
233 let (upper_amount, upper) = amounts_swapped
234 .amounts_and_results
235 .get(insert_at)?;
236 let (lower, upper) = (lower.as_ref()?, upper.as_ref()?);
237
238 let amount_gap = upper_amount - lower_amount;
239 if &amount_gap * 100u32 > amount_in * INTERPOLATION_GAP_PERCENT {
240 return None;
241 }
242 if upper.amount_out < lower.amount_out {
244 return None;
245 }
246
247 let output_gap = &upper.amount_out - &lower.amount_out;
248 let amount_past_lower = amount_in - lower_amount;
249 let amount_out = &lower.amount_out + output_gap * amount_past_lower / amount_gap;
250 Some(SwapResult { amount_out, gas: upper.gas.clone() })
251 }
252}
253
254#[cfg(test)]
255mod tests {
256 use super::*;
257 use crate::algorithm::test_utils::addr;
258
259 const RANKING: &str = "ranking";
262 const COMMITTING: &str = "chunking";
263 const EXCHANGE: &str = "exchange";
264 const INTERPOLATES: bool = true;
265 const NO_INTERPOLATION: bool = false;
266
267 fn hop(amount_out: u64, gas: u64) -> SwapResult {
270 SwapResult { amount_out: BigUint::from(amount_out), gas: BigUint::from(gas) }
271 }
272
273 fn cache_holding(amounts: Vec<(u64, Option<SwapResult>)>) -> SwappedAmounts {
275 SwappedAmounts {
276 amounts_and_results: amounts
277 .into_iter()
278 .map(|(amount, outcome)| (BigUint::from(amount), outcome))
279 .collect(),
280 failed_at: None,
281 }
282 }
283
284 fn insert_at(swapped: &SwappedAmounts, amount: &BigUint) -> usize {
286 swapped
287 .amounts_and_results
288 .binary_search_by(|(known, _)| known.cmp(amount))
289 .expect_err("amount must not already be recorded")
290 }
291
292 fn read_across(swapped: &SwappedAmounts, amount: u64) -> Option<SwapResult> {
293 let amount = BigUint::from(amount);
294 SwapCache::interpolate(swapped, insert_at(swapped, &amount), &amount)
295 }
296
297 #[test]
299 fn test_interpolate_reads_across_two_amounts() {
300 let swapped = cache_holding(vec![(1000, Some(hop(2000, 50))), (1040, Some(hop(2080, 90)))]);
301
302 let across = read_across(&swapped, 1020).expect("bracketed and inside the gap");
303
304 assert_eq!(across.amount_out, BigUint::from(2040u64));
305 }
306
307 #[test]
310 fn test_interpolate_takes_gas_from_the_larger_amount() {
311 let swapped = cache_holding(vec![(1000, Some(hop(2000, 50))), (1040, Some(hop(2080, 90)))]);
312
313 let nearer_the_lower = read_across(&swapped, 1001).expect("bracketed and inside the gap");
314
315 assert_eq!(nearer_the_lower.gas, BigUint::from(90u64));
316 }
317
318 #[test]
320 fn test_interpolate_declines_a_wide_gap() {
321 let swapped = cache_holding(vec![(1000, Some(hop(2000, 50))), (1500, Some(hop(2600, 90)))]);
322
323 assert!(read_across(&swapped, 1200).is_none());
324 }
325
326 #[test]
329 fn test_interpolate_declines_above_the_largest_amount() {
330 let swapped = cache_holding(vec![(1000, Some(hop(2000, 50))), (1040, Some(hop(2080, 90)))]);
331
332 assert!(read_across(&swapped, 1050).is_none());
333 }
334
335 #[test]
337 fn test_interpolate_declines_when_output_falls() {
338 let swapped = cache_holding(vec![(1000, Some(hop(2000, 50))), (1040, Some(hop(1900, 90)))]);
339
340 assert!(read_across(&swapped, 1020).is_none());
341 }
342
343 #[test]
345 fn test_interpolate_declines_across_a_refusal() {
346 let swapped = cache_holding(vec![(1000, Some(hop(2000, 50))), (1040, None)]);
347
348 assert!(read_across(&swapped, 1020).is_none());
349 }
350
351 fn record(swapped: &mut SwappedAmounts, amount: u64, outcome: Result<SwapResult, Refusal>) {
354 let amount = BigUint::from(amount);
355 let at = insert_at(swapped, &amount);
356 swapped.record(at, &amount, outcome);
357 }
358
359 #[test]
361 fn test_refusal_above_a_served_amount_reaches_upwards() {
362 let mut swapped = cache_holding(vec![(1000, Some(hop(2000, 50)))]);
363
364 record(&mut swapped, 2000, Err(Refusal::OverLimit));
365
366 assert!(swapped.refuses(&BigUint::from(2000u64)));
367 assert!(swapped.refuses(&BigUint::from(5000u64)));
368 assert!(!swapped.refuses(&BigUint::from(1500u64)));
369 }
370
371 #[test]
374 fn test_refusal_with_nothing_served_below_stands_alone() {
375 let mut swapped = cache_holding(vec![]);
376
377 record(&mut swapped, 1000, Err(Refusal::OverLimit));
378
379 assert!(!swapped.refuses(&BigUint::from(5000u64)));
380 }
381
382 #[test]
384 fn test_refusal_below_a_served_amount_does_not_reach_upwards() {
385 let mut swapped =
386 cache_holding(vec![(1000, Some(hop(2000, 50))), (3000, Some(hop(5000, 50)))]);
387
388 record(&mut swapped, 2000, Err(Refusal::OverLimit));
389
390 assert!(!swapped.refuses(&BigUint::from(4000u64)));
391 }
392
393 struct CountingPool {
398 component_id: ComponentId,
399 address_in: Address,
400 address_out: Address,
401 calls: std::cell::Cell<usize>,
402 answer: Option<SwapResult>,
403 }
404
405 impl CountingPool {
406 fn paying(amount_out: u64) -> Self {
407 Self {
408 component_id: ComponentId::from("pool"),
409 address_in: addr(0x01),
410 address_out: addr(0x02),
411 calls: std::cell::Cell::new(0),
412 answer: Some(hop(amount_out, 10)),
413 }
414 }
415
416 fn refusing() -> Self {
417 Self { answer: None, ..Self::paying(0) }
418 }
419
420 fn direction(&self) -> PoolDirection<'_> {
421 PoolDirection {
422 component_id: &self.component_id,
423 address_in: &self.address_in,
424 address_out: &self.address_out,
425 }
426 }
427
428 fn ask<'a>(
429 &'a self,
430 cache: &mut SwapCache<'a>,
431 amount: u64,
432 label: &'static str,
433 may_interpolate: bool,
434 ) -> Option<SwapResult> {
435 cache.swap(
436 self.direction(),
437 &BigUint::from(amount),
438 label,
439 || {
440 self.calls.set(self.calls.get() + 1);
441 self.answer
442 .clone()
443 .ok_or(Refusal::OverLimit)
444 },
445 may_interpolate,
446 )
447 }
448 }
449
450 #[test]
452 fn test_swap_answers_a_repeated_amount_without_calling() {
453 let pool = CountingPool::paying(2000);
454 let mut cache = SwapCache::new();
455
456 let first = pool.ask(&mut cache, 1000, COMMITTING, NO_INTERPOLATION);
457 let second = pool.ask(&mut cache, 1000, COMMITTING, NO_INTERPOLATION);
458
459 assert_eq!(pool.calls.get(), 1);
460 assert_eq!(first.map(|h| h.amount_out), Some(BigUint::from(2000u64)));
461 assert_eq!(second.map(|h| h.amount_out), Some(BigUint::from(2000u64)));
462 }
463
464 #[test]
467 fn test_swap_short_circuits_above_a_refusal() {
468 let pool = CountingPool::refusing();
469 let mut cache = SwapCache::new();
470 cache.swap(
472 pool.direction(),
473 &BigUint::from(500u64),
474 COMMITTING,
475 || Ok(hop(1000, 10)),
476 NO_INTERPOLATION,
477 );
478 pool.ask(&mut cache, 1000, COMMITTING, NO_INTERPOLATION);
479 let calls_after_refusal = pool.calls.get();
480
481 let larger = pool.ask(&mut cache, 5000, COMMITTING, NO_INTERPOLATION);
482
483 assert!(larger.is_none());
484 assert_eq!(pool.calls.get(), calls_after_refusal, "the pool was asked again");
485 }
486
487 #[test]
491 fn test_swap_interpolates_only_for_a_pass_that_allows_it() {
492 let interpolating = CountingPool::paying(0);
493 let mut cache = SwapCache::new();
494 cache.swap(
495 interpolating.direction(),
496 &BigUint::from(1000u64),
497 RANKING,
498 || Ok(hop(1000, 10)),
499 INTERPOLATES,
500 );
501 cache.swap(
502 interpolating.direction(),
503 &BigUint::from(1100u64),
504 RANKING,
505 || Ok(hop(1100, 10)),
506 INTERPOLATES,
507 );
508 let calls_before = interpolating.calls.get();
509
510 let read_across = interpolating.ask(&mut cache, 1050, RANKING, INTERPOLATES);
511 assert_eq!(interpolating.calls.get(), calls_before, "ranking should not have called");
512 assert_eq!(read_across.map(|h| h.amount_out), Some(BigUint::from(1050u64)));
513
514 let simulated = interpolating.ask(&mut cache, 1060, EXCHANGE, NO_INTERPOLATION);
515 assert_eq!(interpolating.calls.get(), calls_before + 1, "exchange must call the pool");
516 assert_eq!(simulated.map(|h| h.amount_out), Some(BigUint::from(0u64)));
517 }
518
519 #[test]
522 fn test_swap_does_not_store_an_interpolated_answer() {
523 let pool = CountingPool::paying(7777);
524 let mut cache = SwapCache::new();
525 cache.swap(
526 pool.direction(),
527 &BigUint::from(1000u64),
528 RANKING,
529 || Ok(hop(1000, 10)),
530 INTERPOLATES,
531 );
532 cache.swap(
533 pool.direction(),
534 &BigUint::from(1100u64),
535 RANKING,
536 || Ok(hop(1100, 10)),
537 INTERPOLATES,
538 );
539 pool.ask(&mut cache, 1050, RANKING, INTERPOLATES);
540
541 let asked_again = pool.ask(&mut cache, 1050, EXCHANGE, NO_INTERPOLATION);
542
543 assert_eq!(pool.calls.get(), 1, "the interpolated answer should not have been stored");
544 assert_eq!(asked_again.map(|h| h.amount_out), Some(BigUint::from(7777u64)));
545 }
546
547 #[test]
551 fn test_failure_that_is_not_a_limit_does_not_reach_upwards() {
552 let mut swapped = cache_holding(vec![(1000, Some(hop(2000, 50)))]);
553
554 record(&mut swapped, 2000, Err(Refusal::Failed));
555
556 assert!(!swapped.refuses(&BigUint::from(2000u64)));
557 assert!(!swapped.refuses(&BigUint::from(5000u64)));
558 }
559
560 #[test]
562 fn test_lower_refusal_moves_the_refusal_point_down() {
563 let mut swapped = cache_holding(vec![(1000, Some(hop(2000, 50)))]);
564
565 record(&mut swapped, 3000, Err(Refusal::OverLimit));
566 record(&mut swapped, 2000, Err(Refusal::OverLimit));
567
568 assert!(swapped.refuses(&BigUint::from(2000u64)));
569 }
570}