Skip to main content

nautilus_data/defi/
engine.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
16//! DeFi-specific data engine functionality.
17//!
18//! This module provides DeFi processing methods for the `DataEngine`.
19//! All code in this module requires the `defi` feature flag.
20
21use std::{rc::Rc, sync::Arc};
22
23use nautilus_common::{
24    defi,
25    messages::defi::{
26        DefiRequestCommand, DefiSubscribeCommand, DefiUnsubscribeCommand, RequestPoolSnapshot,
27    },
28    msgbus::{self, TypedHandler},
29};
30use nautilus_core::UUID4;
31use nautilus_model::{
32    defi::{
33        Blockchain, DefiData, PoolProfiler,
34        data::{DexPoolData, block::BlockPosition},
35    },
36    identifiers::{ClientId, InstrumentId},
37};
38
39use crate::engine::{
40    DataEngine,
41    pool::{
42        PoolCollectHandler, PoolFlashHandler, PoolLiquidityHandler, PoolSwapHandler, PoolUpdater,
43    },
44};
45
46/// Extracts the block position tuple from a `DexPoolData` event.
47fn get_event_block_position(event: &DexPoolData) -> (u64, u32, u32) {
48    match event {
49        DexPoolData::Swap(s) => (s.block, s.transaction_index, s.log_index),
50        DexPoolData::LiquidityUpdate(u) => (u.block, u.transaction_index, u.log_index),
51        DexPoolData::FeeCollect(c) => (c.block, c.transaction_index, c.log_index),
52        DexPoolData::FeeProtocolUpdate(u) => (u.block, u.transaction_index, u.log_index),
53        DexPoolData::FeeProtocolCollect(c) => (c.block, c.transaction_index, c.log_index),
54        DexPoolData::Flash(f) => (f.block, f.transaction_index, f.log_index),
55    }
56}
57
58/// Converts buffered `DefiData` events to `DexPoolData` and sorts by block position.
59fn convert_and_sort_buffered_events(buffered_events: Vec<DefiData>) -> Vec<DexPoolData> {
60    let mut events: Vec<DexPoolData> = buffered_events
61        .into_iter()
62        .filter_map(|event| match event {
63            DefiData::PoolSwap(swap) => Some(DexPoolData::Swap(swap)),
64            DefiData::PoolLiquidityUpdate(update) => Some(DexPoolData::LiquidityUpdate(update)),
65            DefiData::PoolFeeCollect(collect) => Some(DexPoolData::FeeCollect(collect)),
66            DefiData::PoolFeeProtocolUpdate(update) => Some(DexPoolData::FeeProtocolUpdate(update)),
67            DefiData::PoolFeeProtocolCollect(collect) => {
68                Some(DexPoolData::FeeProtocolCollect(collect))
69            }
70            DefiData::PoolFlash(flash) => Some(DexPoolData::Flash(flash)),
71            _ => None,
72        })
73        .collect();
74
75    events.sort_by(|a, b| {
76        let pos_a = get_event_block_position(a);
77        let pos_b = get_event_block_position(b);
78        pos_a.cmp(&pos_b)
79    });
80
81    events
82}
83
84impl DataEngine {
85    /// Returns all blockchains for which blocks subscriptions exist.
86    #[must_use]
87    pub fn subscribed_blocks(&self) -> Vec<Blockchain> {
88        self.collect_subscriptions(|client| &client.subscriptions_blocks)
89    }
90
91    /// Returns all instrument IDs for which pool subscriptions exist.
92    #[must_use]
93    pub fn subscribed_pools(&self) -> Vec<InstrumentId> {
94        self.collect_subscriptions(|client| &client.subscriptions_pools)
95    }
96
97    /// Returns all instrument IDs for which swap subscriptions exist.
98    #[must_use]
99    pub fn subscribed_pool_swaps(&self) -> Vec<InstrumentId> {
100        self.collect_subscriptions(|client| &client.subscriptions_pool_swaps)
101    }
102
103    /// Returns all instrument IDs for which liquidity update subscriptions exist.
104    #[must_use]
105    pub fn subscribed_pool_liquidity_updates(&self) -> Vec<InstrumentId> {
106        self.collect_subscriptions(|client| &client.subscriptions_pool_liquidity_updates)
107    }
108
109    /// Returns all instrument IDs for which fee collect subscriptions exist.
110    #[must_use]
111    pub fn subscribed_pool_fee_collects(&self) -> Vec<InstrumentId> {
112        self.collect_subscriptions(|client| &client.subscriptions_pool_fee_collects)
113    }
114
115    /// Returns all instrument IDs for which flash loan subscriptions exist.
116    #[must_use]
117    pub fn subscribed_pool_flash(&self) -> Vec<InstrumentId> {
118        self.collect_subscriptions(|client| &client.subscriptions_pool_flash)
119    }
120
121    /// Handles a subscribe command, updating internal state and forwarding to the client.
122    ///
123    /// # Errors
124    ///
125    /// Returns an error if the subscription is invalid (e.g., synthetic instrument for book data),
126    /// or if the underlying client operation fails.
127    pub fn execute_defi_subscribe(&mut self, cmd: DefiSubscribeCommand) -> anyhow::Result<()> {
128        if let Some(client_id) = cmd.client_id()
129            && self.external_clients.contains(client_id)
130        {
131            if self.config.debug {
132                log::debug!("Skipping defi subscribe for external client {client_id}: {cmd:?}");
133            }
134            return Ok(());
135        }
136
137        if let Some(client) = self.get_client(cmd.client_id(), cmd.venue()) {
138            log::info!("Forwarding subscription to client {}", client.client_id);
139            client.execute_defi_subscribe(cmd.clone());
140        } else {
141            log::error!(
142                "Cannot handle command: no client found for client_id={:?}, venue={:?}",
143                cmd.client_id(),
144                cmd.venue(),
145            );
146        }
147
148        match cmd {
149            DefiSubscribeCommand::Pool(cmd) => {
150                self.setup_pool_updater(&cmd.instrument_id, cmd.client_id.as_ref());
151            }
152            DefiSubscribeCommand::PoolSwaps(cmd) => {
153                self.setup_pool_updater(&cmd.instrument_id, cmd.client_id.as_ref());
154            }
155            DefiSubscribeCommand::PoolLiquidityUpdates(cmd) => {
156                self.setup_pool_updater(&cmd.instrument_id, cmd.client_id.as_ref());
157            }
158            DefiSubscribeCommand::PoolFeeCollects(cmd) => {
159                self.setup_pool_updater(&cmd.instrument_id, cmd.client_id.as_ref());
160            }
161            DefiSubscribeCommand::PoolFlashEvents(cmd) => {
162                self.setup_pool_updater(&cmd.instrument_id, cmd.client_id.as_ref());
163            }
164            DefiSubscribeCommand::Blocks(_) => {} // No pool setup needed for blocks
165        }
166
167        Ok(())
168    }
169
170    /// Handles an unsubscribe command, updating internal state and forwarding to the client.
171    ///
172    /// # Errors
173    ///
174    /// Returns an error if the underlying client operation fails.
175    pub fn execute_defi_unsubscribe(&mut self, cmd: &DefiUnsubscribeCommand) -> anyhow::Result<()> {
176        if let Some(client_id) = cmd.client_id()
177            && self.external_clients.contains(client_id)
178        {
179            if self.config.debug {
180                log::debug!("Skipping defi unsubscribe for external client {client_id}: {cmd:?}");
181            }
182            return Ok(());
183        }
184
185        if let Some(client) = self.get_client(cmd.client_id(), cmd.venue()) {
186            client.execute_defi_unsubscribe(cmd);
187        } else {
188            log::error!(
189                "Cannot handle command: no client found for client_id={:?}, venue={:?}",
190                cmd.client_id(),
191                cmd.venue(),
192            );
193        }
194
195        Ok(())
196    }
197
198    /// Sends a [`DefiRequestCommand`] to a suitable data client implementation.
199    ///
200    /// # Errors
201    ///
202    /// Returns an error if no client is found for the given client ID or venue,
203    /// or if the client fails to process the request.
204    pub fn execute_defi_request(&mut self, req: DefiRequestCommand) -> anyhow::Result<()> {
205        // Skip requests for external clients
206        if let Some(cid) = req.client_id()
207            && self.external_clients.contains(cid)
208        {
209            if self.config.debug {
210                log::debug!("Skipping defi data request for external client {cid}: {req:?}");
211            }
212            return Ok(());
213        }
214
215        if let Some(client) = self.get_client(req.client_id(), req.venue()) {
216            client.execute_defi_request(req)
217        } else {
218            anyhow::bail!(
219                "Cannot handle request: no client found for {:?} {:?}",
220                req.client_id(),
221                req.venue()
222            );
223        }
224    }
225
226    /// Processes DeFi-specific data events.
227    pub fn process_defi_data(&mut self, data: DefiData) {
228        self.increment_data_count();
229
230        match data {
231            DefiData::Block(block) => {
232                let topic = defi::switchboard::get_defi_blocks_topic(block.chain());
233                msgbus::publish_defi_block(topic, &block);
234            }
235            DefiData::Pool(pool) => {
236                if let Err(e) = self.cache.borrow_mut().add_pool(pool.clone()) {
237                    log::error!("Failed to add Pool to cache: {e}");
238                }
239
240                // Check if pool profiler creation was deferred. When a snapshot is also
241                // pending, leave the updater marker in place so a second subscription cannot
242                // race ahead and install a duplicate updater. The snapshot handler clears
243                // both flags and creates the updater once.
244                if self.pool_updaters_pending.contains(&pool.instrument_id) {
245                    if self.pool_snapshot_pending.contains(&pool.instrument_id) {
246                        log::debug!(
247                            "Pool {} loaded; deferring profiler creation to snapshot handler",
248                            pool.instrument_id
249                        );
250                    } else {
251                        self.pool_updaters_pending.remove(&pool.instrument_id);
252                        log::info!(
253                            "Pool {} now loaded, creating deferred pool profiler",
254                            pool.instrument_id
255                        );
256                        self.setup_pool_updater(&pool.instrument_id, None);
257                    }
258                }
259
260                let topic = defi::switchboard::get_defi_pool_topic(pool.instrument_id);
261                msgbus::publish_defi_pool(topic, &pool);
262            }
263            DefiData::PoolSnapshot(snapshot) => {
264                let instrument_id = snapshot.instrument_id;
265                log::info!(
266                    "Received pool snapshot for {instrument_id} at block {} with {} positions and {} ticks",
267                    snapshot.block_position.number,
268                    snapshot.positions.len(),
269                    snapshot.ticks.len()
270                );
271
272                // Validate we're expecting this snapshot
273                if !self.pool_snapshot_pending.contains(&instrument_id) {
274                    log::warn!(
275                        "Received unexpected pool snapshot for {instrument_id} (not in pending set)"
276                    );
277                    return;
278                }
279
280                // Get pool from cache
281                let pool = match self.cache.borrow().pool(&instrument_id) {
282                    Some(pool) => Arc::new(pool.clone()),
283                    None => {
284                        log::error!(
285                            "Pool {instrument_id} not found in cache when processing snapshot"
286                        );
287                        return;
288                    }
289                };
290
291                // Defensive: refuse stub snapshots that slipped past the bootstrap-side
292                // guards. Installing one would leave Python actors observing an initialized
293                // profiler with zero liquidity; better to leave the pool without a profiler
294                // so the bad state is visible.
295                if snapshot.positions.is_empty()
296                    && snapshot.ticks.is_empty()
297                    && snapshot.block_position.number == pool.creation_block
298                {
299                    log::warn!(
300                        "Refusing empty stub snapshot for {instrument_id} at pool creation block {}; pool will remain without profiler",
301                        snapshot.block_position.number,
302                    );
303                    self.pool_snapshot_pending.remove(&instrument_id);
304                    self.pool_updaters_pending.remove(&instrument_id);
305                    self.pool_event_buffers.remove(&instrument_id);
306                    return;
307                }
308
309                // Create profiler and restore from snapshot
310                let mut profiler = PoolProfiler::new(pool);
311                if let Err(e) = profiler.restore_from_snapshot(snapshot.clone()) {
312                    log::error!(
313                        "Failed to restore profiler from snapshot for {instrument_id}: {e}"
314                    );
315                    return;
316                }
317                log::debug!("Restored pool profiler for {instrument_id} from snapshot");
318
319                // Process buffered events
320                let buffered_events = self
321                    .pool_event_buffers
322                    .remove(&instrument_id)
323                    .unwrap_or_default();
324
325                if !buffered_events.is_empty() {
326                    log::info!(
327                        "Processing {} buffered events for {instrument_id}",
328                        buffered_events.len()
329                    );
330
331                    let events_to_apply = convert_and_sort_buffered_events(buffered_events);
332                    let applied_count = Self::apply_buffered_events_to_profiler(
333                        &mut profiler,
334                        events_to_apply,
335                        &snapshot.block_position,
336                        instrument_id,
337                    );
338
339                    log::info!(
340                        "Applied {applied_count} buffered events to profiler for {instrument_id}"
341                    );
342                }
343
344                // Add profiler to cache
345                if let Err(e) = self.cache.borrow_mut().add_pool_profiler(profiler) {
346                    log::error!("Failed to add pool profiler to cache for {instrument_id}: {e}");
347                    return;
348                }
349
350                // Create updater and subscribe to topics
351                self.pool_snapshot_pending.remove(&instrument_id);
352                self.pool_updaters_pending.remove(&instrument_id);
353                let updater = Rc::new(PoolUpdater::new(&instrument_id, self.cache.clone()));
354
355                self.subscribe_pool_updater_topics(instrument_id, updater.clone());
356                self.pool_updaters.insert(instrument_id, updater);
357
358                log::info!(
359                    "Pool profiler setup completed for {instrument_id}, now processing live events"
360                );
361            }
362            DefiData::PoolSwap(swap) => {
363                let instrument_id = swap.instrument_id;
364                // Buffer if waiting for snapshot, otherwise publish
365                if self.pool_snapshot_pending.contains(&instrument_id) {
366                    log::debug!("Buffering swap event for {instrument_id} (waiting for snapshot)");
367                    self.pool_event_buffers
368                        .entry(instrument_id)
369                        .or_default()
370                        .push(DefiData::PoolSwap(swap));
371                } else {
372                    let topic = defi::switchboard::get_defi_pool_swaps_topic(instrument_id);
373                    msgbus::publish_defi_swap(topic, &swap);
374                }
375            }
376            DefiData::PoolLiquidityUpdate(update) => {
377                let instrument_id = update.instrument_id;
378                // Buffer if waiting for snapshot, otherwise publish
379                if self.pool_snapshot_pending.contains(&instrument_id) {
380                    log::debug!(
381                        "Buffering liquidity update event for {instrument_id} (waiting for snapshot)"
382                    );
383                    self.pool_event_buffers
384                        .entry(instrument_id)
385                        .or_default()
386                        .push(DefiData::PoolLiquidityUpdate(update));
387                } else {
388                    let topic = defi::switchboard::get_defi_liquidity_topic(instrument_id);
389                    msgbus::publish_defi_liquidity(topic, &update);
390                }
391            }
392            DefiData::PoolFeeCollect(collect) => {
393                let instrument_id = collect.instrument_id;
394                // Buffer if waiting for snapshot, otherwise publish
395                if self.pool_snapshot_pending.contains(&instrument_id) {
396                    log::debug!(
397                        "Buffering fee collect event for {instrument_id} (waiting for snapshot)"
398                    );
399                    self.pool_event_buffers
400                        .entry(instrument_id)
401                        .or_default()
402                        .push(DefiData::PoolFeeCollect(collect));
403                } else {
404                    let topic = defi::switchboard::get_defi_collect_topic(instrument_id);
405                    msgbus::publish_defi_collect(topic, &collect);
406                }
407            }
408            DefiData::PoolFeeProtocolUpdate(update) => {
409                let instrument_id = update.instrument_id;
410                // A protocol-fee change is profiler infrastructure, not a subscriber data stream:
411                // it only needs to keep the profiler's `fee_protocol` consistent for swap and flash
412                // fee splitting. Buffer until the snapshot lands, otherwise apply in place.
413                if self.pool_snapshot_pending.contains(&instrument_id) {
414                    log::debug!(
415                        "Buffering fee protocol update for {instrument_id} (waiting for snapshot)"
416                    );
417                    self.pool_event_buffers
418                        .entry(instrument_id)
419                        .or_default()
420                        .push(DefiData::PoolFeeProtocolUpdate(update));
421                } else if let Some(profiler) =
422                    self.cache.borrow_mut().pool_profiler_mut(&instrument_id)
423                    && let Err(e) = profiler.process_fee_protocol_update(&update)
424                {
425                    log::error!("Failed to process pool fee protocol update: {e}");
426                }
427            }
428            DefiData::PoolFeeProtocolCollect(collect) => {
429                let instrument_id = collect.instrument_id;
430                // A protocol-fee withdrawal is profiler infrastructure, not a subscriber data
431                // stream: it only decrements the profiler's accrued protocol-fee balances. Buffer
432                // until the snapshot lands, otherwise apply in place.
433                if self.pool_snapshot_pending.contains(&instrument_id) {
434                    log::debug!(
435                        "Buffering fee protocol collect event for {instrument_id} (waiting for snapshot)"
436                    );
437                    self.pool_event_buffers
438                        .entry(instrument_id)
439                        .or_default()
440                        .push(DefiData::PoolFeeProtocolCollect(collect));
441                } else if let Some(profiler) =
442                    self.cache.borrow_mut().pool_profiler_mut(&instrument_id)
443                    && let Err(e) = profiler.process_fee_protocol_collect(&collect)
444                {
445                    log::error!("Failed to process pool fee protocol collect event: {e}");
446                }
447            }
448            DefiData::PoolFlash(flash) => {
449                let instrument_id = flash.instrument_id;
450                // Buffer if waiting for snapshot, otherwise publish
451                if self.pool_snapshot_pending.contains(&instrument_id) {
452                    log::debug!("Buffering flash event for {instrument_id} (waiting for snapshot)");
453                    self.pool_event_buffers
454                        .entry(instrument_id)
455                        .or_default()
456                        .push(DefiData::PoolFlash(flash));
457                } else {
458                    let topic = defi::switchboard::get_defi_flash_topic(instrument_id);
459                    msgbus::publish_defi_flash(topic, &flash);
460                }
461            }
462        }
463    }
464
465    /// Subscribes a pool updater to all relevant pool data topics using typed handlers.
466    fn subscribe_pool_updater_topics(&self, instrument_id: InstrumentId, updater: Rc<PoolUpdater>) {
467        let priority = Some(self.msgbus_priority);
468
469        // Subscribe swap handler
470        let swap_topic = defi::switchboard::get_defi_pool_swaps_topic(instrument_id);
471        let swap_handler = TypedHandler(Rc::new(PoolSwapHandler::new(updater.clone())));
472        msgbus::subscribe_defi_swaps(swap_topic.into(), swap_handler, priority);
473
474        // Subscribe liquidity handler
475        let liq_topic = defi::switchboard::get_defi_liquidity_topic(instrument_id);
476        let liq_handler = TypedHandler(Rc::new(PoolLiquidityHandler::new(updater.clone())));
477        msgbus::subscribe_defi_liquidity(liq_topic.into(), liq_handler, priority);
478
479        // Subscribe collect handler
480        let collect_topic = defi::switchboard::get_defi_collect_topic(instrument_id);
481        let collect_handler = TypedHandler(Rc::new(PoolCollectHandler::new(updater.clone())));
482        msgbus::subscribe_defi_collects(collect_topic.into(), collect_handler, priority);
483
484        // Subscribe flash handler
485        let flash_topic = defi::switchboard::get_defi_flash_topic(instrument_id);
486        let flash_handler = TypedHandler(Rc::new(PoolFlashHandler::new(updater)));
487        msgbus::subscribe_defi_flash(flash_topic.into(), flash_handler, priority);
488    }
489
490    /// Applies buffered events to a pool profiler, filtering to events after the snapshot.
491    ///
492    /// Returns the count of successfully applied events.
493    fn apply_buffered_events_to_profiler(
494        profiler: &mut PoolProfiler,
495        events: Vec<DexPoolData>,
496        snapshot_block: &BlockPosition,
497        instrument_id: InstrumentId,
498    ) -> usize {
499        let mut applied_count = 0;
500
501        for event in events {
502            let event_block = get_event_block_position(&event);
503
504            // Only apply events that occurred after the snapshot
505            let is_after_snapshot = event_block.0 > snapshot_block.number
506                || (event_block.0 == snapshot_block.number
507                    && event_block.1 > snapshot_block.transaction_index)
508                || (event_block.0 == snapshot_block.number
509                    && event_block.1 == snapshot_block.transaction_index
510                    && event_block.2 > snapshot_block.log_index);
511
512            if is_after_snapshot {
513                if let Err(e) = profiler.process(&event) {
514                    log::error!(
515                        "Failed to apply buffered event to profiler for {instrument_id}: {e}"
516                    );
517                } else {
518                    applied_count += 1;
519                }
520            }
521        }
522
523        applied_count
524    }
525
526    fn setup_pool_updater(&mut self, instrument_id: &InstrumentId, client_id: Option<&ClientId>) {
527        // Early return if updater already exists or we are in the middle of setting it up.
528        if self.pool_updaters.contains_key(instrument_id)
529            || self.pool_updaters_pending.contains(instrument_id)
530        {
531            log::debug!("Pool updater for {instrument_id} already exists");
532            return;
533        }
534
535        log::info!("Setting up pool updater for {instrument_id}");
536
537        // Check cache state and ensure profiler exists
538        {
539            let mut cache = self.cache.borrow_mut();
540
541            if cache.pool_profiler(instrument_id).is_some() {
542                // Profiler already exists, proceed to create updater
543                log::debug!("Pool profiler already exists for {instrument_id}");
544            } else if let Some(pool) = cache.pool(instrument_id) {
545                // Pool exists but no profiler, create profiler from pool
546                let pool = Arc::new(pool.clone());
547                let mut pool_profiler = PoolProfiler::new(pool.clone());
548
549                if let Some(initial_sqrt_price_x96) = pool.initial_sqrt_price_x96 {
550                    if let Err(e) = pool_profiler.initialize(initial_sqrt_price_x96) {
551                        log::error!("Failed to initialize pool profiler for {instrument_id}: {e}");
552                        drop(cache);
553                        return;
554                    }
555                    log::debug!(
556                        "Initialized pool profiler for {instrument_id} with sqrt_price {initial_sqrt_price_x96}"
557                    );
558                } else {
559                    log::debug!("Created pool profiler for {instrument_id}");
560                }
561
562                if let Err(e) = cache.add_pool_profiler(pool_profiler) {
563                    log::error!("Failed to add pool profiler for {instrument_id}: {e}");
564                    drop(cache);
565                    return;
566                }
567                drop(cache);
568            } else {
569                // Neither profiler nor pool exists, request snapshot
570                drop(cache);
571
572                let request_id = UUID4::new();
573                let ts_init = self.clock.borrow().timestamp_ns();
574                let request = RequestPoolSnapshot::new(
575                    *instrument_id,
576                    client_id.copied(),
577                    request_id,
578                    ts_init,
579                    None,
580                );
581
582                if let Err(e) = self.execute_defi_request(DefiRequestCommand::PoolSnapshot(request))
583                {
584                    log::warn!("Failed to request pool snapshot for {instrument_id}: {e}");
585                } else {
586                    log::debug!("Requested pool snapshot for {instrument_id}");
587                    self.pool_snapshot_pending.insert(*instrument_id);
588                    self.pool_updaters_pending.insert(*instrument_id);
589                    self.pool_event_buffers.entry(*instrument_id).or_default();
590                }
591                return;
592            }
593        }
594
595        // Profiler exists, create updater and subscribe to topics
596        let updater = Rc::new(PoolUpdater::new(instrument_id, self.cache.clone()));
597
598        self.subscribe_pool_updater_topics(*instrument_id, updater.clone());
599        self.pool_updaters.insert(*instrument_id, updater);
600
601        log::debug!("Created PoolUpdater for instrument ID {instrument_id}");
602    }
603}
604
605#[cfg(test)]
606mod tests {
607    use std::sync::Arc;
608
609    use alloy_primitives::{Address, I256, U160, U256};
610    use nautilus_core::UnixNanos;
611    use nautilus_model::{
612        defi::{
613            Chain, DefiData, PoolFeeCollect, PoolFeeProtocolUpdate, PoolFlash, PoolIdentifier,
614            PoolLiquidityUpdate, PoolLiquidityUpdateType, PoolSwap,
615            chain::chains,
616            data::DexPoolData,
617            dex::{AmmType, Dex, DexType},
618        },
619        identifiers::{InstrumentId, Symbol, Venue},
620    };
621    use rstest::*;
622
623    use super::*;
624
625    #[fixture]
626    fn test_instrument_id() -> InstrumentId {
627        InstrumentId::new(Symbol::from("ETH/USDC"), Venue::from("UNISWAPV3"))
628    }
629
630    #[fixture]
631    fn test_chain() -> Arc<Chain> {
632        Arc::new(chains::ETHEREUM.clone())
633    }
634
635    #[fixture]
636    fn test_dex(test_chain: Arc<Chain>) -> Arc<Dex> {
637        Arc::new(Dex::new(
638            (*test_chain).clone(),
639            DexType::UniswapV3,
640            "0x1F98431c8aD98523631AE4a59f267346ea31F984",
641            12369621,
642            AmmType::CLAMM,
643            "PoolCreated(address,address,uint24,int24,address)",
644            "Swap(address,address,int256,int256,uint160,uint128,int24)",
645            "Mint(address,address,int24,int24,uint128,uint256,uint256)",
646            "Burn(address,int24,int24,uint128,uint256,uint256)",
647            "Collect(address,address,int24,int24,uint128,uint128)",
648        ))
649    }
650
651    fn create_test_swap(
652        test_instrument_id: InstrumentId,
653        test_chain: Arc<Chain>,
654        test_dex: Arc<Dex>,
655        block: u64,
656        tx_index: u32,
657        log_index: u32,
658    ) -> PoolSwap {
659        PoolSwap::new(
660            test_chain,
661            test_dex,
662            test_instrument_id,
663            PoolIdentifier::from_address(Address::ZERO),
664            block,
665            format!("0x{block:064x}"),
666            tx_index,
667            log_index,
668            UnixNanos::default(),
669            UnixNanos::default(),
670            Address::ZERO,
671            Address::ZERO,
672            I256::ZERO,
673            I256::ZERO,
674            U160::ZERO,
675            0,
676            0,
677        )
678    }
679
680    fn create_test_liquidity_update(
681        test_instrument_id: InstrumentId,
682        test_chain: Arc<Chain>,
683        test_dex: Arc<Dex>,
684        block: u64,
685        tx_index: u32,
686        log_index: u32,
687    ) -> PoolLiquidityUpdate {
688        PoolLiquidityUpdate::new(
689            test_chain,
690            test_dex,
691            test_instrument_id,
692            PoolIdentifier::from_address(Address::ZERO),
693            PoolLiquidityUpdateType::Mint,
694            block,
695            format!("0x{block:064x}"),
696            tx_index,
697            log_index,
698            None,
699            Address::ZERO,
700            0,
701            U256::ZERO,
702            U256::ZERO,
703            0,
704            0,
705            UnixNanos::default(),
706            UnixNanos::default(),
707        )
708    }
709
710    fn create_test_fee_collect(
711        test_instrument_id: InstrumentId,
712        test_chain: Arc<Chain>,
713        test_dex: Arc<Dex>,
714        block: u64,
715        tx_index: u32,
716        log_index: u32,
717    ) -> PoolFeeCollect {
718        PoolFeeCollect::new(
719            test_chain,
720            test_dex,
721            test_instrument_id,
722            PoolIdentifier::from_address(Address::ZERO),
723            block,
724            format!("0x{block:064x}"),
725            tx_index,
726            log_index,
727            Address::ZERO,
728            0,
729            0,
730            0,
731            0,
732            UnixNanos::default(),
733            UnixNanos::default(),
734        )
735    }
736
737    fn create_test_flash(
738        test_instrument_id: InstrumentId,
739        test_chain: Arc<Chain>,
740        test_dex: Arc<Dex>,
741        block: u64,
742        tx_index: u32,
743        log_index: u32,
744    ) -> PoolFlash {
745        PoolFlash::new(
746            test_chain,
747            test_dex,
748            test_instrument_id,
749            PoolIdentifier::from_address(Address::ZERO),
750            block,
751            format!("0x{block:064x}"),
752            tx_index,
753            log_index,
754            UnixNanos::default(),
755            UnixNanos::default(),
756            Address::ZERO,
757            Address::ZERO,
758            U256::ZERO,
759            U256::ZERO,
760            U256::ZERO,
761            U256::ZERO,
762        )
763    }
764
765    fn create_test_fee_protocol_update(
766        test_instrument_id: InstrumentId,
767        test_chain: Arc<Chain>,
768        test_dex: Arc<Dex>,
769        block: u64,
770        tx_index: u32,
771        log_index: u32,
772    ) -> PoolFeeProtocolUpdate {
773        PoolFeeProtocolUpdate::new(
774            test_chain,
775            test_dex,
776            test_instrument_id,
777            PoolIdentifier::from_address(Address::ZERO),
778            block,
779            format!("0x{block:064x}"),
780            tx_index,
781            log_index,
782            4,
783            4,
784            UnixNanos::default(),
785            UnixNanos::default(),
786        )
787    }
788
789    #[rstest]
790    fn test_get_event_block_position_swap(
791        test_instrument_id: InstrumentId,
792        test_chain: Arc<Chain>,
793        test_dex: Arc<Dex>,
794    ) {
795        let swap = create_test_swap(test_instrument_id, test_chain, test_dex, 100, 5, 3);
796        let pos = get_event_block_position(&DexPoolData::Swap(swap));
797        assert_eq!(pos, (100, 5, 3));
798    }
799
800    #[rstest]
801    fn test_get_event_block_position_liquidity_update(
802        test_instrument_id: InstrumentId,
803        test_chain: Arc<Chain>,
804        test_dex: Arc<Dex>,
805    ) {
806        let update =
807            create_test_liquidity_update(test_instrument_id, test_chain, test_dex, 200, 10, 7);
808        let pos = get_event_block_position(&DexPoolData::LiquidityUpdate(update));
809        assert_eq!(pos, (200, 10, 7));
810    }
811
812    #[rstest]
813    fn test_get_event_block_position_fee_collect(
814        test_instrument_id: InstrumentId,
815        test_chain: Arc<Chain>,
816        test_dex: Arc<Dex>,
817    ) {
818        let collect = create_test_fee_collect(test_instrument_id, test_chain, test_dex, 300, 15, 2);
819        let pos = get_event_block_position(&DexPoolData::FeeCollect(collect));
820        assert_eq!(pos, (300, 15, 2));
821    }
822
823    #[rstest]
824    fn test_get_event_block_position_flash(
825        test_instrument_id: InstrumentId,
826        test_chain: Arc<Chain>,
827        test_dex: Arc<Dex>,
828    ) {
829        let flash = create_test_flash(test_instrument_id, test_chain, test_dex, 400, 20, 8);
830        let pos = get_event_block_position(&DexPoolData::Flash(flash));
831        assert_eq!(pos, (400, 20, 8));
832    }
833
834    #[rstest]
835    fn test_get_event_block_position_fee_protocol_update(
836        test_instrument_id: InstrumentId,
837        test_chain: Arc<Chain>,
838        test_dex: Arc<Dex>,
839    ) {
840        let update =
841            create_test_fee_protocol_update(test_instrument_id, test_chain, test_dex, 500, 25, 4);
842        let pos = get_event_block_position(&DexPoolData::FeeProtocolUpdate(update));
843        assert_eq!(pos, (500, 25, 4));
844    }
845
846    #[rstest]
847    fn test_convert_and_sort_empty_events() {
848        let events = convert_and_sort_buffered_events(vec![]);
849        assert!(events.is_empty());
850    }
851
852    #[rstest]
853    fn test_convert_and_sort_filters_non_pool_events(
854        test_instrument_id: InstrumentId,
855        test_chain: Arc<Chain>,
856        test_dex: Arc<Dex>,
857    ) {
858        let events = vec![
859            DefiData::PoolSwap(create_test_swap(
860                test_instrument_id,
861                test_chain,
862                test_dex,
863                100,
864                0,
865                0,
866            )),
867            // Block events would be filtered out
868        ];
869        let sorted = convert_and_sort_buffered_events(events);
870        assert_eq!(sorted.len(), 1);
871    }
872
873    #[rstest]
874    fn test_convert_and_sort_single_event(
875        test_instrument_id: InstrumentId,
876        test_chain: Arc<Chain>,
877        test_dex: Arc<Dex>,
878    ) {
879        let swap = create_test_swap(test_instrument_id, test_chain, test_dex, 100, 5, 3);
880        let events = vec![DefiData::PoolSwap(swap)];
881        let sorted = convert_and_sort_buffered_events(events);
882        assert_eq!(sorted.len(), 1);
883        assert_eq!(get_event_block_position(&sorted[0]), (100, 5, 3));
884    }
885
886    #[rstest]
887    fn test_convert_and_sort_already_sorted(
888        test_instrument_id: InstrumentId,
889        test_chain: Arc<Chain>,
890        test_dex: Arc<Dex>,
891    ) {
892        let events = vec![
893            DefiData::PoolSwap(create_test_swap(
894                test_instrument_id,
895                test_chain.clone(),
896                test_dex.clone(),
897                100,
898                0,
899                0,
900            )),
901            DefiData::PoolSwap(create_test_swap(
902                test_instrument_id,
903                test_chain.clone(),
904                test_dex.clone(),
905                100,
906                0,
907                1,
908            )),
909            DefiData::PoolSwap(create_test_swap(
910                test_instrument_id,
911                test_chain,
912                test_dex,
913                100,
914                1,
915                0,
916            )),
917        ];
918        let sorted = convert_and_sort_buffered_events(events);
919        assert_eq!(sorted.len(), 3);
920        assert_eq!(get_event_block_position(&sorted[0]), (100, 0, 0));
921        assert_eq!(get_event_block_position(&sorted[1]), (100, 0, 1));
922        assert_eq!(get_event_block_position(&sorted[2]), (100, 1, 0));
923    }
924
925    #[rstest]
926    fn test_convert_and_sort_reverse_order(
927        test_instrument_id: InstrumentId,
928        test_chain: Arc<Chain>,
929        test_dex: Arc<Dex>,
930    ) {
931        let events = vec![
932            DefiData::PoolSwap(create_test_swap(
933                test_instrument_id,
934                test_chain.clone(),
935                test_dex.clone(),
936                100,
937                2,
938                5,
939            )),
940            DefiData::PoolSwap(create_test_swap(
941                test_instrument_id,
942                test_chain.clone(),
943                test_dex.clone(),
944                100,
945                1,
946                3,
947            )),
948            DefiData::PoolSwap(create_test_swap(
949                test_instrument_id,
950                test_chain,
951                test_dex,
952                100,
953                0,
954                1,
955            )),
956        ];
957        let sorted = convert_and_sort_buffered_events(events);
958        assert_eq!(sorted.len(), 3);
959        assert_eq!(get_event_block_position(&sorted[0]), (100, 0, 1));
960        assert_eq!(get_event_block_position(&sorted[1]), (100, 1, 3));
961        assert_eq!(get_event_block_position(&sorted[2]), (100, 2, 5));
962    }
963
964    #[rstest]
965    fn test_convert_and_sort_mixed_blocks(
966        test_instrument_id: InstrumentId,
967        test_chain: Arc<Chain>,
968        test_dex: Arc<Dex>,
969    ) {
970        let events = vec![
971            DefiData::PoolSwap(create_test_swap(
972                test_instrument_id,
973                test_chain.clone(),
974                test_dex.clone(),
975                102,
976                0,
977                0,
978            )),
979            DefiData::PoolSwap(create_test_swap(
980                test_instrument_id,
981                test_chain.clone(),
982                test_dex.clone(),
983                100,
984                5,
985                2,
986            )),
987            DefiData::PoolSwap(create_test_swap(
988                test_instrument_id,
989                test_chain,
990                test_dex,
991                101,
992                3,
993                1,
994            )),
995        ];
996        let sorted = convert_and_sort_buffered_events(events);
997        assert_eq!(sorted.len(), 3);
998        assert_eq!(get_event_block_position(&sorted[0]), (100, 5, 2));
999        assert_eq!(get_event_block_position(&sorted[1]), (101, 3, 1));
1000        assert_eq!(get_event_block_position(&sorted[2]), (102, 0, 0));
1001    }
1002
1003    #[rstest]
1004    fn test_convert_and_sort_mixed_event_types(
1005        test_instrument_id: InstrumentId,
1006        test_chain: Arc<Chain>,
1007        test_dex: Arc<Dex>,
1008    ) {
1009        let events = vec![
1010            DefiData::PoolSwap(create_test_swap(
1011                test_instrument_id,
1012                test_chain.clone(),
1013                test_dex.clone(),
1014                100,
1015                2,
1016                0,
1017            )),
1018            DefiData::PoolLiquidityUpdate(create_test_liquidity_update(
1019                test_instrument_id,
1020                test_chain.clone(),
1021                test_dex.clone(),
1022                100,
1023                0,
1024                0,
1025            )),
1026            DefiData::PoolFeeCollect(create_test_fee_collect(
1027                test_instrument_id,
1028                test_chain.clone(),
1029                test_dex.clone(),
1030                100,
1031                1,
1032                0,
1033            )),
1034            DefiData::PoolFlash(create_test_flash(
1035                test_instrument_id,
1036                test_chain.clone(),
1037                test_dex.clone(),
1038                100,
1039                3,
1040                0,
1041            )),
1042            DefiData::PoolFeeProtocolUpdate(create_test_fee_protocol_update(
1043                test_instrument_id,
1044                test_chain,
1045                test_dex,
1046                100,
1047                4,
1048                0,
1049            )),
1050        ];
1051        let sorted = convert_and_sort_buffered_events(events);
1052        assert_eq!(sorted.len(), 5);
1053        assert_eq!(get_event_block_position(&sorted[0]), (100, 0, 0));
1054        assert_eq!(get_event_block_position(&sorted[1]), (100, 1, 0));
1055        assert_eq!(get_event_block_position(&sorted[2]), (100, 2, 0));
1056        assert_eq!(get_event_block_position(&sorted[3]), (100, 3, 0));
1057        assert_eq!(get_event_block_position(&sorted[4]), (100, 4, 0));
1058        assert!(matches!(sorted[4], DexPoolData::FeeProtocolUpdate(_)));
1059    }
1060
1061    #[rstest]
1062    fn test_convert_and_sort_same_block_and_tx_different_log_index(
1063        test_instrument_id: InstrumentId,
1064        test_chain: Arc<Chain>,
1065        test_dex: Arc<Dex>,
1066    ) {
1067        let events = vec![
1068            DefiData::PoolSwap(create_test_swap(
1069                test_instrument_id,
1070                test_chain.clone(),
1071                test_dex.clone(),
1072                100,
1073                5,
1074                10,
1075            )),
1076            DefiData::PoolSwap(create_test_swap(
1077                test_instrument_id,
1078                test_chain.clone(),
1079                test_dex.clone(),
1080                100,
1081                5,
1082                5,
1083            )),
1084            DefiData::PoolSwap(create_test_swap(
1085                test_instrument_id,
1086                test_chain,
1087                test_dex,
1088                100,
1089                5,
1090                1,
1091            )),
1092        ];
1093        let sorted = convert_and_sort_buffered_events(events);
1094        assert_eq!(sorted.len(), 3);
1095        assert_eq!(get_event_block_position(&sorted[0]), (100, 5, 1));
1096        assert_eq!(get_event_block_position(&sorted[1]), (100, 5, 5));
1097        assert_eq!(get_event_block_position(&sorted[2]), (100, 5, 10));
1098    }
1099}