nautilus_backtest/
data_iterator.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//! Multi-stream, time-ordered data iterator for replaying historical data.
17
18use std::collections::BinaryHeap;
19
20use ahash::AHashMap;
21use nautilus_core::UnixNanos;
22use nautilus_model::data::{Data, HasTsInit};
23
24#[cfg(feature = "defi")]
25use crate::defi::replay::replay_position;
26
27// TODO: block_number/transaction_index/log_index/phase are DeFi-only (zero for all other data,
28// even in non-DeFi builds); they exist to order same-block DeFi events in canonical chain order.
29// This leaks DeFi-specific shape into a general key, so it could be cfg-gated or moved behind an
30// opaque secondary key later (non-breaking, no correctness or perf cost).
31#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd)]
32struct ReplayKey {
33    ts: UnixNanos,
34    block_number: u64,
35    transaction_index: u32,
36    log_index: u32,
37    phase: u8,
38}
39
40fn replay_key(data: &Data) -> ReplayKey {
41    match data {
42        #[cfg(feature = "defi")]
43        Data::Defi(defi) => {
44            let (block_number, transaction_index, log_index, phase) = replay_position(defi);
45            ReplayKey {
46                ts: defi.ts_init(),
47                block_number,
48                transaction_index,
49                log_index,
50                phase,
51            }
52        }
53        _ => ReplayKey {
54            ts: data.ts_init(),
55            block_number: 0,
56            transaction_index: 0,
57            log_index: 0,
58            phase: 0,
59        },
60    }
61}
62
63/// Internal convenience struct to keep heap entries ordered by replay key and priority.
64#[derive(Debug, Eq, PartialEq)]
65struct HeapEntry {
66    key: ReplayKey,
67    priority: i32,
68    index: usize,
69}
70
71impl Ord for HeapEntry {
72    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
73        // min-heap on replay key, then priority sign (+/-) then index
74        self.key
75            .cmp(&other.key)
76            .then_with(|| self.priority.cmp(&other.priority))
77            .then_with(|| self.index.cmp(&other.index))
78            .reverse() // BinaryHeap is max by default -> reverse for min behaviour
79    }
80}
81
82impl PartialOrd for HeapEntry {
83    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
84        Some(self.cmp(other))
85    }
86}
87
88/// Multi-stream, time-ordered data iterator used by the backtest engine.
89#[derive(Debug, Default)]
90pub struct BacktestDataIterator {
91    streams: AHashMap<i32, Vec<Data>>, // key: priority, value: Vec<Data>
92    names: AHashMap<i32, String>,      // priority -> name
93    priorities: AHashMap<String, i32>, // name -> priority
94    indices: AHashMap<i32, usize>,     // cursor per stream
95    heap: BinaryHeap<HeapEntry>,
96    single_priority: Option<i32>,
97    next_priority_counter: i32, // monotonically increasing counter used to assign priorities
98}
99
100impl BacktestDataIterator {
101    /// Creates a new empty [`BacktestDataIterator`].
102    #[must_use]
103    pub fn new() -> Self {
104        Self {
105            streams: AHashMap::new(),
106            names: AHashMap::new(),
107            priorities: AHashMap::new(),
108            indices: AHashMap::new(),
109            heap: BinaryHeap::new(),
110            single_priority: None,
111            next_priority_counter: 0,
112        }
113    }
114
115    /// Adds (or replaces) a named data stream.
116    ///
117    /// When `append_data` is true the stream gets lower priority on timestamp
118    /// ties; when false (prepend) it wins ties.
119    pub fn add_data(&mut self, name: &str, mut data: Vec<Data>, append_data: bool) {
120        if data.is_empty() {
121            return;
122        }
123
124        data.sort_by_key(replay_key);
125
126        self.add_stream(name, data, append_data);
127    }
128
129    fn add_stream(&mut self, name: &str, data: Vec<Data>, append_data: bool) {
130        let priority = if let Some(p) = self.priorities.get(name) {
131            // Replace existing stream – remove previous traces then re-insert below.
132            *p
133        } else {
134            self.next_priority_counter += 1;
135            let sign = if append_data { 1 } else { -1 };
136            sign * self.next_priority_counter
137        };
138
139        // Remove old state if any
140        self.remove_data(name, true);
141
142        self.streams.insert(priority, data);
143        self.names.insert(priority, name.to_string());
144        self.priorities.insert(name.to_string(), priority);
145        self.indices.insert(priority, 0);
146
147        self.rebuild_heap();
148    }
149
150    /// Removes a named data stream.
151    pub fn remove_data(&mut self, name: &str, complete_remove: bool) {
152        if let Some(priority) = self.priorities.remove(name) {
153            self.streams.remove(&priority);
154            self.indices.remove(&priority);
155            self.names.remove(&priority);
156
157            // Rebuild heap sans removed priority
158            self.heap.retain(|e| e.priority != priority);
159
160            if self.heap.is_empty() {
161                self.single_priority = None;
162            }
163        }
164
165        if complete_remove {
166            // Placeholder for future generator cleanup
167        }
168    }
169
170    /// Sets the cursor of a named stream to `index` (0-based).
171    pub fn set_index(&mut self, name: &str, index: usize) {
172        if let Some(priority) = self.priorities.get(name) {
173            self.indices.insert(*priority, index);
174            self.rebuild_heap();
175        }
176    }
177
178    /// Resets all stream cursors to the beginning.
179    pub fn reset_all_cursors(&mut self) {
180        for idx in self.indices.values_mut() {
181            *idx = 0;
182        }
183        self.rebuild_heap();
184    }
185
186    /// Returns the next backtest data element across all streams in replay order.
187    pub(crate) fn next_item(&mut self) -> Option<Data> {
188        // Fast path for single stream
189        if let Some(p) = self.single_priority {
190            let data = self.streams.get_mut(&p)?;
191            let idx = self.indices.get_mut(&p)?;
192            if *idx >= data.len() {
193                return None;
194            }
195            let element = data[*idx].clone();
196            *idx += 1;
197            return Some(element);
198        }
199
200        // Multi-stream path using heap
201        let entry = self.heap.pop()?;
202        let stream_vec = self.streams.get(&entry.priority)?;
203        let element = stream_vec[entry.index].clone();
204
205        // Advance cursor and push next entry
206        let next_index = entry.index + 1;
207        self.indices.insert(entry.priority, next_index);
208        if next_index < stream_vec.len() {
209            self.heap.push(HeapEntry {
210                key: replay_key(&stream_vec[next_index]),
211                priority: entry.priority,
212                index: next_index,
213            });
214        }
215
216        Some(element)
217    }
218
219    /// Returns the next market [`Data`] element across all streams in chronological order.
220    #[expect(clippy::should_implement_trait)]
221    pub fn next(&mut self) -> Option<Data> {
222        self.next_item()
223    }
224
225    /// Returns whether all streams have been fully consumed.
226    #[must_use]
227    pub fn is_done(&self) -> bool {
228        if let Some(p) = self.single_priority {
229            if let Some(idx) = self.indices.get(&p)
230                && let Some(vec) = self.streams.get(&p)
231            {
232                return *idx >= vec.len();
233            }
234            true
235        } else {
236            self.heap.is_empty()
237        }
238    }
239
240    fn rebuild_heap(&mut self) {
241        self.heap.clear();
242
243        // Determine if we’re in single-stream mode
244        if self.streams.len() == 1 {
245            self.single_priority = self.streams.keys().next().copied();
246            return;
247        }
248        self.single_priority = None;
249
250        for (&priority, vec) in &self.streams {
251            let idx = *self.indices.get(&priority).unwrap_or(&0);
252            if idx < vec.len() {
253                self.heap.push(HeapEntry {
254                    key: replay_key(&vec[idx]),
255                    priority,
256                    index: idx,
257                });
258            }
259        }
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use nautilus_model::{
266        data::QuoteTick,
267        identifiers::InstrumentId,
268        types::{Price, Quantity},
269    };
270    #[cfg(feature = "defi")]
271    use nautilus_model::{
272        defi::{
273            DefiData,
274            data::block::BlockPosition,
275            pool_analysis::snapshot::{PoolAnalytics, PoolSnapshot, PoolState},
276        },
277        identifiers::{Symbol, Venue},
278    };
279    use rstest::rstest;
280
281    use super::*;
282
283    fn quote(id: &str, ts: u64) -> Data {
284        let inst = InstrumentId::from(id);
285        Data::Quote(QuoteTick::new(
286            inst,
287            Price::from("1.0"),
288            Price::from("1.0"),
289            Quantity::from(100),
290            Quantity::from(100),
291            ts.into(),
292            ts.into(),
293        ))
294    }
295
296    fn collect_ts(it: &mut BacktestDataIterator) -> Vec<u64> {
297        let mut ts = Vec::new();
298        while let Some(d) = it.next() {
299            ts.push(d.ts_init().as_u64());
300        }
301        ts
302    }
303
304    #[cfg(feature = "defi")]
305    fn defi_snapshot(ts: u64, block: u64, transaction_index: u32, log_index: u32) -> Data {
306        let instrument_id = InstrumentId::new(Symbol::from("ETH/USDC"), Venue::from("UNISWAPV3"));
307        let snapshot = PoolSnapshot::new(
308            instrument_id,
309            PoolState::default(),
310            Vec::new(),
311            Vec::new(),
312            PoolAnalytics::default(),
313            BlockPosition::new(block, format!("0x{block:x}"), transaction_index, log_index),
314            UnixNanos::from(ts),
315            UnixNanos::from(ts),
316        );
317
318        Data::Defi(Box::new(DefiData::PoolSnapshot(snapshot)))
319    }
320
321    #[rstest]
322    fn test_single_stream_yields_in_order() {
323        let mut it = BacktestDataIterator::new();
324        it.add_data(
325            "s",
326            vec![quote("A.B", 100), quote("A.B", 200), quote("A.B", 300)],
327            true,
328        );
329
330        assert_eq!(collect_ts(&mut it), vec![100, 200, 300]);
331        assert!(it.is_done());
332    }
333
334    #[rstest]
335    fn test_single_stream_exhaustion_returns_none() {
336        let mut it = BacktestDataIterator::new();
337        it.add_data("s", vec![quote("A.B", 1), quote("A.B", 3)], true);
338        assert_eq!(it.next().unwrap().ts_init(), UnixNanos::from(1));
339        assert_eq!(it.next().unwrap().ts_init(), UnixNanos::from(3));
340        assert!(it.next().is_none());
341    }
342
343    #[rstest]
344    fn test_single_stream_sorts_unsorted_input() {
345        let mut it = BacktestDataIterator::new();
346        it.add_data(
347            "s",
348            vec![quote("A.B", 300), quote("A.B", 100), quote("A.B", 200)],
349            true,
350        );
351
352        assert_eq!(collect_ts(&mut it), vec![100, 200, 300]);
353    }
354
355    #[rstest]
356    fn test_two_stream_merge_chronological() {
357        let mut it = BacktestDataIterator::new();
358        it.add_data("s1", vec![quote("A.B", 1), quote("A.B", 4)], true);
359        it.add_data("s2", vec![quote("C.D", 2), quote("C.D", 3)], false);
360
361        assert_eq!(collect_ts(&mut it), vec![1, 2, 3, 4]);
362    }
363
364    #[rstest]
365    fn test_three_stream_merge_sorted() {
366        let mut it = BacktestDataIterator::new();
367        let data_len = 5;
368        let d0: Vec<Data> = (0..data_len).map(|k| quote("A.B", 3 * k)).collect();
369        let d1: Vec<Data> = (0..data_len).map(|k| quote("C.D", 3 * k + 1)).collect();
370        let d2: Vec<Data> = (0..data_len).map(|k| quote("E.F", 3 * k + 2)).collect();
371        it.add_data("d0", d0, true);
372        it.add_data("d1", d1, true);
373        it.add_data("d2", d2, true);
374
375        let ts = collect_ts(&mut it);
376        assert_eq!(ts.len(), 15);
377        for i in 0..ts.len() - 1 {
378            assert!(ts[i] <= ts[i + 1], "Not sorted at index {i}");
379        }
380    }
381
382    #[rstest]
383    fn test_multiple_streams_merge_order() {
384        let mut it = BacktestDataIterator::new();
385        it.add_data("s1", vec![quote("A.B", 100), quote("A.B", 300)], true);
386        it.add_data("s2", vec![quote("C.D", 200), quote("C.D", 400)], true);
387
388        assert_eq!(collect_ts(&mut it), vec![100, 200, 300, 400]);
389    }
390
391    #[rstest]
392    fn test_append_data_priority_default_fifo() {
393        let mut it = BacktestDataIterator::new();
394        it.add_data("a", vec![quote("A.B", 100)], true);
395        it.add_data("b", vec![quote("C.D", 100)], true);
396
397        // Both at same timestamp, FIFO order (a before b)
398        let ts = collect_ts(&mut it);
399        assert_eq!(ts, vec![100, 100]);
400    }
401
402    #[rstest]
403    fn test_prepend_priority_wins_ties() {
404        let mut it = BacktestDataIterator::new();
405        // "a" is appended (lower priority), "b" is prepended (higher priority)
406        it.add_data("a", vec![quote("A.B", 100)], true);
407        it.add_data("b", vec![quote("C.D", 100)], false);
408
409        // "b" (prepend) should come first despite being added second
410        let first = it.next().unwrap();
411        let second = it.next().unwrap();
412        // Prepend stream (negative priority) wins ties over append (positive)
413        assert_eq!(first.instrument_id(), InstrumentId::from("C.D"));
414        assert_eq!(second.instrument_id(), InstrumentId::from("A.B"));
415    }
416
417    #[rstest]
418    fn test_is_done_empty_iterator() {
419        let it = BacktestDataIterator::new();
420        assert!(it.is_done());
421    }
422
423    #[rstest]
424    fn test_is_done_after_consumption() {
425        let mut it = BacktestDataIterator::new();
426        it.add_data("s", vec![quote("A.B", 1)], true);
427
428        assert!(!it.is_done());
429        it.next();
430        assert!(it.is_done());
431    }
432
433    #[rstest]
434    fn test_is_done_multi_stream() {
435        let mut it = BacktestDataIterator::new();
436        it.add_data("s1", vec![quote("A.B", 1)], true);
437        it.add_data("s2", vec![quote("C.D", 2)], true);
438
439        assert!(!it.is_done());
440        it.next();
441        assert!(!it.is_done());
442        it.next();
443        assert!(it.is_done());
444    }
445
446    #[rstest]
447    fn test_partial_consumption_then_complete() {
448        let mut it = BacktestDataIterator::new();
449        it.add_data(
450            "s",
451            vec![
452                quote("A.B", 0),
453                quote("A.B", 1),
454                quote("A.B", 2),
455                quote("A.B", 3),
456            ],
457            true,
458        );
459
460        assert_eq!(it.next().unwrap().ts_init().as_u64(), 0);
461        assert_eq!(it.next().unwrap().ts_init().as_u64(), 1);
462
463        let remaining = collect_ts(&mut it);
464        assert_eq!(remaining, vec![2, 3]);
465        assert!(it.is_done());
466    }
467
468    #[rstest]
469    fn test_remove_stream_reduces_output() {
470        let mut it = BacktestDataIterator::new();
471        it.add_data("a", vec![quote("A.B", 1)], true);
472        it.add_data("b", vec![quote("C.D", 2)], true);
473
474        it.remove_data("a", false);
475
476        assert_eq!(collect_ts(&mut it), vec![2]);
477    }
478
479    #[rstest]
480    fn test_remove_all_streams_yields_empty() {
481        let mut it = BacktestDataIterator::new();
482        it.add_data("x", vec![quote("A.B", 1)], true);
483        it.add_data("y", vec![quote("C.D", 2)], true);
484
485        it.remove_data("x", false);
486        it.remove_data("y", false);
487
488        assert!(it.next().is_none());
489        assert!(it.is_done());
490    }
491
492    #[rstest]
493    fn test_remove_nonexistent_stream_is_noop() {
494        let mut it = BacktestDataIterator::new();
495        it.add_data("s", vec![quote("A.B", 1)], true);
496
497        it.remove_data("nonexistent", false);
498
499        assert_eq!(collect_ts(&mut it), vec![1]);
500    }
501
502    #[rstest]
503    fn test_remove_after_full_consumption() {
504        let mut it = BacktestDataIterator::new();
505        it.add_data("s", vec![quote("A.B", 1), quote("A.B", 2)], true);
506
507        collect_ts(&mut it);
508
509        it.remove_data("s", true);
510        assert!(it.is_done());
511    }
512
513    #[rstest]
514    fn test_set_index_rewinds_stream() {
515        let mut it = BacktestDataIterator::new();
516        it.add_data(
517            "s",
518            vec![quote("A.B", 10), quote("A.B", 20), quote("A.B", 30)],
519            true,
520        );
521
522        assert_eq!(it.next().unwrap().ts_init().as_u64(), 10);
523
524        it.set_index("s", 0);
525
526        assert_eq!(collect_ts(&mut it), vec![10, 20, 30]);
527    }
528
529    #[rstest]
530    fn test_set_index_skips_forward() {
531        let mut it = BacktestDataIterator::new();
532        it.add_data(
533            "s",
534            vec![quote("A.B", 10), quote("A.B", 20), quote("A.B", 30)],
535            true,
536        );
537
538        it.set_index("s", 2);
539
540        assert_eq!(collect_ts(&mut it), vec![30]);
541    }
542
543    #[rstest]
544    fn test_set_index_nonexistent_stream_is_noop() {
545        let mut it = BacktestDataIterator::new();
546        it.add_data("s", vec![quote("A.B", 1)], true);
547
548        it.set_index("nonexistent", 0);
549
550        assert_eq!(collect_ts(&mut it), vec![1]);
551    }
552
553    #[rstest]
554    fn test_reset_all_cursors_single_stream() {
555        let mut it = BacktestDataIterator::new();
556        it.add_data("s", vec![quote("A.B", 1), quote("A.B", 2)], true);
557
558        collect_ts(&mut it);
559        assert!(it.is_done());
560
561        it.reset_all_cursors();
562        assert!(!it.is_done());
563        assert_eq!(collect_ts(&mut it), vec![1, 2]);
564    }
565
566    #[rstest]
567    fn test_reset_all_cursors_multi_stream() {
568        let mut it = BacktestDataIterator::new();
569        it.add_data("s1", vec![quote("A.B", 1), quote("A.B", 3)], true);
570        it.add_data("s2", vec![quote("C.D", 2), quote("C.D", 4)], true);
571
572        collect_ts(&mut it);
573        assert!(it.is_done());
574
575        it.reset_all_cursors();
576        assert_eq!(collect_ts(&mut it), vec![1, 2, 3, 4]);
577    }
578
579    #[rstest]
580    fn test_readding_data_replaces_stream() {
581        let mut it = BacktestDataIterator::new();
582        it.add_data("X", vec![quote("A.B", 1), quote("A.B", 2)], true);
583        it.add_data("X", vec![quote("A.B", 10)], true);
584
585        assert_eq!(collect_ts(&mut it), vec![10]);
586    }
587
588    #[rstest]
589    fn test_add_empty_data_is_noop() {
590        let mut it = BacktestDataIterator::new();
591        it.add_data("empty", vec![], true);
592
593        assert!(it.is_done());
594        assert!(it.next().is_none());
595    }
596
597    #[rstest]
598    fn test_empty_iterator_returns_none() {
599        let mut it = BacktestDataIterator::new();
600        assert!(it.next().is_none());
601        assert!(it.is_done());
602    }
603
604    #[rstest]
605    fn test_multiple_add_data_calls_with_different_names() {
606        let mut it = BacktestDataIterator::new();
607        it.add_data("batch_0", vec![quote("A.B", 1), quote("A.B", 3)], true);
608        it.add_data("batch_1", vec![quote("A.B", 2), quote("A.B", 4)], true);
609
610        assert_eq!(collect_ts(&mut it), vec![1, 2, 3, 4]);
611    }
612
613    #[rstest]
614    fn test_prepend_stream_always_wins_ties_across_batches() {
615        // Verifies that a prepend stream (negative priority) wins ties
616        // even when added after multiple append streams
617        let mut it = BacktestDataIterator::new();
618        it.add_data("append_a", vec![quote("A.B", 100)], true);
619        it.add_data("append_b", vec![quote("C.D", 100)], true);
620        it.add_data("prepend", vec![quote("E.F", 100)], false);
621
622        let first = it.next().unwrap();
623        assert_eq!(
624            first.instrument_id(),
625            InstrumentId::from("E.F"),
626            "Prepend stream should always come first in ties"
627        );
628    }
629
630    #[rstest]
631    fn test_equal_timestamps_across_many_streams_preserves_priority_order() {
632        // All items at the same timestamp — ordering is strictly by priority
633        let mut it = BacktestDataIterator::new();
634        it.add_data("s1", vec![quote("A.B", 50)], true);
635        it.add_data("s2", vec![quote("C.D", 50)], true);
636        it.add_data("s3", vec![quote("E.F", 50)], true);
637        it.add_data("s4", vec![quote("G.H", 50)], true);
638
639        let mut ids = Vec::new();
640        while let Some(d) = it.next() {
641            ids.push(d.instrument_id());
642        }
643
644        assert_eq!(ids.len(), 4);
645
646        // All should be yielded (no duplicates dropped, no items lost)
647        assert!(ids.contains(&InstrumentId::from("A.B")));
648        assert!(ids.contains(&InstrumentId::from("C.D")));
649        assert!(ids.contains(&InstrumentId::from("E.F")));
650        assert!(ids.contains(&InstrumentId::from("G.H")));
651    }
652
653    #[cfg(feature = "defi")]
654    #[rstest]
655    fn test_defi_data_orders_equal_timestamps_by_block_position() {
656        let mut it = BacktestDataIterator::new();
657        it.add_data(
658            "defi",
659            vec![
660                defi_snapshot(100, 12, 4, 1),
661                defi_snapshot(100, 11, 9, 9),
662                defi_snapshot(100, 12, 2, 7),
663            ],
664            true,
665        );
666
667        let mut positions = Vec::new();
668        while let Some(Data::Defi(data)) = it.next_item() {
669            positions.push(data.block_position());
670        }
671
672        assert_eq!(positions, vec![(11, 9, 9), (12, 2, 7), (12, 4, 1)]);
673    }
674}
nautilus_backtest/data_iterator.rs

nautilus_backtest/
data_iterator.rs
