nautilus_blockchain/data/
subscription.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use ahash::{AHashMap, AHashSet};
17use alloy::primitives::{Address, keccak256};
18use nautilus_core::hex;
19use nautilus_model::defi::DexType;
20
21/// Manages subscriptions to DeFi protocol events (swaps, mints, burns, collects) across different DEXs.
22///
23/// This manager counts the owners of each pool address per event type, so overlapping
24/// subscriptions (a complete pool and a narrower event subscription) release independently.
25/// It also maintains the event signature encodings for efficient filtering.
26#[derive(Debug, Default)]
27pub struct DefiDataSubscriptionManager {
28    block_demand_explicit: bool,
29    block_demand_pool_events: bool,
30    block_feed_backend: Option<BlockFeedBackend>,
31    pool_swap_event_encoded: AHashMap<DexType, String>,
32    pool_mint_event_encoded: AHashMap<DexType, String>,
33    pool_burn_event_encoded: AHashMap<DexType, String>,
34    pool_collect_event_encoded: AHashMap<DexType, String>,
35    pool_flash_event_encoded: AHashMap<DexType, String>,
36    pool_fee_protocol_update_event_encoded: AHashMap<DexType, String>,
37    pool_fee_protocol_collect_event_encoded: AHashMap<DexType, String>,
38    subscribed_pool_swaps: AHashMap<DexType, PoolEventDemand>,
39    subscribed_pool_mints: AHashMap<DexType, PoolEventDemand>,
40    subscribed_pool_burns: AHashMap<DexType, PoolEventDemand>,
41    subscribed_pool_collects: AHashMap<DexType, PoolEventDemand>,
42    subscribed_pool_flashes: AHashMap<DexType, PoolEventDemand>,
43    subscribed_pool_fee_protocol_updates: AHashMap<DexType, PoolEventDemand>,
44    subscribed_pool_fee_protocol_collects: AHashMap<DexType, PoolEventDemand>,
45}
46
47impl DefiDataSubscriptionManager {
48    /// Creates a new [`DefiDataSubscriptionManager`] instance.
49    #[must_use]
50    pub fn new() -> Self {
51        Self::default()
52    }
53
54    pub(crate) fn add_block_demand(
55        &mut self,
56        owner: BlockFeedOwner,
57        preferred_backend: BlockFeedBackend,
58    ) -> Option<BlockFeedBackend> {
59        match owner {
60            BlockFeedOwner::Explicit => self.block_demand_explicit = true,
61            BlockFeedOwner::PoolEvents => self.block_demand_pool_events = true,
62        }
63
64        self.block_feed_backend
65            .is_none()
66            .then_some(preferred_backend)
67    }
68
69    pub(crate) fn block_feed_started(&mut self, backend: BlockFeedBackend) {
70        self.block_feed_backend = Some(backend);
71    }
72
73    pub(crate) fn remove_block_demand(
74        &mut self,
75        owner: BlockFeedOwner,
76    ) -> Option<BlockFeedBackend> {
77        match owner {
78            BlockFeedOwner::Explicit => self.block_demand_explicit = false,
79            BlockFeedOwner::PoolEvents => self.block_demand_pool_events = false,
80        }
81
82        if self.has_block_demand() {
83            None
84        } else {
85            self.block_feed_backend
86        }
87    }
88
89    pub(crate) fn block_feed_stopped(&mut self, backend: BlockFeedBackend) {
90        if self.block_feed_backend == Some(backend) {
91            self.block_feed_backend = None;
92        }
93    }
94
95    pub(crate) fn clear_block_demand(&mut self) {
96        self.block_demand_explicit = false;
97        self.block_demand_pool_events = false;
98        self.block_feed_backend = None;
99    }
100
101    pub(crate) fn has_pool_event_subscriptions(&self) -> bool {
102        self.subscribed_pool_swaps
103            .values()
104            .chain(self.subscribed_pool_mints.values())
105            .chain(self.subscribed_pool_burns.values())
106            .chain(self.subscribed_pool_collects.values())
107            .chain(self.subscribed_pool_flashes.values())
108            .chain(self.subscribed_pool_fee_protocol_updates.values())
109            .chain(self.subscribed_pool_fee_protocol_collects.values())
110            .any(|addresses| !addresses.is_empty())
111    }
112
113    fn has_block_demand(&self) -> bool {
114        self.block_demand_explicit || self.block_demand_pool_events
115    }
116
117    /// Gets all unique contract addresses subscribed for any event type for a given DEX.
118    #[must_use]
119    pub fn get_subscribed_dex_contract_addresses(&self, dex: &DexType) -> Vec<Address> {
120        let mut unique_addresses = AHashSet::new();
121
122        if let Some(addresses) = self.subscribed_pool_swaps.get(dex) {
123            unique_addresses.extend(addresses.keys().copied());
124        }
125
126        if let Some(addresses) = self.subscribed_pool_mints.get(dex) {
127            unique_addresses.extend(addresses.keys().copied());
128        }
129
130        if let Some(addresses) = self.subscribed_pool_burns.get(dex) {
131            unique_addresses.extend(addresses.keys().copied());
132        }
133
134        if let Some(addresses) = self.subscribed_pool_collects.get(dex) {
135            unique_addresses.extend(addresses.keys().copied());
136        }
137
138        if let Some(addresses) = self.subscribed_pool_flashes.get(dex) {
139            unique_addresses.extend(addresses.keys().copied());
140        }
141
142        if let Some(addresses) = self.subscribed_pool_fee_protocol_updates.get(dex) {
143            unique_addresses.extend(addresses.keys().copied());
144        }
145
146        if let Some(addresses) = self.subscribed_pool_fee_protocol_collects.get(dex) {
147            unique_addresses.extend(addresses.keys().copied());
148        }
149
150        unique_addresses.into_iter().collect()
151    }
152
153    /// Gets all event signatures (keccak256 hashes) registered for a given DEX.
154    #[must_use]
155    pub fn get_subscribed_dex_event_signatures(&self, dex: &DexType) -> Vec<String> {
156        let mut result = Vec::new();
157
158        if let Some(swap_event_signature) = self.pool_swap_event_encoded.get(dex) {
159            result.push(swap_event_signature.clone());
160        }
161
162        if let Some(mint_event_signature) = self.pool_mint_event_encoded.get(dex) {
163            result.push(mint_event_signature.clone());
164        }
165
166        if let Some(burn_event_signature) = self.pool_burn_event_encoded.get(dex) {
167            result.push(burn_event_signature.clone());
168        }
169
170        if let Some(collect_event_signature) = self.pool_collect_event_encoded.get(dex) {
171            result.push(collect_event_signature.clone());
172        }
173
174        if let Some(flash_event_signature) = self.pool_flash_event_encoded.get(dex) {
175            result.push(flash_event_signature.clone());
176        }
177
178        if let Some(fee_protocol_update_event_signature) =
179            self.pool_fee_protocol_update_event_encoded.get(dex)
180        {
181            result.push(fee_protocol_update_event_signature.clone());
182        }
183
184        if let Some(fee_protocol_collect_event_signature) =
185            self.pool_fee_protocol_collect_event_encoded.get(dex)
186        {
187            result.push(fee_protocol_collect_event_signature.clone());
188        }
189
190        result
191    }
192
193    /// Gets event signatures with at least one subscribed pool address for a given DEX.
194    #[must_use]
195    pub fn get_active_subscribed_dex_event_signatures(&self, dex: &DexType) -> Vec<String> {
196        let mut result = Vec::new();
197
198        if self
199            .subscribed_pool_swaps
200            .get(dex)
201            .is_some_and(|addresses| !addresses.is_empty())
202            && let Some(signature) = self.pool_swap_event_encoded.get(dex)
203        {
204            result.push(signature.clone());
205        }
206
207        if self
208            .subscribed_pool_mints
209            .get(dex)
210            .is_some_and(|addresses| !addresses.is_empty())
211            && let Some(signature) = self.pool_mint_event_encoded.get(dex)
212        {
213            result.push(signature.clone());
214        }
215
216        if self
217            .subscribed_pool_burns
218            .get(dex)
219            .is_some_and(|addresses| !addresses.is_empty())
220            && let Some(signature) = self.pool_burn_event_encoded.get(dex)
221        {
222            result.push(signature.clone());
223        }
224
225        if self
226            .subscribed_pool_collects
227            .get(dex)
228            .is_some_and(|addresses| !addresses.is_empty())
229            && let Some(signature) = self.pool_collect_event_encoded.get(dex)
230        {
231            result.push(signature.clone());
232        }
233
234        if self
235            .subscribed_pool_flashes
236            .get(dex)
237            .is_some_and(|addresses| !addresses.is_empty())
238            && let Some(signature) = self.pool_flash_event_encoded.get(dex)
239        {
240            result.push(signature.clone());
241        }
242
243        if self
244            .subscribed_pool_fee_protocol_updates
245            .get(dex)
246            .is_some_and(|addresses| !addresses.is_empty())
247            && let Some(signature) = self.pool_fee_protocol_update_event_encoded.get(dex)
248        {
249            result.push(signature.clone());
250        }
251
252        if self
253            .subscribed_pool_fee_protocol_collects
254            .get(dex)
255            .is_some_and(|addresses| !addresses.is_empty())
256            && let Some(signature) = self.pool_fee_protocol_collect_event_encoded.get(dex)
257        {
258            result.push(signature.clone());
259        }
260
261        result
262    }
263
264    /// Gets pool addresses subscribed to swap events for a given DEX.
265    #[must_use]
266    pub fn get_subscribed_pool_swap_addresses(&self, dex: &DexType) -> Vec<Address> {
267        self.subscribed_pool_swaps
268            .get(dex)
269            .map(|addresses| addresses.keys().copied().collect())
270            .unwrap_or_default()
271    }
272
273    /// Gets pool addresses subscribed to mint events for a given DEX.
274    #[must_use]
275    pub fn get_subscribed_pool_mint_addresses(&self, dex: &DexType) -> Vec<Address> {
276        self.subscribed_pool_mints
277            .get(dex)
278            .map(|addresses| addresses.keys().copied().collect())
279            .unwrap_or_default()
280    }
281
282    /// Gets pool addresses subscribed to burn events for a given DEX.
283    #[must_use]
284    pub fn get_subscribed_pool_burn_addresses(&self, dex: &DexType) -> Vec<Address> {
285        self.subscribed_pool_burns
286            .get(dex)
287            .map(|addresses| addresses.keys().copied().collect())
288            .unwrap_or_default()
289    }
290
291    /// Gets pool addresses subscribed to collect events for a given DEX.
292    #[must_use]
293    pub fn get_subscribed_pool_collect_addresses(&self, dex: &DexType) -> Vec<Address> {
294        self.subscribed_pool_collects
295            .get(dex)
296            .map(|addresses| addresses.keys().copied().collect())
297            .unwrap_or_default()
298    }
299
300    /// Gets pool addresses subscribed to flash events for a given DEX.
301    #[must_use]
302    pub fn get_subscribed_pool_flash_addresses(&self, dex: &DexType) -> Vec<Address> {
303        self.subscribed_pool_flashes
304            .get(dex)
305            .map(|addresses| addresses.keys().copied().collect())
306            .unwrap_or_default()
307    }
308
309    /// Gets pool addresses subscribed to fee-protocol update events for a given DEX.
310    #[must_use]
311    pub fn get_subscribed_pool_fee_protocol_update_addresses(&self, dex: &DexType) -> Vec<Address> {
312        self.subscribed_pool_fee_protocol_updates
313            .get(dex)
314            .map(|addresses| addresses.keys().copied().collect())
315            .unwrap_or_default()
316    }
317
318    /// Gets pool addresses subscribed to fee-protocol collect events for a given DEX.
319    #[must_use]
320    pub fn get_subscribed_pool_fee_protocol_collect_addresses(
321        &self,
322        dex: &DexType,
323    ) -> Vec<Address> {
324        self.subscribed_pool_fee_protocol_collects
325            .get(dex)
326            .map(|addresses| addresses.keys().copied().collect())
327            .unwrap_or_default()
328    }
329
330    /// Gets the swap event signature for a specific DEX.
331    #[must_use]
332    pub fn get_dex_pool_swap_event_signature(&self, dex: &DexType) -> Option<String> {
333        self.pool_swap_event_encoded.get(dex).cloned()
334    }
335
336    /// Gets the mint event signature for a specific DEX.
337    #[must_use]
338    pub fn get_dex_pool_mint_event_signature(&self, dex: &DexType) -> Option<String> {
339        self.pool_mint_event_encoded.get(dex).cloned()
340    }
341    /// Gets the burn event signature for a specific DEX.
342    #[must_use]
343    pub fn get_dex_pool_burn_event_signature(&self, dex: &DexType) -> Option<String> {
344        self.pool_burn_event_encoded.get(dex).cloned()
345    }
346
347    /// Gets the collect event signature for a specific DEX.
348    #[must_use]
349    pub fn get_dex_pool_collect_event_signature(&self, dex: &DexType) -> Option<String> {
350        self.pool_collect_event_encoded.get(dex).cloned()
351    }
352
353    /// Gets the flash event signature for a specific DEX.
354    #[must_use]
355    pub fn get_dex_pool_flash_event_signature(&self, dex: &DexType) -> Option<String> {
356        self.pool_flash_event_encoded.get(dex).cloned()
357    }
358
359    /// Gets the fee-protocol update event signature for a specific DEX.
360    #[must_use]
361    pub fn get_dex_pool_fee_protocol_update_event_signature(
362        &self,
363        dex: &DexType,
364    ) -> Option<String> {
365        self.pool_fee_protocol_update_event_encoded
366            .get(dex)
367            .cloned()
368    }
369
370    /// Gets the fee-protocol collect event signature for a specific DEX.
371    #[must_use]
372    pub fn get_dex_pool_fee_protocol_collect_event_signature(
373        &self,
374        dex: &DexType,
375    ) -> Option<String> {
376        self.pool_fee_protocol_collect_event_encoded
377            .get(dex)
378            .cloned()
379    }
380
381    /// Normalizes an event signature to a consistent format.
382    ///
383    /// Accepts:
384    /// - A raw event signature like "Swap(address,address,int256,int256,uint160,uint128,int24)".
385    /// - A pre-encoded topic like "0xc42079f94a6350d7e6235f29174924f928cc2ac818eb64fed8004e115fbcca67".
386    /// - A hex string without 0x prefix.
387    ///
388    /// Returns a normalized "0x..." format string.
389    fn normalize_topic(sig: &str) -> String {
390        let s = sig.trim();
391
392        // Check if it's already a properly formatted hex string with 0x prefix
393        if let Some(rest) = s.strip_prefix("0x")
394            && rest.len() == 64
395            && rest.chars().all(|c| c.is_ascii_hexdigit())
396        {
397            return format!("0x{}", rest.to_ascii_lowercase());
398        }
399
400        // Check if it's a hex string without 0x prefix
401        if s.len() == 64 && s.chars().all(|c| c.is_ascii_hexdigit()) {
402            return format!("0x{}", s.to_ascii_lowercase());
403        }
404
405        // Otherwise, it's a raw signature that needs hashing
406        hex::encode_prefixed(keccak256(s.as_bytes()))
407    }
408
409    /// Registers a DEX with its event signatures for subscription management.
410    ///
411    /// This must be called before subscribing to any events for a DEX.
412    /// Event signatures can be either raw signatures or pre-encoded keccak256 hashes.
413    pub fn register_dex_for_subscriptions(
414        &mut self,
415        dex: DexType,
416        swap_event_signature: &str,
417        mint_event_signature: &str,
418        burn_event_signature: &str,
419        collect_event_signature: &str,
420        flash_event_signature: Option<&str>,
421    ) {
422        self.subscribed_pool_swaps
423            .insert(dex, PoolEventDemand::new());
424        self.pool_swap_event_encoded
425            .insert(dex, Self::normalize_topic(swap_event_signature));
426
427        self.subscribed_pool_mints
428            .insert(dex, PoolEventDemand::new());
429        self.pool_mint_event_encoded
430            .insert(dex, Self::normalize_topic(mint_event_signature));
431
432        self.subscribed_pool_burns
433            .insert(dex, PoolEventDemand::new());
434        self.pool_burn_event_encoded
435            .insert(dex, Self::normalize_topic(burn_event_signature));
436
437        self.subscribed_pool_collects
438            .insert(dex, PoolEventDemand::new());
439        self.pool_collect_event_encoded
440            .insert(dex, Self::normalize_topic(collect_event_signature));
441
442        if let Some(flash_event_signature) = flash_event_signature {
443            self.subscribed_pool_flashes
444                .insert(dex, PoolEventDemand::new());
445            self.pool_flash_event_encoded
446                .insert(dex, Self::normalize_topic(flash_event_signature));
447        }
448
449        log::debug!("Registered DEX for subscriptions: {dex:?}");
450    }
451
452    /// Registers optional fee-protocol event signatures for subscription management.
453    pub fn register_dex_fee_protocol_events(
454        &mut self,
455        dex: DexType,
456        fee_protocol_update_event_signature: Option<&str>,
457        fee_protocol_collect_event_signature: Option<&str>,
458    ) {
459        if let Some(fee_protocol_update_event_signature) = fee_protocol_update_event_signature {
460            self.subscribed_pool_fee_protocol_updates
461                .insert(dex, PoolEventDemand::new());
462            self.pool_fee_protocol_update_event_encoded.insert(
463                dex,
464                Self::normalize_topic(fee_protocol_update_event_signature),
465            );
466        }
467
468        if let Some(fee_protocol_collect_event_signature) = fee_protocol_collect_event_signature {
469            self.subscribed_pool_fee_protocol_collects
470                .insert(dex, PoolEventDemand::new());
471            self.pool_fee_protocol_collect_event_encoded.insert(
472                dex,
473                Self::normalize_topic(fee_protocol_collect_event_signature),
474            );
475        }
476    }
477
478    /// Subscribes to every event type of a specific pool address on a DEX.
479    pub fn subscribe_pool(&mut self, dex: DexType, address: Address) {
480        self.subscribe_swaps(dex, address);
481        self.subscribe_mints(dex, address);
482        self.subscribe_burns(dex, address);
483        self.subscribe_collects(dex, address);
484        self.subscribe_flashes(dex, address);
485        self.subscribe_fee_protocol_updates(dex, address);
486        self.subscribe_fee_protocol_collects(dex, address);
487    }
488
489    /// Unsubscribes from every event type of a specific pool address on a DEX.
490    pub fn unsubscribe_pool(&mut self, dex: DexType, address: Address) {
491        self.unsubscribe_swaps(dex, address);
492        self.unsubscribe_mints(dex, address);
493        self.unsubscribe_burns(dex, address);
494        self.unsubscribe_collects(dex, address);
495        self.unsubscribe_flashes(dex, address);
496        self.unsubscribe_fee_protocol_updates(dex, address);
497        self.unsubscribe_fee_protocol_collects(dex, address);
498    }
499
500    /// Subscribes to swap events for a specific pool address on a DEX.
501    pub fn subscribe_swaps(&mut self, dex: DexType, address: Address) {
502        if let Some(pool_set) = self.subscribed_pool_swaps.get_mut(&dex) {
503            acquire_pool_event(pool_set, address);
504        } else {
505            log::error!("DEX not registered for swap subscriptions: {dex:?}");
506        }
507    }
508
509    /// Subscribes to mint events for a specific pool address on a DEX.
510    pub fn subscribe_mints(&mut self, dex: DexType, address: Address) {
511        if let Some(pool_set) = self.subscribed_pool_mints.get_mut(&dex) {
512            acquire_pool_event(pool_set, address);
513        } else {
514            log::error!("DEX not registered for mint subscriptions: {dex:?}");
515        }
516    }
517
518    /// Subscribes to burn events for a specific pool address on a DEX.
519    pub fn subscribe_burns(&mut self, dex: DexType, address: Address) {
520        if let Some(pool_set) = self.subscribed_pool_burns.get_mut(&dex) {
521            acquire_pool_event(pool_set, address);
522        } else {
523            log::warn!("DEX not registered for burn subscriptions: {dex:?}");
524        }
525    }
526
527    /// Unsubscribes from swap events for a specific pool address on a DEX.
528    pub fn unsubscribe_swaps(&mut self, dex: DexType, address: Address) {
529        if let Some(pool_set) = self.subscribed_pool_swaps.get_mut(&dex) {
530            release_pool_event(pool_set, address);
531        } else {
532            log::error!("DEX not registered for swap subscriptions: {dex:?}");
533        }
534    }
535
536    /// Unsubscribes from mint events for a specific pool address on a DEX.
537    pub fn unsubscribe_mints(&mut self, dex: DexType, address: Address) {
538        if let Some(pool_set) = self.subscribed_pool_mints.get_mut(&dex) {
539            release_pool_event(pool_set, address);
540        } else {
541            log::error!("DEX not registered for mint subscriptions: {dex:?}");
542        }
543    }
544
545    /// Unsubscribes from burn events for a specific pool address on a DEX.
546    pub fn unsubscribe_burns(&mut self, dex: DexType, address: Address) {
547        if let Some(pool_set) = self.subscribed_pool_burns.get_mut(&dex) {
548            release_pool_event(pool_set, address);
549        } else {
550            log::error!("DEX not registered for burn subscriptions: {dex:?}");
551        }
552    }
553
554    /// Subscribes to collect events for a specific pool address on a DEX.
555    pub fn subscribe_collects(&mut self, dex: DexType, address: Address) {
556        if let Some(pool_set) = self.subscribed_pool_collects.get_mut(&dex) {
557            acquire_pool_event(pool_set, address);
558        } else {
559            log::error!("DEX not registered for collect subscriptions: {dex:?}");
560        }
561    }
562
563    /// Unsubscribes from collect events for a specific pool address on a DEX.
564    pub fn unsubscribe_collects(&mut self, dex: DexType, address: Address) {
565        if let Some(pool_set) = self.subscribed_pool_collects.get_mut(&dex) {
566            release_pool_event(pool_set, address);
567        } else {
568            log::error!("DEX not registered for collect subscriptions: {dex:?}");
569        }
570    }
571
572    /// Subscribes to flash events for a specific pool address on a DEX.
573    pub fn subscribe_flashes(&mut self, dex: DexType, address: Address) {
574        if let Some(pool_set) = self.subscribed_pool_flashes.get_mut(&dex) {
575            acquire_pool_event(pool_set, address);
576        } else {
577            log::error!("DEX not registered for flash subscriptions: {dex:?}");
578        }
579    }
580
581    /// Unsubscribes from flash events for a specific pool address on a DEX.
582    pub fn unsubscribe_flashes(&mut self, dex: DexType, address: Address) {
583        if let Some(pool_set) = self.subscribed_pool_flashes.get_mut(&dex) {
584            release_pool_event(pool_set, address);
585        } else {
586            log::error!("DEX not registered for flash subscriptions: {dex:?}");
587        }
588    }
589
590    /// Subscribes to fee-protocol update events for a specific pool address on a DEX.
591    pub fn subscribe_fee_protocol_updates(&mut self, dex: DexType, address: Address) {
592        if let Some(pool_set) = self.subscribed_pool_fee_protocol_updates.get_mut(&dex) {
593            acquire_pool_event(pool_set, address);
594        }
595    }
596
597    /// Unsubscribes from fee-protocol update events for a specific pool address on a DEX.
598    pub fn unsubscribe_fee_protocol_updates(&mut self, dex: DexType, address: Address) {
599        if let Some(pool_set) = self.subscribed_pool_fee_protocol_updates.get_mut(&dex) {
600            release_pool_event(pool_set, address);
601        }
602    }
603
604    /// Subscribes to fee-protocol collect events for a specific pool address on a DEX.
605    pub fn subscribe_fee_protocol_collects(&mut self, dex: DexType, address: Address) {
606        if let Some(pool_set) = self.subscribed_pool_fee_protocol_collects.get_mut(&dex) {
607            acquire_pool_event(pool_set, address);
608        }
609    }
610
611    /// Unsubscribes from fee-protocol collect events for a specific pool address on a DEX.
612    pub fn unsubscribe_fee_protocol_collects(&mut self, dex: DexType, address: Address) {
613        if let Some(pool_set) = self.subscribed_pool_fee_protocol_collects.get_mut(&dex) {
614            release_pool_event(pool_set, address);
615        }
616    }
617}
618
619/// Owner count per pool address for one event kind; an address is present only while owned.
620type PoolEventDemand = AHashMap<Address, usize>;
621
622fn acquire_pool_event(demand: &mut PoolEventDemand, address: Address) {
623    *demand.entry(address).or_default() += 1;
624}
625
626fn release_pool_event(demand: &mut PoolEventDemand, address: Address) {
627    if let Some(owners) = demand.get_mut(&address) {
628        *owners -= 1;
629        if *owners == 0 {
630            demand.remove(&address);
631        }
632    }
633}
634
635#[derive(Debug, Clone, Copy, PartialEq, Eq)]
636pub(crate) enum BlockFeedBackend {
637    Rpc,
638    HyperSync,
639}
640
641#[derive(Debug, Clone, Copy, PartialEq, Eq)]
642pub(crate) enum BlockFeedOwner {
643    Explicit,
644    PoolEvents,
645}
646
647#[cfg(test)]
648mod tests {
649    use alloy::primitives::address;
650    use rstest::{fixture, rstest};
651
652    use super::*;
653
654    #[fixture]
655    fn manager() -> DefiDataSubscriptionManager {
656        DefiDataSubscriptionManager::new()
657    }
658
659    #[fixture]
660    fn registered_manager() -> DefiDataSubscriptionManager {
661        let mut manager = DefiDataSubscriptionManager::new();
662        manager.register_dex_for_subscriptions(
663            DexType::UniswapV3,
664            "Swap(address,address,int256,int256,uint160,uint128,int24)",
665            "Mint(address,address,int24,int24,uint128,uint256,uint256)",
666            "Burn(address,int24,int24,uint128,uint256,uint256)",
667            "Collect(address,address,int24,int24,uint128,uint128)",
668            Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
669        );
670        manager
671    }
672
673    #[rstest]
674    #[case(BlockFeedBackend::Rpc)]
675    #[case(BlockFeedBackend::HyperSync)]
676    fn block_feed_pool_only_stops_with_final_owner(#[case] backend: BlockFeedBackend) {
677        let mut manager = DefiDataSubscriptionManager::new();
678
679        assert_eq!(
680            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
681            Some(backend)
682        );
683        manager.block_feed_started(backend);
684
685        assert_eq!(
686            manager.remove_block_demand(BlockFeedOwner::PoolEvents),
687            Some(backend)
688        );
689        manager.block_feed_stopped(backend);
690
691        assert_eq!(
692            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
693            Some(backend)
694        );
695    }
696
697    #[rstest]
698    #[case(BlockFeedBackend::Rpc)]
699    #[case(BlockFeedBackend::HyperSync)]
700    fn block_feed_repeated_commands_are_idempotent(#[case] backend: BlockFeedBackend) {
701        let mut manager = DefiDataSubscriptionManager::new();
702
703        assert_eq!(
704            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
705            Some(backend)
706        );
707        manager.block_feed_started(backend);
708        assert_eq!(
709            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
710            None
711        );
712
713        assert_eq!(
714            manager.remove_block_demand(BlockFeedOwner::Explicit),
715            Some(backend)
716        );
717        manager.block_feed_stopped(backend);
718        assert_eq!(manager.remove_block_demand(BlockFeedOwner::Explicit), None);
719    }
720
721    #[rstest]
722    #[case(BlockFeedBackend::Rpc)]
723    #[case(BlockFeedBackend::HyperSync)]
724    fn block_feed_explicit_demand_survives_pool_unsubscribe(#[case] backend: BlockFeedBackend) {
725        let mut manager = DefiDataSubscriptionManager::new();
726
727        assert_eq!(
728            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
729            Some(backend)
730        );
731        manager.block_feed_started(backend);
732        assert_eq!(
733            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
734            None
735        );
736
737        assert_eq!(
738            manager.remove_block_demand(BlockFeedOwner::PoolEvents),
739            None
740        );
741        assert_eq!(
742            manager.remove_block_demand(BlockFeedOwner::Explicit),
743            Some(backend)
744        );
745    }
746
747    #[rstest]
748    #[case(BlockFeedBackend::Rpc)]
749    #[case(BlockFeedBackend::HyperSync)]
750    fn block_feed_pool_demand_survives_explicit_unsubscribe(#[case] backend: BlockFeedBackend) {
751        let mut manager = DefiDataSubscriptionManager::new();
752
753        assert_eq!(
754            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
755            Some(backend)
756        );
757        manager.block_feed_started(backend);
758        assert_eq!(
759            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
760            None
761        );
762
763        assert_eq!(manager.remove_block_demand(BlockFeedOwner::Explicit), None);
764        assert_eq!(
765            manager.remove_block_demand(BlockFeedOwner::PoolEvents),
766            Some(backend)
767        );
768    }
769
770    #[rstest]
771    #[case(BlockFeedBackend::Rpc)]
772    #[case(BlockFeedBackend::HyperSync)]
773    fn block_feed_pool_demand_spans_dexes_addresses_and_events(#[case] backend: BlockFeedBackend) {
774        let mut manager = DefiDataSubscriptionManager::new();
775        let pool1 = address!("1111111111111111111111111111111111111111");
776        let pool2 = address!("2222222222222222222222222222222222222222");
777        let pool3 = address!("3333333333333333333333333333333333333333");
778
779        for dex in [DexType::UniswapV3, DexType::PancakeSwapV3] {
780            manager.register_dex_for_subscriptions(
781                dex,
782                "Swap(address,address,int256,int256,uint160,uint128,int24)",
783                "Mint(address,address,int24,int24,uint128,uint256,uint256)",
784                "Burn(address,int24,int24,uint128,uint256,uint256)",
785                "Collect(address,address,int24,int24,uint128,uint128)",
786                Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
787            );
788        }
789
790        manager.subscribe_swaps(DexType::UniswapV3, pool1);
791        manager.subscribe_mints(DexType::UniswapV3, pool2);
792        manager.subscribe_flashes(DexType::PancakeSwapV3, pool3);
793
794        assert!(manager.has_pool_event_subscriptions());
795        assert_eq!(
796            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
797            Some(backend)
798        );
799        manager.block_feed_started(backend);
800
801        manager.unsubscribe_swaps(DexType::UniswapV3, pool1);
802        assert!(manager.has_pool_event_subscriptions());
803        assert_eq!(
804            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
805            None
806        );
807
808        manager.unsubscribe_mints(DexType::UniswapV3, pool2);
809        assert!(manager.has_pool_event_subscriptions());
810        assert_eq!(
811            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
812            None
813        );
814
815        manager.unsubscribe_flashes(DexType::PancakeSwapV3, pool3);
816        assert!(!manager.has_pool_event_subscriptions());
817        assert_eq!(
818            manager.remove_block_demand(BlockFeedOwner::PoolEvents),
819            Some(backend)
820        );
821    }
822
823    #[rstest]
824    #[case(BlockFeedBackend::Rpc)]
825    #[case(BlockFeedBackend::HyperSync)]
826    fn block_feed_disconnect_clears_owners_and_backend(#[case] backend: BlockFeedBackend) {
827        let mut manager = DefiDataSubscriptionManager::new();
828
829        assert_eq!(
830            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
831            Some(backend)
832        );
833        manager.block_feed_started(backend);
834        assert_eq!(
835            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
836            None
837        );
838
839        manager.clear_block_demand();
840
841        assert!(!manager.block_demand_explicit);
842        assert!(!manager.block_demand_pool_events);
843        assert_eq!(manager.block_feed_backend, None);
844        assert_eq!(
845            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
846            Some(backend)
847        );
848    }
849
850    #[rstest]
851    fn block_feed_rpc_fallback_stops_hypersync_backend() {
852        let mut manager = DefiDataSubscriptionManager::new();
853
854        assert_eq!(
855            manager.add_block_demand(BlockFeedOwner::Explicit, BlockFeedBackend::Rpc),
856            Some(BlockFeedBackend::Rpc)
857        );
858        manager.block_feed_started(BlockFeedBackend::HyperSync);
859
860        assert_eq!(
861            manager.remove_block_demand(BlockFeedOwner::Explicit),
862            Some(BlockFeedBackend::HyperSync)
863        );
864    }
865
866    #[rstest]
867    fn test_new_creates_empty_manager(manager: DefiDataSubscriptionManager) {
868        assert_eq!(
869            manager
870                .get_subscribed_dex_contract_addresses(&DexType::UniswapV3)
871                .len(),
872            0
873        );
874        assert_eq!(
875            manager
876                .get_subscribed_dex_event_signatures(&DexType::UniswapV3)
877                .len(),
878            0
879        );
880        assert!(
881            manager
882                .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
883                .is_none()
884        );
885        assert!(
886            manager
887                .get_dex_pool_mint_event_signature(&DexType::UniswapV3)
888                .is_none()
889        );
890        assert!(
891            manager
892                .get_dex_pool_burn_event_signature(&DexType::UniswapV3)
893                .is_none()
894        );
895    }
896
897    #[rstest]
898    fn test_register_dex_for_subscriptions(registered_manager: DefiDataSubscriptionManager) {
899        // Should have all four event signatures
900        let signatures =
901            registered_manager.get_subscribed_dex_event_signatures(&DexType::UniswapV3);
902        assert_eq!(signatures.len(), 5);
903
904        // Each signature should be properly encoded
905        assert!(
906            registered_manager
907                .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
908                .is_some()
909        );
910        assert!(
911            registered_manager
912                .get_dex_pool_mint_event_signature(&DexType::UniswapV3)
913                .is_some()
914        );
915        assert!(
916            registered_manager
917                .get_dex_pool_burn_event_signature(&DexType::UniswapV3)
918                .is_some()
919        );
920    }
921
922    #[rstest]
923    fn test_subscribe_and_get_addresses(mut registered_manager: DefiDataSubscriptionManager) {
924        let pool_address = address!("1234567890123456789012345678901234567890");
925
926        // Subscribe to swap events
927        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
928
929        let addresses =
930            registered_manager.get_subscribed_dex_contract_addresses(&DexType::UniswapV3);
931        assert_eq!(addresses.len(), 1);
932        assert_eq!(addresses[0], pool_address);
933    }
934
935    #[rstest]
936    fn test_subscribe_to_unregistered_dex(mut manager: DefiDataSubscriptionManager) {
937        let pool_address = address!("1234567890123456789012345678901234567890");
938
939        // Try to subscribe without registering - should log warning but not panic
940        manager.subscribe_swaps(DexType::UniswapV3, pool_address);
941        manager.subscribe_mints(DexType::UniswapV3, pool_address);
942        manager.subscribe_burns(DexType::UniswapV3, pool_address);
943
944        // Should return empty results
945        let addresses = manager.get_subscribed_dex_contract_addresses(&DexType::UniswapV3);
946        assert_eq!(addresses.len(), 0);
947    }
948
949    #[rstest]
950    fn test_unsubscribe_removes_address(mut registered_manager: DefiDataSubscriptionManager) {
951        let pool_address = address!("1234567890123456789012345678901234567890");
952
953        // Subscribe
954        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
955
956        // Verify subscription
957        assert_eq!(
958            registered_manager
959                .get_subscribed_dex_contract_addresses(&DexType::UniswapV3)
960                .len(),
961            1
962        );
963
964        // Unsubscribe
965        registered_manager.unsubscribe_swaps(DexType::UniswapV3, pool_address);
966
967        // Verify removal
968        assert_eq!(
969            registered_manager
970                .get_subscribed_dex_contract_addresses(&DexType::UniswapV3)
971                .len(),
972            0
973        );
974    }
975
976    #[rstest]
977    fn test_get_event_signatures(registered_manager: DefiDataSubscriptionManager) {
978        let swap_sig = registered_manager.get_dex_pool_swap_event_signature(&DexType::UniswapV3);
979        let mint_sig = registered_manager.get_dex_pool_mint_event_signature(&DexType::UniswapV3);
980        let burn_sig = registered_manager.get_dex_pool_burn_event_signature(&DexType::UniswapV3);
981
982        // All should be Some and start with 0x
983        assert!(swap_sig.is_some() && swap_sig.unwrap().starts_with("0x"));
984        assert!(mint_sig.is_some() && mint_sig.unwrap().starts_with("0x"));
985        assert!(burn_sig.is_some() && burn_sig.unwrap().starts_with("0x"));
986    }
987
988    #[rstest]
989    fn test_active_event_signatures_only_include_subscribed_event_types(
990        mut registered_manager: DefiDataSubscriptionManager,
991    ) {
992        let pool_address = address!("1234567890123456789012345678901234567890");
993        let swap_sig = registered_manager
994            .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
995            .unwrap();
996        let collect_sig = registered_manager
997            .get_dex_pool_collect_event_signature(&DexType::UniswapV3)
998            .unwrap();
999
1000        assert!(
1001            registered_manager
1002                .get_active_subscribed_dex_event_signatures(&DexType::UniswapV3)
1003                .is_empty()
1004        );
1005
1006        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
1007        assert_eq!(
1008            registered_manager.get_active_subscribed_dex_event_signatures(&DexType::UniswapV3),
1009            vec![swap_sig]
1010        );
1011
1012        registered_manager.subscribe_collects(DexType::UniswapV3, pool_address);
1013        let active_signatures =
1014            registered_manager.get_active_subscribed_dex_event_signatures(&DexType::UniswapV3);
1015
1016        assert_eq!(active_signatures.len(), 2);
1017        assert!(active_signatures.contains(&collect_sig));
1018    }
1019
1020    #[rstest]
1021    fn test_fee_protocol_events_are_active_only_when_subscribed(
1022        mut registered_manager: DefiDataSubscriptionManager,
1023    ) {
1024        let pool_address = address!("1234567890123456789012345678901234567890");
1025        registered_manager.register_dex_fee_protocol_events(
1026            DexType::UniswapV3,
1027            Some("SetFeeProtocol(uint8,uint8,uint8,uint8)"),
1028            Some("CollectProtocol(address,address,uint128,uint128)"),
1029        );
1030
1031        let update_sig = registered_manager
1032            .get_dex_pool_fee_protocol_update_event_signature(&DexType::UniswapV3)
1033            .unwrap();
1034        let collect_sig = registered_manager
1035            .get_dex_pool_fee_protocol_collect_event_signature(&DexType::UniswapV3)
1036            .unwrap();
1037
1038        registered_manager.subscribe_fee_protocol_updates(DexType::UniswapV3, pool_address);
1039        registered_manager.subscribe_fee_protocol_collects(DexType::UniswapV3, pool_address);
1040
1041        let active_signatures =
1042            registered_manager.get_active_subscribed_dex_event_signatures(&DexType::UniswapV3);
1043        let addresses =
1044            registered_manager.get_subscribed_dex_contract_addresses(&DexType::UniswapV3);
1045
1046        assert_eq!(addresses.len(), 1);
1047        assert!(addresses.contains(&pool_address));
1048        assert!(active_signatures.contains(&update_sig));
1049        assert!(active_signatures.contains(&collect_sig));
1050    }
1051
1052    #[rstest]
1053    fn test_same_event_demand_releases_with_final_owner(
1054        mut registered_manager: DefiDataSubscriptionManager,
1055    ) {
1056        let pool_address = address!("1234567890123456789012345678901234567890");
1057        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
1058        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
1059
1060        registered_manager.unsubscribe_swaps(DexType::UniswapV3, pool_address);
1061        let after_first =
1062            registered_manager.get_subscribed_pool_swap_addresses(&DexType::UniswapV3);
1063        registered_manager.unsubscribe_swaps(DexType::UniswapV3, pool_address);
1064
1065        assert_eq!(after_first, vec![pool_address]);
1066        assert!(
1067            registered_manager
1068                .get_subscribed_pool_swap_addresses(&DexType::UniswapV3)
1069                .is_empty()
1070        );
1071        assert!(!registered_manager.has_pool_event_subscriptions());
1072    }
1073
1074    #[rstest]
1075    fn test_unsubscribe_without_demand_is_ignored(
1076        mut registered_manager: DefiDataSubscriptionManager,
1077    ) {
1078        let pool_address = address!("1234567890123456789012345678901234567890");
1079        registered_manager.unsubscribe_swaps(DexType::UniswapV3, pool_address);
1080        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
1081
1082        assert_eq!(
1083            registered_manager.get_subscribed_pool_swap_addresses(&DexType::UniswapV3),
1084            vec![pool_address]
1085        );
1086    }
1087
1088    #[derive(Clone, Copy, Debug)]
1089    enum NarrowShape {
1090        Swaps,
1091        LiquidityUpdates,
1092        FeeCollects,
1093        FlashEvents,
1094    }
1095
1096    impl NarrowShape {
1097        fn subscribe(self, manager: &mut DefiDataSubscriptionManager, address: Address) {
1098            match self {
1099                Self::Swaps => manager.subscribe_swaps(DexType::UniswapV3, address),
1100                Self::LiquidityUpdates => {
1101                    manager.subscribe_mints(DexType::UniswapV3, address);
1102                    manager.subscribe_burns(DexType::UniswapV3, address);
1103                }
1104                Self::FeeCollects => manager.subscribe_collects(DexType::UniswapV3, address),
1105                Self::FlashEvents => manager.subscribe_flashes(DexType::UniswapV3, address),
1106            }
1107        }
1108
1109        fn unsubscribe(self, manager: &mut DefiDataSubscriptionManager, address: Address) {
1110            match self {
1111                Self::Swaps => manager.unsubscribe_swaps(DexType::UniswapV3, address),
1112                Self::LiquidityUpdates => {
1113                    manager.unsubscribe_mints(DexType::UniswapV3, address);
1114                    manager.unsubscribe_burns(DexType::UniswapV3, address);
1115                }
1116                Self::FeeCollects => manager.unsubscribe_collects(DexType::UniswapV3, address),
1117                Self::FlashEvents => manager.unsubscribe_flashes(DexType::UniswapV3, address),
1118            }
1119        }
1120
1121        fn signatures(self, manager: &DefiDataSubscriptionManager) -> Vec<String> {
1122            let dex = DexType::UniswapV3;
1123
1124            let signatures = match self {
1125                Self::Swaps => vec![manager.get_dex_pool_swap_event_signature(&dex)],
1126                Self::LiquidityUpdates => vec![
1127                    manager.get_dex_pool_mint_event_signature(&dex),
1128                    manager.get_dex_pool_burn_event_signature(&dex),
1129                ],
1130                Self::FeeCollects => vec![manager.get_dex_pool_collect_event_signature(&dex)],
1131                Self::FlashEvents => vec![manager.get_dex_pool_flash_event_signature(&dex)],
1132            };
1133
1134            signatures.into_iter().map(Option::unwrap).collect()
1135        }
1136    }
1137
1138    fn sorted(mut values: Vec<String>) -> Vec<String> {
1139        values.sort();
1140        values
1141    }
1142
1143    #[rstest]
1144    fn test_complete_pool_overlaps_narrow_shape(
1145        mut registered_manager: DefiDataSubscriptionManager,
1146        #[values(
1147            NarrowShape::Swaps,
1148            NarrowShape::LiquidityUpdates,
1149            NarrowShape::FeeCollects,
1150            NarrowShape::FlashEvents
1151        )]
1152        narrow: NarrowShape,
1153        #[values(true, false)] pool_first: bool,
1154    ) {
1155        let dex = DexType::UniswapV3;
1156        let pool_address = address!("1234567890123456789012345678901234567890");
1157        registered_manager.register_dex_fee_protocol_events(
1158            dex,
1159            Some("SetFeeProtocol(uint8,uint8,uint8,uint8)"),
1160            Some("CollectProtocol(address,address,uint128,uint128)"),
1161        );
1162        let all_signatures = sorted(registered_manager.get_subscribed_dex_event_signatures(&dex));
1163        let narrow_signatures = sorted(narrow.signatures(&registered_manager));
1164        registered_manager.subscribe_pool(dex, pool_address);
1165        narrow.subscribe(&mut registered_manager, pool_address);
1166
1167        if pool_first {
1168            registered_manager.unsubscribe_pool(dex, pool_address);
1169        } else {
1170            narrow.unsubscribe(&mut registered_manager, pool_address);
1171        }
1172
1173        let addresses_after_first = registered_manager.get_subscribed_dex_contract_addresses(&dex);
1174        let signatures_after_first =
1175            sorted(registered_manager.get_active_subscribed_dex_event_signatures(&dex));
1176
1177        if pool_first {
1178            narrow.unsubscribe(&mut registered_manager, pool_address);
1179        } else {
1180            registered_manager.unsubscribe_pool(dex, pool_address);
1181        }
1182
1183        assert_eq!(all_signatures.len(), 7);
1184        assert_eq!(addresses_after_first, vec![pool_address]);
1185
1186        if pool_first {
1187            assert_eq!(signatures_after_first, narrow_signatures);
1188        } else {
1189            assert_eq!(signatures_after_first, all_signatures);
1190        }
1191
1192        assert!(
1193            registered_manager
1194                .get_subscribed_dex_contract_addresses(&dex)
1195                .is_empty()
1196        );
1197        assert!(
1198            registered_manager
1199                .get_active_subscribed_dex_event_signatures(&dex)
1200                .is_empty()
1201        );
1202        assert!(!registered_manager.has_pool_event_subscriptions());
1203    }
1204
1205    #[rstest]
1206    fn test_get_combined_addresses_from_all_events(
1207        mut registered_manager: DefiDataSubscriptionManager,
1208    ) {
1209        let pool1 = address!("1111111111111111111111111111111111111111");
1210        let pool2 = address!("2222222222222222222222222222222222222222");
1211        let pool3 = address!("3333333333333333333333333333333333333333");
1212
1213        // Subscribe different pools to different events
1214        registered_manager.subscribe_swaps(DexType::UniswapV3, pool1);
1215        registered_manager.subscribe_mints(DexType::UniswapV3, pool2);
1216        registered_manager.subscribe_burns(DexType::UniswapV3, pool3);
1217
1218        // Should get all unique addresses
1219        let addresses =
1220            registered_manager.get_subscribed_dex_contract_addresses(&DexType::UniswapV3);
1221        assert_eq!(addresses.len(), 3);
1222        assert!(addresses.contains(&pool1));
1223        assert!(addresses.contains(&pool2));
1224        assert!(addresses.contains(&pool3));
1225    }
1226
1227    #[rstest]
1228    fn test_event_signature_encoding(registered_manager: DefiDataSubscriptionManager) {
1229        // Known event signature and its expected keccak256 hash
1230        // Swap(address,address,int256,int256,uint160,uint128,int24) for UniswapV3
1231        let swap_sig = registered_manager
1232            .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
1233            .unwrap();
1234
1235        // Should be properly formatted hex string
1236        assert!(swap_sig.starts_with("0x"));
1237        assert_eq!(swap_sig.len(), 66); // 0x + 64 hex chars (32 bytes)
1238
1239        // Verify it's valid hex
1240        let hex_part = &swap_sig[2..];
1241        assert!(hex_part.chars().all(|c| c.is_ascii_hexdigit()));
1242    }
1243
1244    #[rstest]
1245    #[case(DexType::UniswapV3)]
1246    #[case(DexType::UniswapV2)]
1247    fn test_complete_subscription_workflow(#[case] dex_type: DexType) {
1248        let mut manager = DefiDataSubscriptionManager::new();
1249        let pool1 = address!("1111111111111111111111111111111111111111");
1250        let pool2 = address!("2222222222222222222222222222222222222222");
1251
1252        // Step 1: Register DEX
1253        manager.register_dex_for_subscriptions(
1254            dex_type,
1255            "Swap(address,uint256,uint256)",
1256            "Mint(address,uint256)",
1257            "Burn(address,uint256)",
1258            "Collect(address,uint256,uint256)",
1259            Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
1260        );
1261
1262        // Step 2: Subscribe to events
1263        manager.subscribe_swaps(dex_type, pool1);
1264        manager.subscribe_swaps(dex_type, pool2);
1265        manager.subscribe_mints(dex_type, pool1);
1266        manager.subscribe_burns(dex_type, pool2);
1267
1268        // Step 3: Verify subscriptions
1269        let addresses = manager.get_subscribed_dex_contract_addresses(&dex_type);
1270        assert_eq!(addresses.len(), 2);
1271        assert!(addresses.contains(&pool1));
1272        assert!(addresses.contains(&pool2));
1273
1274        // Step 4: Get event signatures
1275        let signatures = manager.get_subscribed_dex_event_signatures(&dex_type);
1276        assert_eq!(signatures.len(), 5);
1277
1278        // Step 5: Unsubscribe from some events
1279        manager.unsubscribe_swaps(dex_type, pool1);
1280        manager.unsubscribe_burns(dex_type, pool2);
1281
1282        // Step 6: Verify remaining subscriptions (only pool1 mint remains)
1283        let remaining = manager.get_subscribed_dex_contract_addresses(&dex_type);
1284        assert!(remaining.contains(&pool1)); // Still has mint subscription
1285        assert!(remaining.contains(&pool2)); // Still has swap subscription
1286    }
1287
1288    #[rstest]
1289    fn test_register_with_raw_signatures() {
1290        let mut manager = DefiDataSubscriptionManager::new();
1291
1292        // Register with raw event signatures
1293        manager.register_dex_for_subscriptions(
1294            DexType::UniswapV3,
1295            "Swap(address,address,int256,int256,uint160,uint128,int24)",
1296            "Mint(address,address,int24,int24,uint128,uint256,uint256)",
1297            "Burn(address,int24,int24,uint128,uint256,uint256)",
1298            "Collect(address,address,int24,int24,uint128,uint128)",
1299            Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
1300        );
1301
1302        // Known keccak256 hashes for UniswapV3 events
1303        let swap_sig = manager
1304            .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
1305            .unwrap();
1306        let mint_sig = manager
1307            .get_dex_pool_mint_event_signature(&DexType::UniswapV3)
1308            .unwrap();
1309        let burn_sig = manager
1310            .get_dex_pool_burn_event_signature(&DexType::UniswapV3)
1311            .unwrap();
1312
1313        // Verify the exact hash values
1314        assert_eq!(
1315            swap_sig,
1316            "0xc42079f94a6350d7e6235f29174924f928cc2ac818eb64fed8004e115fbcca67"
1317        );
1318        assert_eq!(
1319            mint_sig,
1320            "0x7a53080ba414158be7ec69b987b5fb7d07dee101fe85488f0853ae16239d0bde"
1321        );
1322        assert_eq!(
1323            burn_sig,
1324            "0x0c396cd989a39f4459b5fa1aed6a9a8dcdbc45908acfd67e028cd568da98982c"
1325        );
1326    }
1327
1328    #[rstest]
1329    fn test_register_with_pre_encoded_signatures() {
1330        let mut manager = DefiDataSubscriptionManager::new();
1331
1332        // Register with pre-encoded keccak256 hashes (with 0x prefix)
1333        manager.register_dex_for_subscriptions(
1334            DexType::UniswapV3,
1335            "Swap(address,address,int256,int256,uint160,uint128,int24)",
1336            "Mint(address,address,int24,int24,uint128,uint256,uint256)",
1337            "Burn(address,int24,int24,uint128,uint256,uint256)",
1338            "Collect(address,address,int24,int24,uint128,uint128)",
1339            Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
1340        );
1341
1342        // Should store them unchanged (normalized to lowercase)
1343        let swap_sig = manager
1344            .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
1345            .unwrap();
1346        let mint_sig = manager
1347            .get_dex_pool_mint_event_signature(&DexType::UniswapV3)
1348            .unwrap();
1349        let burn_sig = manager
1350            .get_dex_pool_burn_event_signature(&DexType::UniswapV3)
1351            .unwrap();
1352
1353        assert_eq!(
1354            swap_sig,
1355            "0xc42079f94a6350d7e6235f29174924f928cc2ac818eb64fed8004e115fbcca67"
1356        );
1357        assert_eq!(
1358            mint_sig,
1359            "0x7a53080ba414158be7ec69b987b5fb7d07dee101fe85488f0853ae16239d0bde"
1360        );
1361        assert_eq!(
1362            burn_sig,
1363            "0x0c396cd989a39f4459b5fa1aed6a9a8dcdbc45908acfd67e028cd568da98982c"
1364        );
1365    }
1366
1367    #[rstest]
1368    fn test_register_with_pre_encoded_signatures_no_prefix() {
1369        let mut manager = DefiDataSubscriptionManager::new();
1370
1371        // Register with pre-encoded hashes without 0x prefix
1372        manager.register_dex_for_subscriptions(
1373            DexType::UniswapV3,
1374            "Swap(address,address,int256,int256,uint160,uint128,int24)",
1375            "Mint(address,address,int24,int24,uint128,uint256,uint256)",
1376            "Burn(address,int24,int24,uint128,uint256,uint256)",
1377            "Collect(address,address,int24,int24,uint128,uint128)",
1378            Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
1379        );
1380
1381        // Should add 0x prefix and normalize to lowercase
1382        let swap_sig = manager
1383            .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
1384            .unwrap();
1385        let mint_sig = manager
1386            .get_dex_pool_mint_event_signature(&DexType::UniswapV3)
1387            .unwrap();
1388        let burn_sig = manager
1389            .get_dex_pool_burn_event_signature(&DexType::UniswapV3)
1390            .unwrap();
1391
1392        assert_eq!(
1393            swap_sig,
1394            "0xc42079f94a6350d7e6235f29174924f928cc2ac818eb64fed8004e115fbcca67"
1395        );
1396        assert_eq!(
1397            mint_sig,
1398            "0x7a53080ba414158be7ec69b987b5fb7d07dee101fe85488f0853ae16239d0bde"
1399        );
1400        assert_eq!(
1401            burn_sig,
1402            "0x0c396cd989a39f4459b5fa1aed6a9a8dcdbc45908acfd67e028cd568da98982c"
1403        );
1404    }
1405}
nautilus_blockchain/data/subscription.rs

nautilus_blockchain/data/
subscription.rs
