1pub mod fee_tier_fetcher;
13
14use std::sync::{Arc, RwLock};
15
16use num_bigint::BigUint;
17use rustc_hash::{FxHashMap, FxHashSet};
18use tycho_simulation::tycho_common::models::Address;
19
20use crate::{
21 feed::{events::MarketEvent, market_data::MarketDataView},
22 replay::replay_route,
23 types::{ComponentId, Route, Swap},
24};
25
26pub fn has_pamm_leg(route: &Route) -> bool {
31 route.swaps().iter().any(|swap| {
32 swap.protocol()
33 .starts_with(PROPAMM_FALLBACK_PREFIX)
34 })
35}
36
37pub fn fallback_amount_out(
47 route: &Route,
48 market: &MarketDataView<'_>,
49 fee_tiers: &FeeTiers,
50 index: &FallbackPoolIndex,
51) -> FallbackAmountOut {
52 let mut substituted = Vec::with_capacity(route.swaps().len());
53 for swap in route.swaps() {
54 if !swap
55 .protocol()
56 .starts_with(PROPAMM_FALLBACK_PREFIX)
57 {
58 substituted.push(swap.clone());
59 continue;
60 }
61
62 let fee_tier = fee_tiers.resolved_tier(swap.token_in(), swap.token_out());
63 let Some(pool) = index.pool_for(swap.token_in(), swap.token_out(), fee_tier) else {
64 return FallbackAmountOut::NoFallbackPool {
65 component_id: swap.component_id().to_string(),
66 fee_tier,
67 };
68 };
69 let (Some(component), Some(state)) =
70 (market.get_component(pool), market.get_simulation_state(pool))
71 else {
72 return FallbackAmountOut::NoFallbackPool {
75 component_id: swap.component_id().to_string(),
76 fee_tier,
77 };
78 };
79 substituted.push(
80 Swap::new(
81 pool.clone(),
82 component.protocol_system.clone(),
83 swap.token_in().clone(),
84 swap.token_out().clone(),
85 swap.amount_in().clone(),
86 swap.amount_out().clone(),
87 swap.gas_estimate().clone(),
88 component.clone(),
89 state.clone_box(),
90 )
91 .with_split(*swap.split()),
92 );
93 }
94
95 let replayed_components: Vec<ComponentId> = substituted
98 .iter()
99 .map(|swap| swap.component_id().to_string())
100 .collect();
101 let market_state = market.extract_subset_with_overlay(
102 &replayed_components
103 .iter()
104 .collect::<FxHashSet<_>>(),
105 );
106
107 let substituted = match Route::new(substituted, FxHashMap::default()) {
108 Ok(route) => route,
109 Err(e) => return FallbackAmountOut::NotPriceable { reason: e.to_string() },
110 };
111 match replay_route(&substituted, &market_state) {
112 Ok(replay) => FallbackAmountOut::AmountOut(replay.amount_out),
113 Err(e) => FallbackAmountOut::NotPriceable { reason: e.to_string() },
114 }
115}
116
117#[derive(Debug, Clone, PartialEq, Eq)]
119pub enum FallbackAmountOut {
120 AmountOut(BigUint),
122 NoFallbackPool {
124 component_id: ComponentId,
126 fee_tier: u32,
128 },
129 NotPriceable {
132 reason: String,
134 },
135}
136
137#[derive(Debug, Clone)]
142pub struct FeeTiers {
143 default_tier: u32,
144 per_pair: FxHashMap<(Address, Address), u32>,
145}
146
147impl FeeTiers {
148 pub fn new(default_tier: u32) -> Self {
150 Self { default_tier, per_pair: FxHashMap::default() }
151 }
152
153 pub fn resolved_tier(&self, token_a: &Address, token_b: &Address) -> u32 {
155 self.per_pair
156 .get(&sorted_pair(token_a, token_b))
157 .copied()
158 .unwrap_or(self.default_tier)
159 }
160
161 pub fn set_pair_tier(&mut self, token_a: &Address, token_b: &Address, fee_tier: u32) {
163 let key = sorted_pair(token_a, token_b);
164 if fee_tier == 0 {
165 self.per_pair.remove(&key);
166 } else {
167 self.per_pair.insert(key, fee_tier);
168 }
169 }
170
171 pub fn pair_override_count(&self) -> usize {
173 self.per_pair.len()
174 }
175}
176
177#[derive(Debug, Clone, Default)]
183pub struct SharedFeeTiers(Arc<RwLock<Option<FeeTiers>>>);
184
185impl SharedFeeTiers {
186 pub fn snapshot(&self) -> Option<FeeTiers> {
188 self.0
189 .read()
190 .expect("fallback fee tier lock poisoned")
191 .clone()
192 }
193
194 pub fn set(&self, fee_tiers: FeeTiers) {
196 *self
197 .0
198 .write()
199 .expect("fallback fee tier lock poisoned") = Some(fee_tiers);
200 }
201}
202
203#[derive(Debug, Default, Clone)]
210pub struct FallbackPoolIndex {
211 pools: FxHashMap<(Address, Address, u32), ComponentId>,
212 keys: FxHashMap<ComponentId, (Address, Address, u32)>,
214}
215
216impl FallbackPoolIndex {
217 pub fn build(market: &MarketDataView<'_>) -> Self {
220 let mut index = Self::default();
221 for component_id in market.component_topology().into_keys() {
222 index.insert(market, component_id);
223 }
224 index
225 }
226
227 pub fn apply_event(&mut self, market: &MarketDataView<'_>, event: &MarketEvent) {
232 let MarketEvent::MarketUpdated { added_components, removed_components, .. } = event;
233 for component_id in removed_components {
234 self.remove(component_id);
235 }
236 for component_id in added_components.keys() {
237 self.insert(market, component_id.clone());
238 }
239 }
240
241 fn insert(&mut self, market: &MarketDataView<'_>, component_id: ComponentId) {
247 let Some(component) = market.get_component(&component_id) else { return };
248 if component.protocol_system != FALLBACK_PROTOCOL_SYSTEM {
249 return;
250 }
251 let [token_a, token_b] = component.tokens.as_slice() else { return };
252 let Some(fee_tier) = component
253 .static_attributes
254 .get(FEE_ATTRIBUTE)
255 .and_then(parse_fee_tier)
256 else {
257 return;
258 };
259 let (low, high) = sorted_pair(token_a, token_b);
260 self.keys
261 .insert(component_id.clone(), (low.clone(), high.clone(), fee_tier));
262 self.pools
263 .insert((low, high, fee_tier), component_id);
264 }
265
266 fn remove(&mut self, component_id: &ComponentId) {
268 let Some(key) = self.keys.remove(component_id) else { return };
269 if self.pools.get(&key) == Some(component_id) {
270 self.pools.remove(&key);
271 }
272 }
273
274 pub fn pool_for(
276 &self,
277 token_a: &Address,
278 token_b: &Address,
279 fee_tier: u32,
280 ) -> Option<&ComponentId> {
281 let (low, high) = sorted_pair(token_a, token_b);
282 self.pools.get(&(low, high, fee_tier))
283 }
284}
285
286pub const PROPAMM_FALLBACK_PREFIX: &str = "propammfallback:";
288
289pub const PROPAMM_ROUTER_ADDRESS: &str = "0x4DdF368080CD7946db5b459aD591c350158175e1";
294
295pub const PROPAMM_VENUES: &[&str] = &[
302 "0x5979458912F80B96d30D4220af8E2e4925A33320",
304 "0x71e790dd841c8A9061487cb3E78C288E75cE0B3d",
306];
307
308const FALLBACK_PROTOCOL_SYSTEM: &str = "uniswap_v3";
310
311const FEE_ATTRIBUTE: &str = "fee";
313
314fn sorted_pair(token_a: &Address, token_b: &Address) -> (Address, Address) {
316 if token_a <= token_b {
317 (token_a.clone(), token_b.clone())
318 } else {
319 (token_b.clone(), token_a.clone())
320 }
321}
322
323fn parse_fee_tier(raw: &tycho_simulation::tycho_common::Bytes) -> Option<u32> {
325 let bytes = raw.as_ref();
326 if bytes.is_empty() || bytes.len() > 4 {
327 return None;
328 }
329 let mut padded = [0u8; 4];
330 padded[4 - bytes.len()..].copy_from_slice(bytes);
331 Some(u32::from_be_bytes(padded))
332}
333
334#[cfg(test)]
335mod tests {
336 use tycho_simulation::{
337 tycho_common::Bytes, tycho_core::simulation::protocol_sim::ProtocolSim,
338 };
339
340 use super::*;
341 use crate::algorithm::test_utils::{self as util, addr};
342
343 const DEFAULT_TIER: u32 = 3000;
345
346 #[test]
347 fn test_resolved_tier_defaults_and_overrides() {
348 let mut fee_tiers = FeeTiers::new(DEFAULT_TIER);
349 assert_eq!(fee_tiers.resolved_tier(&addr(1), &addr(2)), DEFAULT_TIER);
350
351 fee_tiers.set_pair_tier(&addr(2), &addr(1), 500);
352 assert_eq!(fee_tiers.resolved_tier(&addr(1), &addr(2)), 500);
354 assert_eq!(fee_tiers.resolved_tier(&addr(2), &addr(1)), 500);
355
356 fee_tiers.set_pair_tier(&addr(1), &addr(2), 0);
358 assert_eq!(fee_tiers.resolved_tier(&addr(1), &addr(2)), DEFAULT_TIER);
359 }
360
361 #[test]
363 fn test_shared_fee_tiers_empty_until_set() {
364 let shared = SharedFeeTiers::default();
365 assert!(shared.snapshot().is_none());
366
367 shared.set(FeeTiers::new(500));
368 assert_eq!(
369 shared
370 .snapshot()
371 .expect("tiers were set")
372 .resolved_tier(&addr(1), &addr(2)),
373 500
374 );
375 }
376
377 #[test]
378 fn test_parse_fee_tier() {
379 assert_eq!(parse_fee_tier(&Bytes::from(3000_i32.to_be_bytes().to_vec())), Some(3000));
380 assert_eq!(parse_fee_tier(&Bytes::from(vec![0x01, 0xf4])), Some(500));
382 assert_eq!(parse_fee_tier(&Bytes::from(Vec::new())), None);
383 assert_eq!(parse_fee_tier(&Bytes::from(vec![0u8; 5])), None);
384 }
385
386 #[test]
387 fn test_pool_for_token_order() {
388 let mut index = FallbackPoolIndex::default();
389 let (low, high) = sorted_pair(&addr(9), &addr(3));
390 index
391 .pools
392 .insert((low, high, 500), "pool".to_string());
393
394 assert_eq!(
395 index
396 .pool_for(&addr(3), &addr(9), 500)
397 .map(String::as_str),
398 Some("pool")
399 );
400 assert_eq!(
401 index
402 .pool_for(&addr(9), &addr(3), 500)
403 .map(String::as_str),
404 Some("pool")
405 );
406 assert!(index
408 .pool_for(&addr(3), &addr(9), 3000)
409 .is_none());
410 }
411
412 #[test]
415 fn test_apply_event_added_and_removed_components() {
416 let market = market_with_fallback_pool(500);
417 let view = market
418 .try_read_blocking()
419 .expect("uncontended");
420 let mut index = FallbackPoolIndex::build(&view);
421 assert_eq!(index.pools.len(), 1);
422
423 index.apply_event(
424 &view,
425 &MarketEvent::MarketUpdated {
426 added_components: FxHashMap::default(),
427 removed_components: vec![FALLBACK_POOL.to_string()],
428 updated_components: Vec::new(),
429 },
430 );
431 assert!(index.pools.is_empty());
432 assert!(index
433 .pool_for(&addr(1), &addr(2), 500)
434 .is_none());
435
436 index.apply_event(
437 &view,
438 &MarketEvent::MarketUpdated {
439 added_components: FxHashMap::from_iter([(FALLBACK_POOL.to_string(), Vec::new())]),
440 removed_components: Vec::new(),
441 updated_components: Vec::new(),
442 },
443 );
444 assert_eq!(
445 index
446 .pool_for(&addr(1), &addr(2), 500)
447 .map(String::as_str),
448 Some(FALLBACK_POOL)
449 );
450 }
451
452 #[test]
454 fn test_apply_event_state_updates() {
455 let market = market_with_fallback_pool(500);
456 let view = market
457 .try_read_blocking()
458 .expect("uncontended");
459 let mut index = FallbackPoolIndex::build(&view);
460
461 index.apply_event(
462 &view,
463 &MarketEvent::MarketUpdated {
464 added_components: FxHashMap::default(),
465 removed_components: Vec::new(),
466 updated_components: vec![FALLBACK_POOL.to_string()],
467 },
468 );
469
470 assert_eq!(index.pools.len(), 1);
471 }
472
473 const PAMM_PRICE: f64 = 2.0;
475 const FALLBACK_PRICE: f64 = 1.0;
476 const FALLBACK_POOL: &str = "0xretry";
477 const PAMM_COMPONENT: &str = "0xpamm";
478 const PAMM_PROTOCOL: &str = "propammfallback:fermiswap";
479
480 fn market_with_fallback_pool(fee_tier: u32) -> crate::feed::market_data::MarketData {
482 let (token_in, token_out) = (util::token(1, "WETH"), util::token(2, "USDC"));
483 let mut component = util::component_with_protocol(
484 FALLBACK_POOL,
485 FALLBACK_PROTOCOL_SYSTEM,
486 &[token_in.clone(), token_out.clone()],
487 );
488 component
489 .static_attributes
490 .insert(FEE_ATTRIBUTE.to_string(), Bytes::from(fee_tier.to_be_bytes().to_vec()));
491
492 let market = crate::feed::market_data::MarketData::new_shared();
493 {
494 let mut state = market.try_write().expect("uncontended");
495 state.upsert_tokens([token_in, token_out]);
496 state.upsert_components([component]);
497 state.update_states([(
498 FALLBACK_POOL.to_string(),
499 Box::new(util::MockProtocolSim::new(FALLBACK_PRICE)) as Box<dyn ProtocolSim>,
500 )]);
501 }
502 market
503 }
504
505 fn pamm_swap() -> Swap {
506 let (token_in, token_out) = (util::token(1, "WETH"), util::token(2, "USDC"));
507 Swap::new(
508 PAMM_COMPONENT.to_string(),
509 PAMM_PROTOCOL.to_string(),
510 token_in.address.clone(),
511 token_out.address.clone(),
512 BigUint::from(1_000u32),
513 BigUint::from(2_000u32),
514 BigUint::from(100_000u32),
515 util::component_with_protocol(PAMM_COMPONENT, PAMM_PROTOCOL, &[token_in, token_out]),
516 Box::new(util::MockProtocolSim::new(PAMM_PRICE)),
517 )
518 }
519
520 fn uniswap_swap() -> Swap {
522 let (token_in, token_out) = (util::token(1, "WETH"), util::token(2, "USDC"));
523 Swap::new(
524 FALLBACK_POOL.to_string(),
525 FALLBACK_PROTOCOL_SYSTEM.to_string(),
526 token_in.address.clone(),
527 token_out.address.clone(),
528 BigUint::from(1_000u32),
529 BigUint::from(1_000u32),
530 BigUint::from(100_000u32),
531 util::component_with_protocol(
532 FALLBACK_POOL,
533 FALLBACK_PROTOCOL_SYSTEM,
534 &[token_in, token_out],
535 ),
536 Box::new(util::MockProtocolSim::new(FALLBACK_PRICE)),
537 )
538 }
539
540 #[test]
542 fn test_has_pamm_leg() {
543 let non_pamm =
544 Route::new(vec![uniswap_swap()], FxHashMap::default()).expect("non-empty route");
545 assert!(!has_pamm_leg(&non_pamm));
546
547 let pamm = Route::new(vec![pamm_swap()], FxHashMap::default()).expect("non-empty route");
548 assert!(has_pamm_leg(&pamm));
549 }
550
551 #[test]
553 fn test_fallback_amount_out_with_fallback_pool() {
554 let market = market_with_fallback_pool(500);
555 let view = market
556 .try_read_blocking()
557 .expect("uncontended");
558 let index = FallbackPoolIndex::build(&view);
559 let mut fee_tiers = FeeTiers::new(DEFAULT_TIER);
560 fee_tiers.set_pair_tier(&addr(1), &addr(2), 500);
561
562 let swap = pamm_swap();
563 let route = Route::new(vec![swap.clone()], FxHashMap::default()).expect("non-empty route");
564 let amount_out = fallback_amount_out(&route, &view, &fee_tiers, &index);
565
566 assert_eq!(amount_out, FallbackAmountOut::AmountOut(BigUint::from(1_000u32)));
569 assert_eq!(*swap.amount_out(), BigUint::from(2_000u32));
570 }
571
572 #[tokio::test]
575 async fn test_fallback_amount_out_uses_the_overlay_state() {
576 let market = market_with_fallback_pool(500);
577 let label = "test_overlay".to_string();
578 let mut overlay: rustc_hash::FxHashMap<ComponentId, Box<dyn ProtocolSim>> =
579 FxHashMap::default();
580 overlay.insert(
581 FALLBACK_POOL.to_string(),
582 Box::new(util::MockProtocolSim::new(FALLBACK_PRICE / 2.0)),
583 );
584 market
585 .register_labeled_state(label.clone(), overlay, u64::MAX)
586 .await;
587
588 let view = market
589 .read_labeled(&label)
590 .await
591 .expect("registered overlay");
592 let index = FallbackPoolIndex::build(&view);
593 let mut fee_tiers = FeeTiers::new(DEFAULT_TIER);
594 fee_tiers.set_pair_tier(&addr(1), &addr(2), 500);
595
596 let route = Route::new(vec![pamm_swap()], FxHashMap::default()).expect("non-empty route");
597 let amount_out = fallback_amount_out(&route, &view, &fee_tiers, &index);
598
599 assert_eq!(amount_out, FallbackAmountOut::AmountOut(BigUint::from(500u32)));
601 }
602
603 #[test]
606 fn test_fallback_amount_out_without_pool_at_resolved_tier() {
607 let market = market_with_fallback_pool(500);
609 let view = market
610 .try_read_blocking()
611 .expect("uncontended");
612 let index = FallbackPoolIndex::build(&view);
613
614 let route = Route::new(vec![pamm_swap()], FxHashMap::default()).expect("non-empty route");
615 let amount_out = fallback_amount_out(&route, &view, &FeeTiers::new(DEFAULT_TIER), &index);
616
617 assert_eq!(
618 amount_out,
619 FallbackAmountOut::NoFallbackPool {
620 component_id: PAMM_COMPONENT.to_string(),
621 fee_tier: DEFAULT_TIER,
622 }
623 );
624 }
625
626 #[test]
628 fn test_fallback_amount_out_split_route() {
629 let market = market_with_fallback_pool(500);
630 let view = market
631 .try_read_blocking()
632 .expect("uncontended");
633 let index = FallbackPoolIndex::build(&view);
634 let mut fee_tiers = FeeTiers::new(DEFAULT_TIER);
635 fee_tiers.set_pair_tier(&addr(1), &addr(2), 500);
636
637 let route =
640 Route::new(vec![pamm_swap().with_split(0.6), pamm_swap()], FxHashMap::default())
641 .expect("non-empty route");
642 let amount_out = fallback_amount_out(&route, &view, &fee_tiers, &index);
643
644 assert!(
647 matches!(&amount_out, FallbackAmountOut::AmountOut(amount) if *amount > BigUint::ZERO),
648 "expected a priced split route, got {amount_out:?}"
649 );
650 }
651
652 #[test]
654 fn test_pool_index_non_uniswap_v3_components() {
655 let market = market_with_fallback_pool(500);
656 {
657 let mut state = market.try_write().expect("uncontended");
658 state.upsert_components([util::component_with_protocol(
659 "0xv2",
660 "uniswap_v2",
661 &[util::token(1, "WETH"), util::token(2, "USDC")],
662 )]);
663 }
664 let view = market
665 .try_read_blocking()
666 .expect("uncontended");
667
668 let index = FallbackPoolIndex::build(&view);
669 assert_eq!(index.pools.len(), 1);
670 assert_eq!(
671 index
672 .pool_for(&addr(1), &addr(2), 500)
673 .map(String::as_str),
674 Some(FALLBACK_POOL)
675 );
676 }
677}