tycho_simulation/price_level_stream/
stream.rs

1use std::{
2    collections::{hash_map::Entry, HashMap, HashSet},
3    time::Duration,
4};
5
6use chrono::Utc;
7use num_bigint::BigUint;
8use tokio_stream::{Stream, StreamExt};
9use tycho_common::{
10    models::{token::Token, Chain},
11    simulation::protocol_sim::ProtocolSim,
12    Bytes,
13};
14
15use super::{
16    config::{
17        default_denied_pamms, default_served_pamms, PriceLevelStreamConfig,
18        DEFAULT_AUTO_DETECTED_GAS_COST,
19    },
20    state::{PriceLevelStreamQuote, PriceLevelStreamState},
21    titan::{
22        self, ConnectionSettings, TitanPairLevels, TitanPammLevels, TitanPriceLevel,
23        TitanPriceLevelMessage, TITAN_PRICE_LEVEL_URL,
24    },
25};
26use crate::protocol::models::{ProtocolComponent, Update};
27
28/// Static attribute under which each emitted component carries its pAMM venue address.
29pub const PAMM_ADDRESS_ATTRIBUTE: &str = "pamm_address";
30
31/// Builds a stream of [`Update`]s from the Titan pAMM price level WebSocket.
32///
33/// A new builder serves no pAMMs: register the known venues via
34/// [`with_known_pamms`](Self::with_known_pamms), individual ones via
35/// [`add_pamm`](Self::add_pamm), or opt into serving unknown streamed venues via
36/// [`auto_detect`](Self::auto_detect); [`with_tokens`](Self::with_tokens) provides the token
37/// metadata pairs are interpreted with.
38///
39/// One component is emitted per (pAMM, token pair), identified by the concatenation
40/// `pamm ++ token0 ++ token1` (tokens sorted ascending), under the protocol system
41/// `pricelevelstream:{pamm}`. The venue address is exposed through the
42/// [`PAMM_ADDRESS_ATTRIBUTE`] static attribute for downstream encoding.
43#[derive(Default)]
44pub struct PriceLevelStreamBuilder {
45    registry: HashMap<Bytes, PriceLevelStreamConfig>,
46    denied: HashSet<Bytes>,
47    tokens: HashMap<Bytes, Token>,
48    url: Option<String>,
49    auto_detect: bool,
50    auto_detected_gas_cost: Option<BigUint>,
51    connection: ConnectionSettings,
52}
53
54impl PriceLevelStreamBuilder {
55    pub fn new() -> Self {
56        Self::default()
57    }
58
59    /// Enables serving pAMMs that are not registered via
60    /// [`with_known_pamms`](Self::with_known_pamms) or [`add_pamm`](Self::add_pamm)
61    /// (disabled by default).
62    ///
63    /// When enabled, any unknown streamed venue — except denied ones (see
64    /// [`deny_pamm`](Self::deny_pamm)) — is served under its full lowercase hex address
65    /// as the name, with the default gas cost. A venue's protocol system therefore changes from
66    /// the address form (`pricelevelstream:{0xaddress}`) to a name (`pricelevelstream:{name}`)
67    /// once it gets registered — via [`add_pamm`](Self::add_pamm) or a release's
68    /// [`default_served_pamms`] recognizing it; the name-independent identifiers — the component id
69    /// and the [`PAMM_ADDRESS_ATTRIBUTE`] — stay stable across such renames.
70    pub fn auto_detect(mut self, enabled: bool) -> Self {
71        self.auto_detect = enabled;
72        self
73    }
74
75    /// Overrides the per-swap gas cost that auto-detected pAMMs (see
76    /// [`auto_detect`](Self::auto_detect)) are served with. Defaults to the maximum over the
77    /// known venue profiles, as the conservative choice. Registered venues are unaffected —
78    /// their gas cost comes from their [`PriceLevelStreamConfig`].
79    pub fn auto_detected_gas_cost(mut self, gas_cost: BigUint) -> Self {
80        self.auto_detected_gas_cost = Some(gas_cost);
81        self
82    }
83
84    /// Overrides the stream endpoint, e.g. to connect to a closer Titan region than the default
85    /// (see <https://docs.titanbuilder.xyz/propamms/takers>).
86    pub fn endpoint(mut self, url: impl Into<String>) -> Self {
87        self.url = Some(url.into());
88        self
89    }
90
91    /// Overrides how long a single connection attempt may take before it is aborted and retried
92    /// (default: 10s), so a hung TCP/TLS handshake cannot block the stream forever.
93    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
94        self.connection.connect_timeout = timeout;
95        self
96    }
97
98    /// Overrides the longest gap between Titan messages tolerated before the connection is
99    /// treated as dead and re-established (default: 30s). Titan pushes several updates per
100    /// second, so a multi-second silence means a stalled or half-open connection.
101    pub fn read_idle_timeout(mut self, timeout: Duration) -> Self {
102        self.connection.read_idle_timeout = timeout;
103        self
104    }
105
106    /// Overrides the cap on the exponential reconnect backoff of `2^attempt` seconds
107    /// (default: 32s).
108    pub fn max_backoff(mut self, max_backoff: Duration) -> Self {
109        self.connection.max_backoff = max_backoff;
110        self
111    }
112
113    /// Registers a pAMM to be served under the given configuration, overriding any default,
114    /// denied, or auto-detected one for the same address.
115    ///
116    /// Between [`add_pamm`](Self::add_pamm) and [`deny_pamm`](Self::deny_pamm) for the same
117    /// address, the later call wins; the defaults applied by
118    /// [`with_known_pamms`](Self::with_known_pamms) never override either, in any call order.
119    pub fn add_pamm(mut self, config: PriceLevelStreamConfig) -> Self {
120        self.denied.remove(&config.address);
121        self.registry
122            .insert(config.address.clone(), config);
123        self
124    }
125
126    /// Excludes a venue from being served: drops its current registration (default or explicit)
127    /// and blocks auto-detecting it.
128    ///
129    /// Between [`add_pamm`](Self::add_pamm) and [`deny_pamm`](Self::deny_pamm) for the same
130    /// address, the later call wins; the defaults applied by
131    /// [`with_known_pamms`](Self::with_known_pamms) never override either, in any call order —
132    /// so denying a venue from the default set works whether the denial comes before or after
133    /// [`with_known_pamms`](Self::with_known_pamms).
134    pub fn deny_pamm(mut self, address: Bytes) -> Self {
135        self.registry.remove(&address);
136        self.denied.insert(address);
137        self
138    }
139
140    /// Applies what is known about the streamed venues: registers the known-good ones
141    /// ([`default_served_pamms`]) to be served and denies the known-bad ones
142    /// ([`default_denied_pamms`]) — venues that stream quotes but whose swaps are not executable.
143    ///
144    /// These defaults never override an explicit [`add_pamm`](Self::add_pamm) or
145    /// [`deny_pamm`](Self::deny_pamm) for the same address, regardless of call order.
146    pub fn with_known_pamms(mut self) -> Self {
147        for config in default_served_pamms() {
148            if self.denied.contains(&config.address) {
149                continue;
150            }
151            self.registry
152                .entry(config.address.clone())
153                .or_insert(config);
154        }
155        for address in default_denied_pamms() {
156            if self.registry.contains_key(&address) {
157                continue;
158            }
159            self.denied.insert(address);
160        }
161        self
162    }
163
164    /// Provides the token metadata used to build components and interpret amounts. Pairs whose
165    /// tokens are missing here are skipped.
166    pub fn with_tokens(mut self, tokens: HashMap<Bytes, Token>) -> Self {
167        self.tokens = tokens;
168        self
169    }
170
171    /// Consumes the builder and opens the stream.
172    ///
173    /// The connection is established lazily on first poll and maintained (with reconnects) for as
174    /// long as the stream is polled; it never terminates on its own, and dropping the stream
175    /// closes the connection. Frames that contain no served pAMM produce no update.
176    ///
177    /// Each streamed frame is a complete snapshot of everything Titan currently streams, so
178    /// every update carries the full set of the frame's pair states, with `new_pairs` /
179    /// `removed_pairs` derived by diffing against the previous frame — a pair (or a whole
180    /// venue) the stream stops serving is removed. Frames older than an already processed one
181    /// are skipped, so updates never move backwards in block number. Pairs whose tokens are
182    /// missing from the provided token metadata are skipped.
183    pub fn build(self) -> impl Stream<Item = Update> + Send {
184        if self.registry.is_empty() && !self.auto_detect {
185            tracing::warn!(
186                "No pAMMs registered and auto-detection is off; the stream will never produce \
187                 an update"
188            );
189        }
190        if self.tokens.is_empty() {
191            tracing::warn!(
192                "No token metadata provided; every streamed pair will be skipped and the stream \
193                 will never produce an update"
194            );
195        }
196        let url = self
197            .url
198            .unwrap_or_else(|| TITAN_PRICE_LEVEL_URL.to_string());
199        let auto_detected_gas_cost = self
200            .auto_detected_gas_cost
201            .unwrap_or_else(|| BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST));
202        let mut tracker = SnapshotTracker::new(
203            self.registry,
204            self.denied,
205            self.tokens,
206            self.auto_detect,
207            auto_detected_gas_cost,
208        );
209
210        titan::messages(url, self.connection).filter_map(move |message| tracker.process(message))
211    }
212}
213
214/// Turns Titan frames into [`Update`]s, tracking the previously emitted components so pair
215/// additions and removals can be diffed against the last snapshot.
216struct SnapshotTracker {
217    registry: HashMap<Bytes, PriceLevelStreamConfig>,
218    /// Venues excluded from auto-detection. The builder keeps this disjoint from the registry:
219    /// denying removes any registration and registering removes any denial.
220    denied: HashSet<Bytes>,
221    tokens: HashMap<Bytes, Token>,
222    /// Whether frames from pAMMs absent from the registry get an address-named configuration
223    /// synthesized (and cached in the registry) instead of being skipped.
224    auto_detect: bool,
225    /// The per-swap gas cost synthesized auto-detected configurations are served with.
226    auto_detected_gas_cost: BigUint,
227    /// Components of the last emitted snapshot, across all pAMMs. A frame is a complete
228    /// snapshot of everything Titan currently streams, so removals are diffed globally: a
229    /// known component a frame does not re-emit is gone — including when its venue vanishes
230    /// from the stream entirely.
231    components: HashMap<String, ProtocolComponent>,
232    /// The newest block number processed so far. Frames targeting an older block (e.g.
233    /// delivered around a reconnect) are stale and skipped wholesale — processing one would
234    /// emit superseded states and churn the global diff.
235    newest_block: u64,
236}
237
238impl SnapshotTracker {
239    fn new(
240        registry: HashMap<Bytes, PriceLevelStreamConfig>,
241        denied: HashSet<Bytes>,
242        tokens: HashMap<Bytes, Token>,
243        auto_detect: bool,
244        auto_detected_gas_cost: BigUint,
245    ) -> Self {
246        Self {
247            registry,
248            denied,
249            tokens,
250            auto_detect,
251            auto_detected_gas_cost,
252            components: HashMap::new(),
253            newest_block: 0,
254        }
255    }
256
257    /// Processes one frame into an [`Update`], or `None` if the frame targets an older block
258    /// than an already processed one or contains nothing relevant (no registered pAMM with at
259    /// least one known pair or a pair removal).
260    fn process(&mut self, message: TitanPriceLevelMessage) -> Option<Update> {
261        if message.block_number < self.newest_block {
262            tracing::warn!(
263                block_number = message.block_number,
264                newest_block = self.newest_block,
265                "Skipping out-of-order price level frame"
266            );
267            return None;
268        }
269        self.newest_block = message.block_number;
270
271        let mut states: HashMap<String, Box<dyn ProtocolSim>> = HashMap::new();
272        let mut new_pairs = HashMap::new();
273        // The frame is a complete snapshot: every known component is presumed gone until the
274        // frame re-emits it below.
275        let mut previous = std::mem::take(&mut self.components);
276
277        for TitanPammLevels { pamm, pairs } in message.pamms {
278            let config = match self.registry.entry(pamm.clone()) {
279                Entry::Occupied(entry) => &*entry.into_mut(),
280                Entry::Vacant(entry) => {
281                    if !self.auto_detect {
282                        tracing::debug!(%pamm, "Skipping unregistered pAMM");
283                        continue;
284                    }
285                    if self.denied.contains(&pamm) {
286                        tracing::debug!(%pamm, "Skipping denied pAMM");
287                        continue;
288                    }
289                    tracing::info!(%pamm, "Serving auto-detected pAMM");
290                    &*entry.insert(PriceLevelStreamConfig::auto_detected(
291                        pamm.clone(),
292                        self.auto_detected_gas_cost.clone(),
293                    ))
294                }
295            };
296
297            // Merge the frame's per-direction ladders into one entry per unordered token pair.
298            let mut merged_pairs: HashMap<(Bytes, Bytes), (Vec<_>, Vec<_>)> = HashMap::new();
299            for TitanPairLevels { token_in, token_out, order_book } in pairs {
300                if !self.tokens.contains_key(&token_in) || !self.tokens.contains_key(&token_out) {
301                    tracing::debug!(%token_in, %token_out, "Skipping pair with unknown token");
302                    continue;
303                }
304                let sells_token0 = token_in < token_out;
305                let key = if sells_token0 {
306                    (token_in.clone(), token_out.clone())
307                } else {
308                    (token_out.clone(), token_in.clone())
309                };
310                let quotes = order_book
311                    .into_iter()
312                    .map(|TitanPriceLevel { amount_in, amount_out }| {
313                        PriceLevelStreamQuote::new(amount_in, amount_out)
314                    })
315                    .collect();
316                let entry = merged_pairs.entry(key).or_default();
317                if sells_token0 {
318                    entry.0 = quotes;
319                } else {
320                    entry.1 = quotes;
321                }
322            }
323
324            for ((token0, token1), (quotes_0_to_1, quotes_1_to_0)) in merged_pairs {
325                let id = component_id(&config.address, &token0, &token1);
326                let id_string = id.to_string();
327                let component = previous
328                    .remove(&id_string)
329                    .unwrap_or_else(|| {
330                        let component = build_component(&self.tokens, config, id, &token0, &token1);
331                        new_pairs.insert(id_string.clone(), component.clone());
332                        component
333                    });
334
335                let state = PriceLevelStreamState::new(
336                    token0,
337                    token1,
338                    quotes_0_to_1,
339                    quotes_1_to_0,
340                    config.gas_cost.clone(),
341                );
342
343                states.insert(id_string.clone(), Box::new(state));
344                self.components
345                    .insert(id_string, component);
346            }
347        }
348
349        // Every re-emitted pair was moved back into `self.components` above — whatever remains
350        // is gone: the pair, or its whole venue, is no longer streamed.
351        let removed_pairs = previous;
352
353        if states.is_empty() && new_pairs.is_empty() && removed_pairs.is_empty() {
354            return None;
355        }
356
357        Some(
358            // Quotes target the block currently being built, hence partial. Sync states stay
359            // empty (like the RFQ path) because no full block header is available.
360            Update::new(message.block_number, states, new_pairs)
361                .set_is_partial(true)
362                .set_removed_pairs(removed_pairs),
363        )
364    }
365}
366
367fn build_component(
368    tokens: &HashMap<Bytes, Token>,
369    config: &PriceLevelStreamConfig,
370    id: Bytes,
371    token0: &Bytes,
372    token1: &Bytes,
373) -> ProtocolComponent {
374    let protocol_system = config.protocol_system();
375    ProtocolComponent::new(
376        id,
377        protocol_system.clone(),
378        protocol_system,
379        // Titan builds Ethereum L1 blocks; the stream carries no other chains.
380        Chain::Ethereum,
381        vec![tokens[token0].clone(), tokens[token1].clone()],
382        vec![config.address.clone()],
383        HashMap::from([(PAMM_ADDRESS_ATTRIBUTE.to_string(), config.address.clone())]),
384        Bytes::default(),
385        Utc::now().naive_utc(),
386    )
387}
388
389/// The component identity of a (pAMM, pair) combination: `pamm ++ token0 ++ token1`.
390fn component_id(pamm: &Bytes, token0: &Bytes, token1: &Bytes) -> Bytes {
391    Bytes::from([pamm.as_ref(), token0.as_ref(), token1.as_ref()].concat())
392}
393
394#[cfg(test)]
395mod tests {
396    use std::str::FromStr;
397
398    use num_bigint::BigUint;
399
400    use super::*;
401
402    const PAMM: &str = "0x5979458912f80b96d30d4220af8e2e4925a33320";
403    const WBTC: &str = "0x2260fac5e5542a773aa44fbcfedf7c193bc2c599";
404    const USDC: &str = "0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48";
405    const WETH: &str = "0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2";
406
407    fn token(address: &str, symbol: &str, decimals: u32) -> Token {
408        Token::new(
409            &Bytes::from_str(address).unwrap(),
410            symbol,
411            decimals,
412            0,
413            &[Some(10_000)],
414            Chain::Ethereum,
415            100,
416        )
417    }
418
419    fn tokens() -> HashMap<Bytes, Token> {
420        [token(WBTC, "WBTC", 8), token(USDC, "USDC", 6), token(WETH, "WETH", 18)]
421            .into_iter()
422            .map(|token| (token.address.clone(), token))
423            .collect()
424    }
425
426    fn tracker() -> SnapshotTracker {
427        let config = PriceLevelStreamConfig::new(
428            "fermiswap",
429            Bytes::from_str(PAMM).unwrap(),
430            BigUint::from(120_000u64),
431        );
432        SnapshotTracker::new(
433            HashMap::from([(config.address.clone(), config)]),
434            HashSet::new(),
435            tokens(),
436            false,
437            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
438        )
439    }
440
441    fn level(amount_in: u64, amount_out: u64) -> TitanPriceLevel {
442        TitanPriceLevel {
443            amount_in: BigUint::from(amount_in),
444            amount_out: BigUint::from(amount_out),
445        }
446    }
447
448    fn pair_levels(
449        token_in: &str,
450        token_out: &str,
451        order_book: Vec<TitanPriceLevel>,
452    ) -> TitanPairLevels {
453        TitanPairLevels {
454            token_in: Bytes::from_str(token_in).unwrap(),
455            token_out: Bytes::from_str(token_out).unwrap(),
456            order_book,
457        }
458    }
459
460    fn message(block_number: u64, pairs: Vec<TitanPairLevels>) -> TitanPriceLevelMessage {
461        TitanPriceLevelMessage {
462            block_number,
463            pamms: vec![TitanPammLevels { pamm: Bytes::from_str(PAMM).unwrap(), pairs }],
464        }
465    }
466
467    fn wbtc_usdc_pairs() -> Vec<TitanPairLevels> {
468        vec![
469            pair_levels(WBTC, USDC, vec![level(100_000_000, 100_000_000_000)]),
470            pair_levels(USDC, WBTC, vec![level(100_000_000_000, 99_000_000)]),
471        ]
472    }
473
474    fn expected_id() -> String {
475        // pamm ++ token0 ++ token1 with WBTC < USDC.
476        format!("{PAMM}{}{}", &WBTC[2..], &USDC[2..])
477    }
478
479    #[test]
480    fn first_snapshot_emits_new_pair_with_both_directions() {
481        let mut tracker = tracker();
482        let Update {
483            block_number_or_timestamp,
484            is_partial,
485            sync_states,
486            states,
487            new_pairs,
488            removed_pairs,
489        } = tracker
490            .process(message(100, wbtc_usdc_pairs()))
491            .expect("update expected");
492
493        assert_eq!(block_number_or_timestamp, 100);
494        assert!(is_partial);
495        assert!(sync_states.is_empty());
496        assert!(removed_pairs.is_empty());
497
498        let id = expected_id();
499        let component = &new_pairs[&id];
500        assert_eq!(component.protocol_system, "pricelevelstream:fermiswap");
501        assert_eq!(
502            component.static_attributes[PAMM_ADDRESS_ATTRIBUTE],
503            Bytes::from_str(PAMM).unwrap()
504        );
505
506        let PriceLevelStreamState { token0, token1, quotes_0_to_1, quotes_1_to_0, gas_cost } =
507            states[&id]
508                .as_any()
509                .downcast_ref::<PriceLevelStreamState>()
510                .expect("price level state");
511        assert_eq!(token0, &Bytes::from_str(WBTC).unwrap());
512        assert_eq!(token1, &Bytes::from_str(USDC).unwrap());
513        assert_eq!(quotes_0_to_1.len(), 1);
514        assert_eq!(quotes_1_to_0.len(), 1);
515        assert_eq!(quotes_0_to_1[0].amount_in, BigUint::from(100_000_000u64));
516        assert_eq!(gas_cost, &BigUint::from(120_000u64));
517    }
518
519    #[test]
520    fn repeated_snapshot_is_not_a_new_pair() {
521        let mut tracker = tracker();
522        tracker
523            .process(message(100, wbtc_usdc_pairs()))
524            .expect("update expected");
525        let update = tracker
526            .process(message(101, wbtc_usdc_pairs()))
527            .expect("update expected");
528
529        assert!(update.new_pairs.is_empty());
530        assert!(update.removed_pairs.is_empty());
531        assert!(update
532            .states
533            .contains_key(&expected_id()));
534    }
535
536    #[test]
537    fn dropped_pair_is_removed() {
538        let mut tracker = tracker();
539        tracker
540            .process(message(100, wbtc_usdc_pairs()))
541            .expect("update expected");
542        let weth_usdc =
543            vec![pair_levels(WETH, USDC, vec![level(1_000_000_000_000_000_000, 3_000_000_000)])];
544        let update = tracker
545            .process(message(101, weth_usdc))
546            .expect("update expected");
547
548        assert_eq!(update.removed_pairs.len(), 1);
549        assert!(update
550            .removed_pairs
551            .contains_key(&expected_id()));
552        assert_eq!(update.new_pairs.len(), 1);
553        assert_eq!(update.states.len(), 1);
554    }
555
556    #[test]
557    fn out_of_order_frame_is_skipped() {
558        let mut tracker = tracker();
559        tracker
560            .process(message(101, wbtc_usdc_pairs()))
561            .expect("update expected");
562
563        // A frame for an older block is stale: no update, and the caches stay untouched even
564        // though the frame's snapshot differs completely.
565        let stale =
566            vec![pair_levels(WETH, USDC, vec![level(1_000_000_000_000_000_000, 3_000_000_000)])];
567        assert!(tracker
568            .process(message(100, stale))
569            .is_none());
570
571        // The next current frame diffs against the pre-stale state: nothing was added or
572        // removed in between.
573        let update = tracker
574            .process(message(102, wbtc_usdc_pairs()))
575            .expect("update expected");
576        assert!(update.new_pairs.is_empty());
577        assert!(update.removed_pairs.is_empty());
578    }
579
580    #[test]
581    fn vanished_pamm_has_its_pairs_removed() {
582        let mut tracker = tracker();
583        tracker
584            .process(message(100, wbtc_usdc_pairs()))
585            .expect("update expected");
586
587        // The next frame no longer contains the pAMM at all: a complete snapshot without a
588        // venue means the venue is gone, pairs and all.
589        let update = tracker
590            .process(TitanPriceLevelMessage { block_number: 101, pamms: vec![] })
591            .expect("update expected");
592        assert!(update.states.is_empty());
593        assert!(update.new_pairs.is_empty());
594        assert_eq!(update.removed_pairs.len(), 1);
595        assert!(update
596            .removed_pairs
597            .contains_key(&expected_id()));
598
599        // Nothing served and nothing changed: no update.
600        assert!(tracker
601            .process(TitanPriceLevelMessage { block_number: 102, pamms: vec![] })
602            .is_none());
603
604        // A venue that reappears is a new pair again.
605        let update = tracker
606            .process(message(103, wbtc_usdc_pairs()))
607            .expect("update expected");
608        assert!(update
609            .new_pairs
610            .contains_key(&expected_id()));
611    }
612
613    #[test]
614    fn unregistered_pamm_produces_no_update_without_auto_detection() {
615        let mut tracker = SnapshotTracker::new(
616            HashMap::new(),
617            HashSet::new(),
618            tokens(),
619            false,
620            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
621        );
622        assert!(tracker
623            .process(message(100, wbtc_usdc_pairs()))
624            .is_none());
625    }
626
627    #[test]
628    fn denied_pamm_is_not_auto_detected() {
629        let denied = HashSet::from([Bytes::from_str(PAMM).unwrap()]);
630        let mut tracker = SnapshotTracker::new(
631            HashMap::new(),
632            denied,
633            tokens(),
634            true,
635            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
636        );
637        assert!(tracker
638            .process(message(100, wbtc_usdc_pairs()))
639            .is_none());
640    }
641
642    #[test]
643    fn explicit_add_and_deny_are_last_wins() {
644        let address = Bytes::from_str(PAMM).unwrap();
645        let custom =
646            || PriceLevelStreamConfig::new("custom", Bytes::from_str(PAMM).unwrap(), 1u64.into());
647
648        let builder = PriceLevelStreamBuilder::new()
649            .add_pamm(custom())
650            .deny_pamm(address.clone());
651        assert!(!builder.registry.contains_key(&address));
652        assert!(builder.denied.contains(&address));
653
654        let builder = PriceLevelStreamBuilder::new()
655            .deny_pamm(address.clone())
656            .add_pamm(custom());
657        assert_eq!(builder.registry[&address].protocol, "custom");
658        assert!(builder.denied.is_empty());
659    }
660
661    #[test]
662    fn defaults_never_override_explicit_calls() {
663        // Denying a venue from the default set works in either call order.
664        let fermiswap_router = Bytes::from_str(PAMM).unwrap();
665        for builder in [
666            PriceLevelStreamBuilder::new()
667                .deny_pamm(fermiswap_router.clone())
668                .with_known_pamms(),
669            PriceLevelStreamBuilder::new()
670                .with_known_pamms()
671                .deny_pamm(fermiswap_router.clone()),
672        ] {
673            assert!(!builder
674                .registry
675                .contains_key(&fermiswap_router));
676            assert!(builder
677                .denied
678                .contains(&fermiswap_router));
679            // The other defaults are unaffected.
680            assert!(!builder.registry.is_empty());
681        }
682
683        // Registering a venue from the default deny set works in either call order.
684        let denied_venue = default_denied_pamms().remove(0);
685        let custom = || PriceLevelStreamConfig::new("custom", denied_venue.clone(), 1u64.into());
686        for builder in [
687            PriceLevelStreamBuilder::new()
688                .add_pamm(custom())
689                .with_known_pamms(),
690            PriceLevelStreamBuilder::new()
691                .with_known_pamms()
692                .add_pamm(custom()),
693        ] {
694            assert_eq!(builder.registry[&denied_venue].protocol, "custom");
695            assert!(!builder.denied.contains(&denied_venue));
696        }
697    }
698
699    #[test]
700    fn auto_detected_pamm_is_served_under_its_address() {
701        let mut tracker = SnapshotTracker::new(
702            HashMap::new(),
703            HashSet::new(),
704            tokens(),
705            true,
706            BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST),
707        );
708        let update = tracker
709            .process(message(100, wbtc_usdc_pairs()))
710            .expect("update expected");
711
712        let component = &update.new_pairs[&expected_id()];
713        assert_eq!(component.protocol_system, format!("pricelevelstream:{PAMM}"));
714        let state = update.states[&expected_id()]
715            .as_any()
716            .downcast_ref::<PriceLevelStreamState>()
717            .expect("price level state");
718        assert_eq!(state.gas_cost, BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST));
719
720        // The synthesized config is cached: the next snapshot is not a new pair again.
721        let update = tracker
722            .process(message(101, wbtc_usdc_pairs()))
723            .expect("update expected");
724        assert!(update.new_pairs.is_empty());
725    }
726
727    #[test]
728    fn auto_detected_gas_cost_override_applies() {
729        let mut tracker = SnapshotTracker::new(
730            HashMap::new(),
731            HashSet::new(),
732            tokens(),
733            true,
734            BigUint::from(42_000u64),
735        );
736        let update = tracker
737            .process(message(100, wbtc_usdc_pairs()))
738            .expect("update expected");
739
740        let state = update.states[&expected_id()]
741            .as_any()
742            .downcast_ref::<PriceLevelStreamState>()
743            .expect("price level state");
744        assert_eq!(state.gas_cost, BigUint::from(42_000u64));
745    }
746
747    #[test]
748    fn with_known_pamms_registers_known_venues() {
749        // PAMM is the FermiSwap router, one of the default venues.
750        let fermiswap_router = Bytes::from_str(PAMM).unwrap();
751
752        let builder = PriceLevelStreamBuilder::new();
753        assert!(builder.registry.is_empty());
754        assert!(builder.denied.is_empty());
755
756        let builder = builder.with_known_pamms();
757        assert_eq!(builder.registry[&fermiswap_router].protocol, "fermiswap");
758        // The known-bad venues get denied alongside, and never overlap the served defaults.
759        assert!(!builder.denied.is_empty());
760        assert!(builder.denied.is_disjoint(
761            &builder
762                .registry
763                .keys()
764                .cloned()
765                .collect()
766        ));
767
768        // An `add_pamm` entry wins over the default for the same address, in either call order.
769        let custom =
770            || PriceLevelStreamConfig::new("custom", fermiswap_router.clone(), BigUint::from(1u64));
771        for builder in [
772            PriceLevelStreamBuilder::new()
773                .add_pamm(custom())
774                .with_known_pamms(),
775            PriceLevelStreamBuilder::new()
776                .with_known_pamms()
777                .add_pamm(custom()),
778        ] {
779            assert_eq!(builder.registry[&fermiswap_router].protocol, "custom");
780            assert_eq!(builder.registry[&fermiswap_router].gas_cost, BigUint::from(1u64));
781        }
782    }
783
784    #[test]
785    fn unknown_tokens_are_skipped() {
786        let mut tracker = tracker();
787        let unknown = vec![pair_levels(
788            "0x1111111111111111111111111111111111111111",
789            USDC,
790            vec![level(1, 1)],
791        )];
792        assert!(tracker
793            .process(message(100, unknown))
794            .is_none());
795    }
796}
tycho_simulation/price_level_stream/stream.rs

tycho_simulation/price_level_stream/
stream.rs
