1use 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#[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#[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 self.key
75 .cmp(&other.key)
76 .then_with(|| self.priority.cmp(&other.priority))
77 .then_with(|| self.index.cmp(&other.index))
78 .reverse() }
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#[derive(Debug, Default)]
90pub struct BacktestDataIterator {
91 streams: AHashMap<i32, Vec<Data>>, names: AHashMap<i32, String>, priorities: AHashMap<String, i32>, indices: AHashMap<i32, usize>, heap: BinaryHeap<HeapEntry>,
96 single_priority: Option<i32>,
97 next_priority_counter: i32, }
99
100impl BacktestDataIterator {
101 #[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 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 *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 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 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 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 }
168 }
169
170 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 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 pub(crate) fn peek(&self) -> Option<&Data> {
188 if let Some(p) = self.single_priority {
189 let data = self.streams.get(&p)?;
190 let idx = *self.indices.get(&p)?;
191 return data.get(idx);
192 }
193
194 let entry = self.heap.peek()?;
195 self.streams.get(&entry.priority)?.get(entry.index)
196 }
197
198 pub(crate) fn next_item(&mut self) -> Option<Data> {
200 if let Some(p) = self.single_priority {
202 let data = self.streams.get_mut(&p)?;
203 let idx = self.indices.get_mut(&p)?;
204 if *idx >= data.len() {
205 return None;
206 }
207 let element = data[*idx].clone();
208 *idx += 1;
209 return Some(element);
210 }
211
212 let entry = self.heap.pop()?;
214 let stream_vec = self.streams.get(&entry.priority)?;
215 let element = stream_vec[entry.index].clone();
216
217 let next_index = entry.index + 1;
219 self.indices.insert(entry.priority, next_index);
220 if next_index < stream_vec.len() {
221 self.heap.push(HeapEntry {
222 key: replay_key(&stream_vec[next_index]),
223 priority: entry.priority,
224 index: next_index,
225 });
226 }
227
228 Some(element)
229 }
230
231 #[expect(clippy::should_implement_trait)]
233 pub fn next(&mut self) -> Option<Data> {
234 self.next_item()
235 }
236
237 #[must_use]
239 pub fn is_done(&self) -> bool {
240 if let Some(p) = self.single_priority {
241 if let Some(idx) = self.indices.get(&p)
242 && let Some(vec) = self.streams.get(&p)
243 {
244 return *idx >= vec.len();
245 }
246 true
247 } else {
248 self.heap.is_empty()
249 }
250 }
251
252 fn rebuild_heap(&mut self) {
253 self.heap.clear();
254
255 if self.streams.len() == 1 {
257 self.single_priority = self.streams.keys().next().copied();
258 return;
259 }
260 self.single_priority = None;
261
262 for (&priority, vec) in &self.streams {
263 let idx = *self.indices.get(&priority).unwrap_or(&0);
264 if idx < vec.len() {
265 self.heap.push(HeapEntry {
266 key: replay_key(&vec[idx]),
267 priority,
268 index: idx,
269 });
270 }
271 }
272 }
273}
274
275#[cfg(test)]
276mod tests {
277 use nautilus_model::{
278 data::QuoteTick,
279 identifiers::InstrumentId,
280 types::{Price, Quantity},
281 };
282 #[cfg(feature = "defi")]
283 use nautilus_model::{
284 defi::{
285 DefiData,
286 data::block::BlockPosition,
287 pool_analysis::snapshot::{PoolAnalytics, PoolSnapshot, PoolState},
288 },
289 identifiers::{Symbol, Venue},
290 };
291 use rstest::rstest;
292
293 use super::*;
294
295 fn quote(id: &str, ts: u64) -> Data {
296 let inst = InstrumentId::from(id);
297 Data::Quote(QuoteTick::new(
298 inst,
299 Price::from("1.0"),
300 Price::from("1.0"),
301 Quantity::from(100),
302 Quantity::from(100),
303 ts.into(),
304 ts.into(),
305 ))
306 }
307
308 fn collect_ts(it: &mut BacktestDataIterator) -> Vec<u64> {
309 let mut ts = Vec::new();
310 while let Some(d) = it.next() {
311 ts.push(d.ts_init().as_u64());
312 }
313 ts
314 }
315
316 #[cfg(feature = "defi")]
317 fn defi_snapshot(ts: u64, block: u64, transaction_index: u32, log_index: u32) -> Data {
318 let instrument_id = InstrumentId::new(Symbol::from("ETH/USDC"), Venue::from("UNISWAPV3"));
319 let snapshot = PoolSnapshot::new(
320 instrument_id,
321 PoolState::default(),
322 Vec::new(),
323 Vec::new(),
324 PoolAnalytics::default(),
325 BlockPosition::new(block, format!("0x{block:x}"), transaction_index, log_index),
326 UnixNanos::from(ts),
327 UnixNanos::from(ts),
328 );
329
330 Data::Defi(Box::new(DefiData::PoolSnapshot(snapshot)))
331 }
332
333 #[rstest]
334 fn test_single_stream_yields_in_order() {
335 let mut it = BacktestDataIterator::new();
336 it.add_data(
337 "s",
338 vec![quote("A.B", 100), quote("A.B", 200), quote("A.B", 300)],
339 true,
340 );
341
342 assert_eq!(collect_ts(&mut it), vec![100, 200, 300]);
343 assert!(it.is_done());
344 }
345
346 #[rstest]
347 fn test_single_stream_exhaustion_returns_none() {
348 let mut it = BacktestDataIterator::new();
349 it.add_data("s", vec![quote("A.B", 1), quote("A.B", 3)], true);
350 assert_eq!(it.next().unwrap().ts_init(), UnixNanos::from(1));
351 assert_eq!(it.next().unwrap().ts_init(), UnixNanos::from(3));
352 assert!(it.next().is_none());
353 }
354
355 #[rstest]
356 fn test_peek_does_not_consume_single_stream_item() {
357 let mut it = BacktestDataIterator::new();
358 it.add_data("s", vec![quote("A.B", 1), quote("A.B", 2)], true);
359
360 assert_eq!(it.peek().unwrap().ts_init(), UnixNanos::from(1));
361 assert_eq!(it.peek().unwrap().ts_init(), UnixNanos::from(1));
362 assert_eq!(it.next().unwrap().ts_init(), UnixNanos::from(1));
363 assert_eq!(it.peek().unwrap().ts_init(), UnixNanos::from(2));
364 }
365
366 #[rstest]
367 fn test_single_stream_sorts_unsorted_input() {
368 let mut it = BacktestDataIterator::new();
369 it.add_data(
370 "s",
371 vec![quote("A.B", 300), quote("A.B", 100), quote("A.B", 200)],
372 true,
373 );
374
375 assert_eq!(collect_ts(&mut it), vec![100, 200, 300]);
376 }
377
378 #[rstest]
379 fn test_two_stream_merge_chronological() {
380 let mut it = BacktestDataIterator::new();
381 it.add_data("s1", vec![quote("A.B", 1), quote("A.B", 4)], true);
382 it.add_data("s2", vec![quote("C.D", 2), quote("C.D", 3)], false);
383
384 assert_eq!(collect_ts(&mut it), vec![1, 2, 3, 4]);
385 }
386
387 #[rstest]
388 fn test_peek_does_not_consume_multi_stream_heap_item() {
389 let mut it = BacktestDataIterator::new();
390 it.add_data("s1", vec![quote("A.B", 1), quote("A.B", 4)], true);
391 it.add_data("s2", vec![quote("C.D", 2), quote("C.D", 3)], true);
392
393 assert_eq!(it.peek().unwrap().ts_init(), UnixNanos::from(1));
394 assert_eq!(it.peek().unwrap().ts_init(), UnixNanos::from(1));
395 assert_eq!(it.next().unwrap().ts_init(), UnixNanos::from(1));
396 assert_eq!(it.peek().unwrap().ts_init(), UnixNanos::from(2));
397 assert_eq!(collect_ts(&mut it), vec![2, 3, 4]);
398 }
399
400 #[rstest]
401 fn test_three_stream_merge_sorted() {
402 let mut it = BacktestDataIterator::new();
403 let data_len = 5;
404 let d0: Vec<Data> = (0..data_len).map(|k| quote("A.B", 3 * k)).collect();
405 let d1: Vec<Data> = (0..data_len).map(|k| quote("C.D", 3 * k + 1)).collect();
406 let d2: Vec<Data> = (0..data_len).map(|k| quote("E.F", 3 * k + 2)).collect();
407 it.add_data("d0", d0, true);
408 it.add_data("d1", d1, true);
409 it.add_data("d2", d2, true);
410
411 let ts = collect_ts(&mut it);
412 assert_eq!(ts.len(), 15);
413 for i in 0..ts.len() - 1 {
414 assert!(ts[i] <= ts[i + 1], "Not sorted at index {i}");
415 }
416 }
417
418 #[rstest]
419 fn test_multiple_streams_merge_order() {
420 let mut it = BacktestDataIterator::new();
421 it.add_data("s1", vec![quote("A.B", 100), quote("A.B", 300)], true);
422 it.add_data("s2", vec![quote("C.D", 200), quote("C.D", 400)], true);
423
424 assert_eq!(collect_ts(&mut it), vec![100, 200, 300, 400]);
425 }
426
427 #[rstest]
428 fn test_append_data_priority_default_fifo() {
429 let mut it = BacktestDataIterator::new();
430 it.add_data("a", vec![quote("A.B", 100)], true);
431 it.add_data("b", vec![quote("C.D", 100)], true);
432
433 let ts = collect_ts(&mut it);
435 assert_eq!(ts, vec![100, 100]);
436 }
437
438 #[rstest]
439 fn test_prepend_priority_wins_ties() {
440 let mut it = BacktestDataIterator::new();
441 it.add_data("a", vec![quote("A.B", 100)], true);
443 it.add_data("b", vec![quote("C.D", 100)], false);
444
445 let first = it.next().unwrap();
447 let second = it.next().unwrap();
448 assert_eq!(first.instrument_id(), InstrumentId::from("C.D"));
450 assert_eq!(second.instrument_id(), InstrumentId::from("A.B"));
451 }
452
453 #[rstest]
454 fn test_is_done_empty_iterator() {
455 let it = BacktestDataIterator::new();
456 assert!(it.is_done());
457 }
458
459 #[rstest]
460 fn test_is_done_after_consumption() {
461 let mut it = BacktestDataIterator::new();
462 it.add_data("s", vec![quote("A.B", 1)], true);
463
464 assert!(!it.is_done());
465 it.next();
466 assert!(it.is_done());
467 }
468
469 #[rstest]
470 fn test_is_done_multi_stream() {
471 let mut it = BacktestDataIterator::new();
472 it.add_data("s1", vec![quote("A.B", 1)], true);
473 it.add_data("s2", vec![quote("C.D", 2)], true);
474
475 assert!(!it.is_done());
476 it.next();
477 assert!(!it.is_done());
478 it.next();
479 assert!(it.is_done());
480 }
481
482 #[rstest]
483 fn test_partial_consumption_then_complete() {
484 let mut it = BacktestDataIterator::new();
485 it.add_data(
486 "s",
487 vec![
488 quote("A.B", 0),
489 quote("A.B", 1),
490 quote("A.B", 2),
491 quote("A.B", 3),
492 ],
493 true,
494 );
495
496 assert_eq!(it.next().unwrap().ts_init().as_u64(), 0);
497 assert_eq!(it.next().unwrap().ts_init().as_u64(), 1);
498
499 let remaining = collect_ts(&mut it);
500 assert_eq!(remaining, vec![2, 3]);
501 assert!(it.is_done());
502 }
503
504 #[rstest]
505 fn test_remove_stream_reduces_output() {
506 let mut it = BacktestDataIterator::new();
507 it.add_data("a", vec![quote("A.B", 1)], true);
508 it.add_data("b", vec![quote("C.D", 2)], true);
509
510 it.remove_data("a", false);
511
512 assert_eq!(collect_ts(&mut it), vec![2]);
513 }
514
515 #[rstest]
516 fn test_remove_all_streams_yields_empty() {
517 let mut it = BacktestDataIterator::new();
518 it.add_data("x", vec![quote("A.B", 1)], true);
519 it.add_data("y", vec![quote("C.D", 2)], true);
520
521 it.remove_data("x", false);
522 it.remove_data("y", false);
523
524 assert!(it.next().is_none());
525 assert!(it.is_done());
526 }
527
528 #[rstest]
529 fn test_remove_nonexistent_stream_is_noop() {
530 let mut it = BacktestDataIterator::new();
531 it.add_data("s", vec![quote("A.B", 1)], true);
532
533 it.remove_data("nonexistent", false);
534
535 assert_eq!(collect_ts(&mut it), vec![1]);
536 }
537
538 #[rstest]
539 fn test_remove_after_full_consumption() {
540 let mut it = BacktestDataIterator::new();
541 it.add_data("s", vec![quote("A.B", 1), quote("A.B", 2)], true);
542
543 collect_ts(&mut it);
544
545 it.remove_data("s", true);
546 assert!(it.is_done());
547 }
548
549 #[rstest]
550 fn test_set_index_rewinds_stream() {
551 let mut it = BacktestDataIterator::new();
552 it.add_data(
553 "s",
554 vec![quote("A.B", 10), quote("A.B", 20), quote("A.B", 30)],
555 true,
556 );
557
558 assert_eq!(it.next().unwrap().ts_init().as_u64(), 10);
559
560 it.set_index("s", 0);
561
562 assert_eq!(collect_ts(&mut it), vec![10, 20, 30]);
563 }
564
565 #[rstest]
566 fn test_set_index_skips_forward() {
567 let mut it = BacktestDataIterator::new();
568 it.add_data(
569 "s",
570 vec![quote("A.B", 10), quote("A.B", 20), quote("A.B", 30)],
571 true,
572 );
573
574 it.set_index("s", 2);
575
576 assert_eq!(collect_ts(&mut it), vec![30]);
577 }
578
579 #[rstest]
580 fn test_set_index_nonexistent_stream_is_noop() {
581 let mut it = BacktestDataIterator::new();
582 it.add_data("s", vec![quote("A.B", 1)], true);
583
584 it.set_index("nonexistent", 0);
585
586 assert_eq!(collect_ts(&mut it), vec![1]);
587 }
588
589 #[rstest]
590 fn test_reset_all_cursors_single_stream() {
591 let mut it = BacktestDataIterator::new();
592 it.add_data("s", vec![quote("A.B", 1), quote("A.B", 2)], true);
593
594 collect_ts(&mut it);
595 assert!(it.is_done());
596
597 it.reset_all_cursors();
598 assert!(!it.is_done());
599 assert_eq!(collect_ts(&mut it), vec![1, 2]);
600 }
601
602 #[rstest]
603 fn test_reset_all_cursors_multi_stream() {
604 let mut it = BacktestDataIterator::new();
605 it.add_data("s1", vec![quote("A.B", 1), quote("A.B", 3)], true);
606 it.add_data("s2", vec![quote("C.D", 2), quote("C.D", 4)], true);
607
608 collect_ts(&mut it);
609 assert!(it.is_done());
610
611 it.reset_all_cursors();
612 assert_eq!(collect_ts(&mut it), vec![1, 2, 3, 4]);
613 }
614
615 #[rstest]
616 fn test_readding_data_replaces_stream() {
617 let mut it = BacktestDataIterator::new();
618 it.add_data("X", vec![quote("A.B", 1), quote("A.B", 2)], true);
619 it.add_data("X", vec![quote("A.B", 10)], true);
620
621 assert_eq!(collect_ts(&mut it), vec![10]);
622 }
623
624 #[rstest]
625 fn test_add_empty_data_is_noop() {
626 let mut it = BacktestDataIterator::new();
627 it.add_data("empty", vec![], true);
628
629 assert!(it.is_done());
630 assert!(it.next().is_none());
631 }
632
633 #[rstest]
634 fn test_empty_iterator_returns_none() {
635 let mut it = BacktestDataIterator::new();
636 assert!(it.next().is_none());
637 assert!(it.is_done());
638 }
639
640 #[rstest]
641 fn test_multiple_add_data_calls_with_different_names() {
642 let mut it = BacktestDataIterator::new();
643 it.add_data("batch_0", vec![quote("A.B", 1), quote("A.B", 3)], true);
644 it.add_data("batch_1", vec![quote("A.B", 2), quote("A.B", 4)], true);
645
646 assert_eq!(collect_ts(&mut it), vec![1, 2, 3, 4]);
647 }
648
649 #[rstest]
650 fn test_prepend_stream_always_wins_ties_across_batches() {
651 let mut it = BacktestDataIterator::new();
654 it.add_data("append_a", vec![quote("A.B", 100)], true);
655 it.add_data("append_b", vec![quote("C.D", 100)], true);
656 it.add_data("prepend", vec![quote("E.F", 100)], false);
657
658 let first = it.next().unwrap();
659 assert_eq!(
660 first.instrument_id(),
661 InstrumentId::from("E.F"),
662 "Prepend stream should always come first in ties"
663 );
664 }
665
666 #[rstest]
667 fn test_equal_timestamps_across_many_streams_preserves_priority_order() {
668 let mut it = BacktestDataIterator::new();
670 it.add_data("s1", vec![quote("A.B", 50)], true);
671 it.add_data("s2", vec![quote("C.D", 50)], true);
672 it.add_data("s3", vec![quote("E.F", 50)], true);
673 it.add_data("s4", vec![quote("G.H", 50)], true);
674
675 let mut ids = Vec::new();
676 while let Some(d) = it.next() {
677 ids.push(d.instrument_id());
678 }
679
680 assert_eq!(ids.len(), 4);
681
682 assert!(ids.contains(&InstrumentId::from("A.B")));
684 assert!(ids.contains(&InstrumentId::from("C.D")));
685 assert!(ids.contains(&InstrumentId::from("E.F")));
686 assert!(ids.contains(&InstrumentId::from("G.H")));
687 }
688
689 #[cfg(feature = "defi")]
690 #[rstest]
691 fn test_defi_data_orders_equal_timestamps_by_block_position() {
692 let mut it = BacktestDataIterator::new();
693 it.add_data(
694 "defi",
695 vec![
696 defi_snapshot(100, 12, 4, 1),
697 defi_snapshot(100, 11, 9, 9),
698 defi_snapshot(100, 12, 2, 7),
699 ],
700 true,
701 );
702
703 let mut positions = Vec::new();
704 while let Some(Data::Defi(data)) = it.next_item() {
705 positions.push(data.block_position());
706 }
707
708 assert_eq!(positions, vec![(11, 9, 9), (12, 2, 7), (12, 4, 1)]);
709 }
710}