mithril_cardano_node_chain/test/double/
chain_data_store.rs

1use std::collections::BTreeSet;
2use std::ops::Range;
3use std::sync::Arc;
4
5use slog::Logger;
6use tokio::sync::Mutex;
7
8use mithril_common::StdResult;
9use mithril_common::crypto_helper::{MKTreeNode, MKTreeStorer};
10use mithril_common::entities::{
11    BlockHash, BlockNumber, BlockRange, CardanoBlock, CardanoBlockTransactionMkTreeNode,
12    CardanoBlockWithTransactions, CardanoTransaction, ChainPoint, SlotNumber, TransactionHash,
13};
14use mithril_common::signable_builder::BlockRangeRootRetriever;
15
16use crate::chain_importer::{BlockRangeImporter, ChainDataStore};
17
18/// In memory Block range root representation, for testing purposes.
19#[derive(Debug, PartialEq, Clone)]
20pub struct InMemoryBlockRangeRoot {
21    /// Range of block numbers covered
22    pub range: BlockRange,
23    /// Merkle root of the block range
24    pub merkle_root: MKTreeNode,
25}
26
27impl From<(BlockRange, MKTreeNode)> for InMemoryBlockRangeRoot {
28    fn from(value: (BlockRange, MKTreeNode)) -> Self {
29        Self {
30            range: value.0,
31            merkle_root: value.1,
32        }
33    }
34}
35
36impl From<InMemoryBlockRangeRoot> for (BlockRange, MKTreeNode) {
37    fn from(value: InMemoryBlockRangeRoot) -> Self {
38        (value.range, value.merkle_root)
39    }
40}
41
42/// In memory implementation of [ChainDataStore], for testing purposes.
43pub struct InMemoryChainDataStore {
44    blocks_with_txs: Mutex<Vec<CardanoBlockWithTransactions>>,
45    block_range_roots: Mutex<Vec<InMemoryBlockRangeRoot>>,
46    legacy_block_range_roots: Mutex<Vec<InMemoryBlockRangeRoot>>,
47}
48
49impl Default for InMemoryChainDataStore {
50    fn default() -> Self {
51        Self::builder().build()
52    }
53}
54
55/// Builder for [InMemoryChainDataStore].
56pub struct InMemoryChainDataStoreBuilder {
57    blocks_with_txs: Vec<CardanoBlockWithTransactions>,
58    block_range_roots: Vec<InMemoryBlockRangeRoot>,
59    legacy_block_range_roots: Vec<InMemoryBlockRangeRoot>,
60}
61
62impl InMemoryChainDataStoreBuilder {
63    /// Set the initial blocks and transactions for the store, replacing any existing data.
64    pub fn with_blocks_and_transactions<T: Into<CardanoBlockWithTransactions> + Clone>(
65        mut self,
66        transactions: &[T],
67    ) -> Self {
68        self.blocks_with_txs = transactions.iter().map(|b| b.clone().into()).collect();
69        self
70    }
71
72    /// Set the initial block range roots for the store, replacing any existing data.
73    pub fn with_block_range_roots<T: Into<InMemoryBlockRangeRoot> + Clone>(
74        mut self,
75        block_range_roots: &[T],
76    ) -> Self {
77        self.block_range_roots = block_range_roots.iter().map(|brr| brr.clone().into()).collect();
78        self
79    }
80
81    /// Set the initial block range roots for the store, replacing any existing data.
82    pub fn with_legacy_block_range_roots<T: Into<InMemoryBlockRangeRoot> + Clone>(
83        mut self,
84        block_range_roots: &[T],
85    ) -> Self {
86        self.legacy_block_range_roots =
87            block_range_roots.iter().map(|brr| brr.clone().into()).collect();
88        self
89    }
90
91    /// Computes the block ranges roots, new and legacy, up to the block range that includes the given
92    /// block number and based on the blocks given in [Self::with_blocks_and_transactions]
93    ///
94    /// Replace any existing block ranges roots and legacy block ranges roots.
95    pub async fn compute_block_ranges(
96        mut self,
97        up_to_block_range_that_include: BlockNumber,
98    ) -> Self {
99        // Leverage tested `BlockRangeImporter` to compute the block ranges roots.
100        let worker_store = Arc::new(InMemoryChainDataStore {
101            blocks_with_txs: Mutex::new(self.blocks_with_txs.clone()),
102            block_range_roots: Default::default(),
103            legacy_block_range_roots: Default::default(),
104        });
105
106        let discard_logger = Logger::root(slog::Discard, slog::o!());
107        let importer = BlockRangeImporter::new(worker_store.clone(), discard_logger);
108        importer.run(up_to_block_range_that_include).await.unwrap();
109        importer.run_legacy(up_to_block_range_that_include).await.unwrap();
110
111        self.block_range_roots = worker_store.get_all_block_range_root().await;
112        self.legacy_block_range_roots = worker_store.get_all_legacy_block_range_root().await;
113
114        self
115    }
116
117    /// Creates a new [InMemoryChainDataStore] with the current builder's data.
118    pub fn build(self) -> InMemoryChainDataStore {
119        InMemoryChainDataStore {
120            blocks_with_txs: Mutex::new(self.blocks_with_txs),
121            block_range_roots: Mutex::new(self.block_range_roots),
122            legacy_block_range_roots: Mutex::new(self.legacy_block_range_roots),
123        }
124    }
125}
126
127impl InMemoryChainDataStore {
128    /// Creates a new in-memory store builder with no initial data.
129    pub fn builder() -> InMemoryChainDataStoreBuilder {
130        InMemoryChainDataStoreBuilder {
131            blocks_with_txs: vec![],
132            block_range_roots: vec![],
133            legacy_block_range_roots: vec![],
134        }
135    }
136
137    /// Returns all transactions in the store.
138    pub async fn get_all_transactions(&self) -> Vec<CardanoTransaction> {
139        let blocks = self.blocks_with_txs.lock().await.clone();
140        blocks.into_iter().flat_map(|b| b.into_transactions()).collect()
141    }
142
143    /// Returns all [CardanoBlockWithTransactions] in the store.
144    pub async fn get_all_block_with_txs(&self) -> Vec<CardanoBlockWithTransactions> {
145        self.blocks_with_txs.lock().await.clone()
146    }
147
148    /// Returns all block range roots in the store.
149    pub async fn get_all_block_range_root(&self) -> Vec<InMemoryBlockRangeRoot> {
150        self.block_range_roots.lock().await.clone()
151    }
152
153    /// Returns all block ranges in the store.
154    pub async fn get_all_block_ranges(&self) -> Vec<BlockRange> {
155        self.block_range_roots
156            .lock()
157            .await
158            .iter()
159            .map(|r| r.range.clone())
160            .collect()
161    }
162
163    /// Returns all legacy block range roots in the store.
164    pub async fn get_all_legacy_block_range_root(&self) -> Vec<InMemoryBlockRangeRoot> {
165        self.legacy_block_range_roots.lock().await.clone()
166    }
167
168    /// Returns all legacy block ranges in the store.
169    pub async fn get_all_legacy_block_ranges(&self) -> Vec<BlockRange> {
170        self.legacy_block_range_roots
171            .lock()
172            .await
173            .iter()
174            .map(|r| r.range.clone())
175            .collect()
176    }
177
178    /// Returns all [CardanoBlock] with the given block hashes.
179    pub async fn get_blocks_by_hashes(&self, block_hashes: &[BlockHash]) -> Vec<CardanoBlock> {
180        let blocks_with_txs = self.blocks_with_txs.lock().await;
181        blocks_with_txs
182            .iter()
183            .filter(|block| block_hashes.contains(&block.block_hash))
184            .cloned()
185            .map(Into::into)
186            .collect()
187    }
188
189    /// Returns all [CardanoTransaction] with the given transaction hashes.
190    pub async fn get_transactions_by_hashes(
191        &self,
192        transaction_hashes: &[TransactionHash],
193    ) -> Vec<CardanoTransaction> {
194        let blocks_with_txs = self.blocks_with_txs.lock().await.clone();
195        blocks_with_txs
196            .into_iter()
197            .flat_map(|block| block.into_transactions())
198            .filter(|tx| transaction_hashes.contains(&tx.transaction_hash))
199            .collect()
200    }
201
202    /// Returns all [CardanoBlockTransactionMkTreeNode] contained in the given block ranges.
203    pub async fn get_blocks_with_transactions_in_block_ranges(
204        &self,
205        block_ranges: &[BlockRange],
206    ) -> BTreeSet<CardanoBlockTransactionMkTreeNode> {
207        let ranges: Vec<Range<BlockNumber>> =
208            block_ranges.iter().map(|range| range.clone().into()).collect();
209        self.get_blocks_with_transactions_in_ranges(&ranges).await
210    }
211
212    /// Returns all [CardanoBlockTransactionMkTreeNode] contained in the given ranges.
213    pub async fn get_blocks_with_transactions_in_ranges(
214        &self,
215        block_ranges: &[Range<BlockNumber>],
216    ) -> BTreeSet<CardanoBlockTransactionMkTreeNode> {
217        let blocks = self.blocks_with_txs.lock().await;
218        blocks
219            .iter()
220            .filter(|block| block_ranges.iter().any(|range| range.contains(&block.block_number)))
221            .cloned()
222            .flat_map(|block| block.into_mk_tree_node())
223            .collect()
224    }
225}
226
227#[async_trait::async_trait]
228impl ChainDataStore for InMemoryChainDataStore {
229    async fn get_highest_beacon(&self) -> StdResult<Option<ChainPoint>> {
230        let txs = self.blocks_with_txs.lock().await;
231        Ok(txs
232            .iter()
233            .max_by_key(|tx| tx.block_number)
234            .map(|tx| ChainPoint::new(tx.slot_number, tx.block_number, tx.block_hash.clone())))
235    }
236
237    async fn get_highest_block_range(&self) -> StdResult<Option<BlockRange>> {
238        let roots = self.block_range_roots.lock().await;
239        Ok(roots.iter().map(|record| record.range.clone()).max_by_key(|r| r.end))
240    }
241
242    async fn get_highest_legacy_block_range(&self) -> StdResult<Option<BlockRange>> {
243        let roots = self.legacy_block_range_roots.lock().await;
244        Ok(roots.iter().map(|record| record.range.clone()).max_by_key(|r| r.end))
245    }
246
247    async fn store_blocks_and_transactions(
248        &self,
249        blocks_and_transactions: Vec<CardanoBlockWithTransactions>,
250    ) -> StdResult<()> {
251        self.blocks_with_txs.lock().await.extend(blocks_and_transactions);
252        Ok(())
253    }
254
255    async fn get_blocks_and_transactions_in_range(
256        &self,
257        range: Range<BlockNumber>,
258    ) -> StdResult<BTreeSet<CardanoBlockTransactionMkTreeNode>> {
259        Ok(self.get_blocks_with_transactions_in_ranges(&[range]).await)
260    }
261
262    async fn get_transactions_in_range(
263        &self,
264        range: Range<BlockNumber>,
265    ) -> StdResult<Vec<CardanoTransaction>> {
266        let txs = self.blocks_with_txs.lock().await;
267        Ok(txs
268            .iter()
269            .filter(|tx| range.contains(&tx.block_number))
270            .cloned()
271            .flat_map(|tx| tx.into_transactions())
272            .collect())
273    }
274
275    async fn store_block_range_roots(
276        &self,
277        block_ranges: Vec<(BlockRange, MKTreeNode)>,
278    ) -> StdResult<()> {
279        self.block_range_roots
280            .lock()
281            .await
282            .extend(block_ranges.into_iter().map(Into::into));
283        Ok(())
284    }
285
286    async fn store_legacy_block_range_roots(
287        &self,
288        block_ranges: Vec<(BlockRange, MKTreeNode)>,
289    ) -> StdResult<()> {
290        self.legacy_block_range_roots
291            .lock()
292            .await
293            .extend(block_ranges.into_iter().map(Into::into));
294        Ok(())
295    }
296
297    async fn remove_rolled_chain_data_and_block_range(
298        &self,
299        slot_number: SlotNumber,
300    ) -> StdResult<()> {
301        self.blocks_with_txs
302            .lock()
303            .await
304            .retain(|b| b.slot_number <= slot_number);
305
306        if let Some(highest_remaining_block_number) =
307            self.blocks_with_txs.lock().await.last().map(|tx| tx.block_number)
308        {
309            self.block_range_roots
310                .lock()
311                .await
312                .retain(|record| record.range.start < highest_remaining_block_number);
313            self.legacy_block_range_roots
314                .lock()
315                .await
316                .retain(|record| record.range.start < highest_remaining_block_number);
317        } else {
318            self.block_range_roots.lock().await.clear();
319            self.legacy_block_range_roots.lock().await.clear();
320        }
321
322        Ok(())
323    }
324
325    async fn optimize(&self) -> StdResult<()> {
326        Ok(())
327    }
328}
329
330#[async_trait::async_trait]
331impl<S: MKTreeStorer> BlockRangeRootRetriever<S> for InMemoryChainDataStore {
332    async fn retrieve_block_range_roots<'a>(
333        &'a self,
334        up_to_beacon: BlockNumber,
335    ) -> StdResult<Box<dyn Iterator<Item = (BlockRange, MKTreeNode)> + 'a>> {
336        let block_ranges = self.block_range_roots.lock().await;
337        let result: Vec<_> = block_ranges
338            .iter()
339            .filter(|r| r.range.start < up_to_beacon)
340            .cloned()
341            .map(|r| (r.range, r.merkle_root))
342            .collect();
343        Ok(Box::new(result.into_iter()))
344    }
345}
346
347#[cfg(test)]
348mod tests {
349    use mithril_common::crypto_helper::MKTreeStoreInMemory;
350
351    use super::*;
352
353    #[tokio::test]
354    async fn builder_can_compute_block_range_roots() {
355        let expected_computed_ranges = vec![
356            BlockRange::from_block_number(BlockNumber(10)),
357            BlockRange::from_block_number(BlockNumber(25)),
358        ];
359        let builder = InMemoryChainDataStore::builder().with_blocks_and_transactions(&[
360            CardanoBlockWithTransactions::new(
361                "block_hash-10",
362                BlockNumber(10),
363                SlotNumber(50),
364                vec!["tx_hash-1"],
365            ),
366            CardanoBlockWithTransactions::new(
367                "block_hash-25",
368                BlockNumber(25),
369                SlotNumber(51),
370                vec!["tx_hash-2", "tx_hash-3"],
371            ),
372            CardanoBlockWithTransactions::new(
373                "block_hash-30",
374                BlockNumber(30),
375                SlotNumber(52),
376                vec!["tx_hash-4", "tx_hash-5"],
377            ),
378        ]);
379
380        assert_eq!(
381            Vec::<InMemoryBlockRangeRoot>::new(),
382            builder.block_range_roots
383        );
384        assert_eq!(
385            Vec::<InMemoryBlockRangeRoot>::new(),
386            builder.legacy_block_range_roots
387        );
388
389        let builder = builder.compute_block_ranges(BlockNumber(29)).await;
390
391        assert_eq!(
392            expected_computed_ranges,
393            builder
394                .block_range_roots
395                .iter()
396                .map(|b| b.range.clone())
397                .collect::<Vec<_>>()
398        );
399        assert_eq!(
400            expected_computed_ranges,
401            builder
402                .legacy_block_range_roots
403                .iter()
404                .map(|b| b.range.clone())
405                .collect::<Vec<_>>()
406        );
407    }
408
409    #[tokio::test]
410    async fn default_store_is_empty() {
411        let store = InMemoryChainDataStore::default();
412
413        assert!(store.get_all_transactions().await.is_empty());
414        assert!(store.get_all_block_with_txs().await.is_empty());
415        assert!(store.get_all_block_ranges().await.is_empty());
416        assert!(store.get_all_block_range_root().await.is_empty());
417        assert!(store.get_all_legacy_block_range_root().await.is_empty());
418        assert!(store.get_all_legacy_block_ranges().await.is_empty());
419    }
420
421    #[tokio::test]
422    async fn store_and_get_blocks_and_transactions() {
423        let store = InMemoryChainDataStore::default();
424
425        let block_with_tx = vec![
426            CardanoBlockWithTransactions::new(
427                "block_hash-123",
428                BlockNumber(10),
429                SlotNumber(50),
430                vec!["tx_hash-123"],
431            ),
432            CardanoBlockWithTransactions::new(
433                "block_hash-456",
434                BlockNumber(11),
435                SlotNumber(51),
436                vec!["tx_hash-456", "tx_hash-789"],
437            ),
438        ];
439
440        store
441            .store_blocks_and_transactions(block_with_tx.clone())
442            .await
443            .unwrap();
444
445        let stored_blocks_with_txs = store.get_all_block_with_txs().await;
446        assert_eq!(block_with_tx, stored_blocks_with_txs);
447
448        let stored_transactions = store.get_all_transactions().await;
449        assert_eq!(
450            vec![
451                CardanoTransaction::new(
452                    "tx_hash-123",
453                    BlockNumber(10),
454                    SlotNumber(50),
455                    "block_hash-123"
456                ),
457                CardanoTransaction::new(
458                    "tx_hash-456",
459                    BlockNumber(11),
460                    SlotNumber(51),
461                    "block_hash-456"
462                ),
463                CardanoTransaction::new(
464                    "tx_hash-789",
465                    BlockNumber(11),
466                    SlotNumber(51),
467                    "block_hash-456"
468                )
469            ],
470            stored_transactions
471        );
472    }
473
474    #[tokio::test]
475    async fn store_transactions_appends_to_existing() {
476        let existing_batch = vec![CardanoBlockWithTransactions::new(
477            "block_hash-123",
478            BlockNumber(10),
479            SlotNumber(50),
480            vec!["tx_hash-123"],
481        )];
482
483        let store = InMemoryChainDataStore::builder()
484            .with_blocks_and_transactions(&existing_batch)
485            .build();
486
487        let second_batch = vec![CardanoBlockWithTransactions::new(
488            "block_hash-456",
489            BlockNumber(11),
490            SlotNumber(51),
491            vec!["tx_hash-456"],
492        )];
493
494        store
495            .store_blocks_and_transactions(second_batch.clone())
496            .await
497            .unwrap();
498
499        let stored_blocks_with_tx = store.get_all_block_with_txs().await;
500        assert_eq!(2, stored_blocks_with_tx.len());
501        assert_eq!(existing_batch[0], stored_blocks_with_tx[0]);
502        assert_eq!(second_batch[0], stored_blocks_with_tx[1]);
503    }
504
505    #[tokio::test]
506    async fn get_highest_beacon_returns_none_when_empty() {
507        let store = InMemoryChainDataStore::default();
508
509        let highest_beacon = store.get_highest_beacon().await.unwrap();
510
511        assert_eq!(None, highest_beacon);
512    }
513
514    #[tokio::test]
515    async fn get_highest_beacon_returns_transaction_with_highest_block_number() {
516        let store = InMemoryChainDataStore::builder()
517            .with_blocks_and_transactions(&[
518                CardanoBlockWithTransactions::new(
519                    "block_hash-10",
520                    BlockNumber(10),
521                    SlotNumber(50),
522                    vec!["tx_hash-123"],
523                ),
524                CardanoBlockWithTransactions::new(
525                    "block_hash-25",
526                    BlockNumber(25),
527                    SlotNumber(51),
528                    vec!["tx_hash-456"],
529                ),
530                CardanoBlockWithTransactions::new(
531                    "block_hash-15",
532                    BlockNumber(15),
533                    SlotNumber(52),
534                    vec!["tx_hash-789"],
535                ),
536            ])
537            .build();
538
539        let highest_beacon = store.get_highest_beacon().await.unwrap();
540
541        assert_eq!(
542            Some(ChainPoint::new(
543                SlotNumber(51),
544                BlockNumber(25),
545                "block_hash-25"
546            )),
547            highest_beacon
548        );
549    }
550
551    #[tokio::test]
552    async fn get_highest_beacon_with_multiple_blocks_with_same_block_number() {
553        let store = InMemoryChainDataStore::builder()
554            .with_blocks_and_transactions(&[
555                CardanoBlockWithTransactions::new(
556                    "block_hash-10",
557                    BlockNumber(10),
558                    SlotNumber(50),
559                    vec!["tx_hash-123"],
560                ),
561                CardanoBlockWithTransactions::new(
562                    "block_hash-25",
563                    BlockNumber(25),
564                    SlotNumber(51),
565                    vec!["tx_hash-456"],
566                ),
567                CardanoBlockWithTransactions::new(
568                    "block_hash-25",
569                    BlockNumber(25),
570                    SlotNumber(51),
571                    vec!["tx_hash-789"],
572                ),
573            ])
574            .build();
575
576        let highest_beacon = store.get_highest_beacon().await.unwrap();
577
578        assert_eq!(
579            Some(ChainPoint::new(
580                SlotNumber(51),
581                BlockNumber(25),
582                "block_hash-25"
583            )),
584            highest_beacon
585        );
586    }
587
588    #[tokio::test]
589    async fn store_and_get_block_range_roots() {
590        let store = InMemoryChainDataStore::default();
591
592        let block_ranges = vec![
593            (
594                BlockRange::from_block_number(BlockNumber(0)),
595                MKTreeNode::from_hex("AAAA").unwrap(),
596            ),
597            (
598                BlockRange::from_block_number(BlockRange::LENGTH),
599                MKTreeNode::from_hex("BBBB").unwrap(),
600            ),
601        ];
602
603        store.store_block_range_roots(block_ranges.clone()).await.unwrap();
604
605        let stored_roots = store.get_all_block_range_root().await;
606        assert_eq!(
607            vec![
608                InMemoryBlockRangeRoot {
609                    range: BlockRange::from_block_number(BlockNumber(0)),
610                    merkle_root: MKTreeNode::from_hex("AAAA").unwrap(),
611                },
612                InMemoryBlockRangeRoot {
613                    range: BlockRange::from_block_number(BlockRange::LENGTH),
614                    merkle_root: MKTreeNode::from_hex("BBBB").unwrap(),
615                },
616            ],
617            stored_roots
618        );
619    }
620
621    #[tokio::test]
622    async fn store_block_range_roots_appends_to_existing() {
623        let store = InMemoryChainDataStore::builder()
624            .with_block_range_roots(&[(
625                BlockRange::from_block_number(BlockNumber(0)),
626                MKTreeNode::from_hex("AAAA").unwrap(),
627            )])
628            .build();
629
630        store
631            .store_block_range_roots(vec![(
632                BlockRange::from_block_number(BlockRange::LENGTH),
633                MKTreeNode::from_hex("BBBB").unwrap(),
634            )])
635            .await
636            .unwrap();
637
638        let stored_roots = store.get_all_block_range_root().await;
639        assert_eq!(
640            vec![
641                InMemoryBlockRangeRoot {
642                    range: BlockRange::from_block_number(BlockNumber(0)),
643                    merkle_root: MKTreeNode::from_hex("AAAA").unwrap()
644                },
645                InMemoryBlockRangeRoot {
646                    range: BlockRange::from_block_number(BlockRange::LENGTH),
647                    merkle_root: MKTreeNode::from_hex("BBBB").unwrap()
648                },
649            ],
650            stored_roots
651        );
652
653        let ranges = store.get_all_block_ranges().await;
654        assert_eq!(
655            vec![
656                BlockRange::from_block_number(BlockNumber(0)),
657                BlockRange::from_block_number(BlockRange::LENGTH),
658            ],
659            ranges
660        );
661    }
662
663    #[tokio::test]
664    async fn get_highest_block_range_returns_none_when_empty() {
665        let store = InMemoryChainDataStore::default();
666
667        let highest_range = store.get_highest_block_range().await.unwrap();
668
669        assert_eq!(None, highest_range);
670    }
671
672    #[tokio::test]
673    async fn store_and_get_legacy_block_range_roots() {
674        let store = InMemoryChainDataStore::default();
675
676        let block_ranges = vec![
677            (
678                BlockRange::from_block_number(BlockNumber(0)),
679                MKTreeNode::from_hex("AAAA").unwrap(),
680            ),
681            (
682                BlockRange::from_block_number(BlockRange::LENGTH),
683                MKTreeNode::from_hex("BBBB").unwrap(),
684            ),
685        ];
686
687        store
688            .store_legacy_block_range_roots(block_ranges.clone())
689            .await
690            .unwrap();
691
692        let stored_roots = store.get_all_legacy_block_range_root().await;
693        assert_eq!(
694            vec![
695                InMemoryBlockRangeRoot {
696                    range: BlockRange::from_block_number(BlockNumber(0)),
697                    merkle_root: MKTreeNode::from_hex("AAAA").unwrap(),
698                },
699                InMemoryBlockRangeRoot {
700                    range: BlockRange::from_block_number(BlockRange::LENGTH),
701                    merkle_root: MKTreeNode::from_hex("BBBB").unwrap(),
702                },
703            ],
704            stored_roots
705        );
706    }
707
708    #[tokio::test]
709    async fn store_legacy_block_range_roots_appends_to_existing() {
710        let store = InMemoryChainDataStore::builder()
711            .with_legacy_block_range_roots(&[(
712                BlockRange::from_block_number(BlockNumber(0)),
713                MKTreeNode::from_hex("AAAA").unwrap(),
714            )])
715            .build();
716
717        store
718            .store_legacy_block_range_roots(vec![(
719                BlockRange::from_block_number(BlockRange::LENGTH),
720                MKTreeNode::from_hex("BBBB").unwrap(),
721            )])
722            .await
723            .unwrap();
724
725        let stored_roots = store.get_all_legacy_block_range_root().await;
726        assert_eq!(
727            vec![
728                InMemoryBlockRangeRoot {
729                    range: BlockRange::from_block_number(BlockNumber(0)),
730                    merkle_root: MKTreeNode::from_hex("AAAA").unwrap()
731                },
732                InMemoryBlockRangeRoot {
733                    range: BlockRange::from_block_number(BlockRange::LENGTH),
734                    merkle_root: MKTreeNode::from_hex("BBBB").unwrap()
735                },
736            ],
737            stored_roots
738        );
739
740        let ranges = store.get_all_legacy_block_ranges().await;
741        assert_eq!(
742            vec![
743                BlockRange::from_block_number(BlockNumber(0)),
744                BlockRange::from_block_number(BlockRange::LENGTH),
745            ],
746            ranges
747        );
748    }
749
750    #[tokio::test]
751    async fn get_highest_legacy_block_range_returns_none_when_empty() {
752        let store = InMemoryChainDataStore::default();
753
754        let highest_range = store.get_highest_legacy_block_range().await.unwrap();
755
756        assert_eq!(None, highest_range);
757    }
758
759    #[tokio::test]
760    async fn get_highest_legacy_block_range_returns_range_with_highest_end() {
761        let store = InMemoryChainDataStore::builder()
762            .with_legacy_block_range_roots(&[
763                (
764                    BlockRange::from_block_number(BlockNumber(0)),
765                    MKTreeNode::from_hex("AAAA").unwrap(),
766                ),
767                (
768                    BlockRange::from_block_number(BlockRange::LENGTH * 2),
769                    MKTreeNode::from_hex("CCCC").unwrap(),
770                ),
771                (
772                    BlockRange::from_block_number(BlockRange::LENGTH),
773                    MKTreeNode::from_hex("BBBB").unwrap(),
774                ),
775            ])
776            .build();
777
778        let highest_range = store.get_highest_legacy_block_range().await.unwrap();
779
780        assert_eq!(
781            Some(BlockRange::from_block_number(BlockRange::LENGTH * 2)),
782            highest_range
783        );
784    }
785
786    #[tokio::test]
787    async fn get_transactions_in_range_returns_empty_when_store_empty() {
788        let store = InMemoryChainDataStore::default();
789
790        let transactions = store
791            .get_transactions_in_range(BlockNumber(0)..BlockNumber(100))
792            .await
793            .unwrap();
794
795        assert!(transactions.is_empty());
796    }
797
798    #[tokio::test]
799    async fn get_transactions_in_range_filters_correctly() {
800        let blocks_with_tx = vec![
801            CardanoBlockWithTransactions::new(
802                "block-hash-1",
803                BlockNumber(10),
804                SlotNumber(50),
805                vec!["tx-hash-1"],
806            ),
807            CardanoBlockWithTransactions::new(
808                "block-hash-2",
809                BlockNumber(11),
810                SlotNumber(51),
811                vec!["tx-hash-2"],
812            ),
813            CardanoBlockWithTransactions::new(
814                "block-hash-3",
815                BlockNumber(12),
816                SlotNumber(52),
817                vec!["tx-hash-3"],
818            ),
819        ];
820        let store = InMemoryChainDataStore::builder()
821            .with_blocks_and_transactions(&blocks_with_tx)
822            .build();
823        let transactions: Vec<CardanoTransaction> = blocks_with_tx
824            .into_iter()
825            .flat_map(|tx| tx.into_transactions())
826            .collect();
827
828        // Range excludes all transactions
829        {
830            let result = store
831                .get_transactions_in_range(BlockNumber(0)..BlockNumber(10))
832                .await
833                .unwrap();
834            assert!(result.is_empty());
835        }
836
837        // Range after all transactions
838        {
839            let result = store
840                .get_transactions_in_range(BlockNumber(13)..BlockNumber(21))
841                .await
842                .unwrap();
843            assert!(result.is_empty());
844        }
845
846        // Range includes the first two transactions (10, 11)
847        {
848            let result = store
849                .get_transactions_in_range(BlockNumber(9)..BlockNumber(12))
850                .await
851                .unwrap();
852            assert_eq!(transactions[0..=1].to_vec(), result);
853        }
854
855        // Range includes all transactions
856        {
857            let result = store
858                .get_transactions_in_range(BlockNumber(10)..BlockNumber(13))
859                .await
860                .unwrap();
861            assert_eq!(transactions, result);
862        }
863
864        // Range includes the last two transactions (11, 12)
865        {
866            let result = store
867                .get_transactions_in_range(BlockNumber(11)..BlockNumber(14))
868                .await
869                .unwrap();
870            assert_eq!(transactions[1..=2].to_vec(), result);
871        }
872    }
873
874    #[tokio::test]
875    async fn get_blocks_with_transactions_in_ranges_returns_empty_when_store_empty() {
876        let store = InMemoryChainDataStore::default();
877
878        let transactions = store
879            .get_blocks_with_transactions_in_ranges(&[BlockNumber(0)..BlockNumber(100)])
880            .await;
881
882        assert!(transactions.is_empty());
883    }
884
885    #[tokio::test]
886    async fn get_blocks_with_transactions_in_ranges_filters_correctly() {
887        fn into_expected_nodes(
888            cbtx: &[CardanoBlockWithTransactions],
889        ) -> BTreeSet<CardanoBlockTransactionMkTreeNode> {
890            cbtx.iter().flat_map(|b| b.clone().into_mk_tree_node()).collect()
891        }
892
893        let blocks_with_tx = vec![
894            CardanoBlockWithTransactions::new(
895                "block-hash-1",
896                BlockNumber(10),
897                SlotNumber(50),
898                vec!["tx-hash-1"],
899            ),
900            CardanoBlockWithTransactions::new(
901                "block-hash-2",
902                BlockNumber(11),
903                SlotNumber(51),
904                vec!["tx-hash-2"],
905            ),
906            CardanoBlockWithTransactions::new(
907                "block-hash-3",
908                BlockNumber(12),
909                SlotNumber(52),
910                vec!["tx-hash-3"],
911            ),
912        ];
913        let store = InMemoryChainDataStore::builder()
914            .with_blocks_and_transactions(&blocks_with_tx)
915            .build();
916
917        // One range that include no transactions
918        {
919            let result = store
920                .get_blocks_with_transactions_in_ranges(&[BlockNumber(0)..BlockNumber(10)])
921                .await;
922            assert!(result.is_empty());
923        }
924
925        // One range after all transactions
926        {
927            let result = store
928                .get_blocks_with_transactions_in_ranges(&[BlockNumber(13)..BlockNumber(21)])
929                .await;
930            assert!(result.is_empty());
931        }
932
933        // One range includes the first two transactions (10, 11)
934        {
935            let result = store
936                .get_blocks_with_transactions_in_ranges(&[BlockNumber(9)..BlockNumber(12)])
937                .await;
938            assert_eq!(into_expected_nodes(&blocks_with_tx[0..=1]), result);
939        }
940
941        // One range includes all transactions
942        {
943            let result = store
944                .get_blocks_with_transactions_in_ranges(&[BlockNumber(10)..BlockNumber(13)])
945                .await;
946            assert_eq!(into_expected_nodes(&blocks_with_tx), result);
947        }
948
949        // One range includes the last two transactions (11, 12)
950        {
951            let result = store
952                .get_blocks_with_transactions_in_ranges(&[BlockNumber(11)..BlockNumber(14)])
953                .await;
954            assert_eq!(into_expected_nodes(&blocks_with_tx[1..=2]), result);
955        }
956
957        // Two ranges, one that include the first transaction (10) and the other that include the last transaction (12)
958        {
959            let result = store
960                .get_blocks_with_transactions_in_ranges(&[
961                    BlockNumber(9)..BlockNumber(11),
962                    BlockNumber(12)..BlockNumber(13),
963                ])
964                .await;
965            assert_eq!(
966                into_expected_nodes(&[blocks_with_tx[0].clone(), blocks_with_tx[2].clone()]),
967                result
968            );
969        }
970    }
971
972    #[tokio::test]
973    async fn get_blocks_by_hashes_filters_correctly() {
974        let blocks_with_tx = vec![
975            CardanoBlockWithTransactions::new(
976                "block-hash-1",
977                BlockNumber(10),
978                SlotNumber(50),
979                vec!["tx-hash-1"],
980            ),
981            CardanoBlockWithTransactions::new(
982                "block-hash-2",
983                BlockNumber(11),
984                SlotNumber(51),
985                vec!["tx-hash-2"],
986            ),
987            CardanoBlockWithTransactions::new(
988                "block-hash-3",
989                BlockNumber(12),
990                SlotNumber(52),
991                vec!["tx-hash-3"],
992            ),
993        ];
994        let store = InMemoryChainDataStore::builder()
995            .with_blocks_and_transactions(&blocks_with_tx)
996            .build();
997
998        // Get one existing hash
999        {
1000            let result = store.get_blocks_by_hashes(&["block-hash-1".to_string()]).await;
1001            assert_eq!(
1002                vec![CardanoBlock::new("block-hash-1", BlockNumber(10), SlotNumber(50))],
1003                result
1004            );
1005        }
1006        // Get one non-existing hash
1007        {
1008            let result = store.get_blocks_by_hashes(&["block-hash-4".to_string()]).await;
1009            assert_eq!(Vec::<CardanoBlock>::new(), result);
1010        }
1011        // Get two existing hashes and one non-existing hash
1012        {
1013            let result = store
1014                .get_blocks_by_hashes(&[
1015                    "block-hash-1".to_string(),
1016                    "block-hash-3".to_string(),
1017                    "block-hash-4".to_string(),
1018                ])
1019                .await;
1020            assert_eq!(
1021                vec![
1022                    CardanoBlock::new("block-hash-1", BlockNumber(10), SlotNumber(50)),
1023                    CardanoBlock::new("block-hash-3", BlockNumber(12), SlotNumber(52))
1024                ],
1025                result
1026            );
1027        }
1028    }
1029
1030    #[tokio::test]
1031    async fn get_transactions_by_hashes_filters_correctly() {
1032        let blocks_with_tx = vec![
1033            CardanoBlockWithTransactions::new(
1034                "block-hash-1",
1035                BlockNumber(10),
1036                SlotNumber(50),
1037                vec!["tx-hash-1"],
1038            ),
1039            CardanoBlockWithTransactions::new(
1040                "block-hash-2",
1041                BlockNumber(11),
1042                SlotNumber(51),
1043                vec!["tx-hash-2"],
1044            ),
1045            CardanoBlockWithTransactions::new(
1046                "block-hash-3",
1047                BlockNumber(12),
1048                SlotNumber(52),
1049                vec!["tx-hash-3"],
1050            ),
1051        ];
1052        let store = InMemoryChainDataStore::builder()
1053            .with_blocks_and_transactions(&blocks_with_tx)
1054            .build();
1055
1056        // Get one existing hash
1057        {
1058            let result = store.get_transactions_by_hashes(&["tx-hash-1".to_string()]).await;
1059            assert_eq!(
1060                vec![CardanoTransaction::new(
1061                    "tx-hash-1",
1062                    BlockNumber(10),
1063                    SlotNumber(50),
1064                    "block-hash-1"
1065                )],
1066                result
1067            );
1068        }
1069        // Get one non-existing hash
1070        {
1071            let result = store.get_transactions_by_hashes(&["tx-hash-4".to_string()]).await;
1072            assert_eq!(Vec::<CardanoTransaction>::new(), result);
1073        }
1074        // Get two existing hashes and one non-existing hash
1075        {
1076            let result = store
1077                .get_transactions_by_hashes(&[
1078                    "tx-hash-1".to_string(),
1079                    "tx-hash-3".to_string(),
1080                    "tx-hash-4".to_string(),
1081                ])
1082                .await;
1083            assert_eq!(
1084                vec![
1085                    CardanoTransaction::new(
1086                        "tx-hash-1",
1087                        BlockNumber(10),
1088                        SlotNumber(50),
1089                        "block-hash-1"
1090                    ),
1091                    CardanoTransaction::new(
1092                        "tx-hash-3",
1093                        BlockNumber(12),
1094                        SlotNumber(52),
1095                        "block-hash-3"
1096                    )
1097                ],
1098                result
1099            );
1100        }
1101    }
1102
1103    #[tokio::test]
1104    async fn remove_rolled_back_chain_data_and_block_range_removes_transactions_above_slot_number()
1105    {
1106        let blocks_with_tx = vec![
1107            CardanoBlockWithTransactions::new(
1108                "block-hash-1",
1109                BlockNumber(10),
1110                SlotNumber(50),
1111                vec!["tx-hash-1"],
1112            ),
1113            CardanoBlockWithTransactions::new(
1114                "block-hash-2",
1115                BlockNumber(11),
1116                SlotNumber(51),
1117                vec!["tx-hash-2"],
1118            ),
1119            CardanoBlockWithTransactions::new(
1120                "block-hash-3",
1121                BlockNumber(12),
1122                SlotNumber(52),
1123                vec!["tx-hash-3"],
1124            ),
1125        ];
1126        let store = InMemoryChainDataStore::builder()
1127            .with_blocks_and_transactions(&blocks_with_tx)
1128            .build();
1129
1130        store
1131            .remove_rolled_chain_data_and_block_range(SlotNumber(51))
1132            .await
1133            .unwrap();
1134
1135        let remaining = store.get_all_block_with_txs().await;
1136        assert_eq!(blocks_with_tx[0..=1].to_vec(), remaining);
1137    }
1138
1139    #[tokio::test]
1140    async fn remove_rolled_back_chain_data_and_block_range_removes_block_ranges_above_highest_remaining_block()
1141     {
1142        let blocks_with_tx = vec![
1143            CardanoBlockWithTransactions::new(
1144                "block-hash-1",
1145                BlockNumber(10),
1146                SlotNumber(50),
1147                vec!["tx-hash-1"],
1148            ),
1149            CardanoBlockWithTransactions::new(
1150                "block-hash-2",
1151                BlockRange::LENGTH * 2,
1152                SlotNumber(100),
1153                vec!["tx-hash-2"],
1154            ),
1155            CardanoBlockWithTransactions::new(
1156                "block-hash-3",
1157                BlockRange::LENGTH * 4,
1158                SlotNumber(200),
1159                vec!["tx-hash-3"],
1160            ),
1161        ];
1162        let block_ranges_roots = vec![
1163            (
1164                BlockRange::from_block_number(BlockNumber(0)),
1165                MKTreeNode::from_hex("AAAA").unwrap(),
1166            ),
1167            (
1168                BlockRange::from_block_number(BlockRange::LENGTH),
1169                MKTreeNode::from_hex("BBBB").unwrap(),
1170            ),
1171            (
1172                BlockRange::from_block_number(BlockRange::LENGTH * 2),
1173                MKTreeNode::from_hex("CCCC").unwrap(),
1174            ),
1175            (
1176                BlockRange::from_block_number(BlockRange::LENGTH * 3),
1177                MKTreeNode::from_hex("DDDD").unwrap(),
1178            ),
1179        ];
1180        let store = InMemoryChainDataStore::builder()
1181            .with_blocks_and_transactions(&blocks_with_tx)
1182            .with_block_range_roots(&block_ranges_roots)
1183            .with_legacy_block_range_roots(&block_ranges_roots)
1184            .build();
1185
1186        // Rollback to slot 100 (keeps transactions with slot <= 100)
1187        store
1188            .remove_rolled_chain_data_and_block_range(SlotNumber(100))
1189            .await
1190            .unwrap();
1191
1192        let remaining_transactions = store.get_all_transactions().await;
1193        assert_eq!(2, remaining_transactions.len());
1194
1195        // Block ranges with start < highest remaining block number (BlockRange::LENGTH * 2) should remain
1196        let remaining_ranges = store.get_all_block_ranges().await;
1197        assert_eq!(
1198            vec![
1199                BlockRange::from_block_number(BlockNumber(0)),
1200                BlockRange::from_block_number(BlockRange::LENGTH),
1201            ],
1202            remaining_ranges
1203        );
1204
1205        // Legacy Block ranges with start < highest remaining block number (BlockRange::LENGTH * 2) should remain
1206        let remaining_legacy_ranges = store.get_all_legacy_block_ranges().await;
1207        assert_eq!(
1208            vec![
1209                BlockRange::from_block_number(BlockNumber(0)),
1210                BlockRange::from_block_number(BlockRange::LENGTH),
1211            ],
1212            remaining_legacy_ranges
1213        );
1214    }
1215
1216    #[tokio::test]
1217    async fn remove_rolled_back_chain_data_and_block_range_with_no_remaining_transactions() {
1218        let store = InMemoryChainDataStore::builder()
1219            .with_blocks_and_transactions(&[CardanoBlockWithTransactions::new(
1220                "block-hash-1",
1221                BlockNumber(10),
1222                SlotNumber(50),
1223                vec!["tx-hash-1"],
1224            )])
1225            .with_legacy_block_range_roots(&[(
1226                BlockRange::from_block_number(BlockNumber(0)),
1227                MKTreeNode::from_hex("AAAA").unwrap(),
1228            )])
1229            .build();
1230
1231        // Rollback to slot before all transactions
1232        store
1233            .remove_rolled_chain_data_and_block_range(SlotNumber(40))
1234            .await
1235            .unwrap();
1236
1237        assert!(store.get_all_block_with_txs().await.is_empty());
1238        assert!(store.get_all_transactions().await.is_empty());
1239        assert!(store.get_all_block_range_root().await.is_empty());
1240        assert!(store.get_all_legacy_block_range_root().await.is_empty());
1241    }
1242
1243    #[tokio::test]
1244    async fn remove_rolled_back_chain_data_and_block_range_keeps_transactions_with_equal_slot_number()
1245     {
1246        let blocks_with_tx = vec![
1247            CardanoBlockWithTransactions::new(
1248                "block-hash-1",
1249                BlockNumber(10),
1250                SlotNumber(50),
1251                vec!["tx-hash-1"],
1252            ),
1253            CardanoBlockWithTransactions::new(
1254                "block-hash-2",
1255                BlockNumber(11),
1256                SlotNumber(51),
1257                vec!["tx-hash-2"],
1258            ),
1259            CardanoBlockWithTransactions::new(
1260                "block-hash-3",
1261                BlockNumber(12),
1262                SlotNumber(52),
1263                vec!["tx-hash-3"],
1264            ),
1265        ];
1266        let store = InMemoryChainDataStore::builder()
1267            .with_blocks_and_transactions(&blocks_with_tx)
1268            .build();
1269
1270        store
1271            .remove_rolled_chain_data_and_block_range(SlotNumber(50))
1272            .await
1273            .unwrap();
1274
1275        let remaining = store.get_all_block_with_txs().await;
1276        assert_eq!(blocks_with_tx[0..1].to_vec(), remaining);
1277    }
1278
1279    mod retrieve_blocks_range_root {
1280        use super::*;
1281
1282        fn test_data_set() -> Vec<(BlockRange, MKTreeNode)> {
1283            vec![
1284                (
1285                    BlockRange::from_block_number(BlockNumber(15)),
1286                    MKTreeNode::from_hex("AAAA").unwrap(),
1287                ),
1288                (
1289                    BlockRange::from_block_number(BlockNumber(30)),
1290                    MKTreeNode::from_hex("BBBB").unwrap(),
1291                ),
1292                (
1293                    BlockRange::from_block_number(BlockNumber(45)),
1294                    MKTreeNode::from_hex("CCCC").unwrap(),
1295                ),
1296                (
1297                    BlockRange::from_block_number(BlockNumber(60)),
1298                    MKTreeNode::from_hex("CCCC").unwrap(),
1299                ),
1300            ]
1301        }
1302
1303        #[tokio::test]
1304        async fn returns_empty_when_store_empty() {
1305            let store = Arc::new(InMemoryChainDataStore::builder().build());
1306            let retriever = store.clone() as Arc<dyn BlockRangeRootRetriever<MKTreeStoreInMemory>>;
1307
1308            let iter = retriever
1309                .retrieve_block_range_roots(BlockNumber(u64::MAX))
1310                .await
1311                .unwrap();
1312            assert_eq!(
1313                Vec::<(BlockRange, MKTreeNode)>::new(),
1314                iter.collect::<Vec<_>>()
1315            );
1316        }
1317
1318        #[tokio::test]
1319        async fn up_to_above_all_stored_ranges_returns_all() {
1320            let store = Arc::new(
1321                InMemoryChainDataStore::builder()
1322                    .with_block_range_roots(&test_data_set())
1323                    .build(),
1324            );
1325            let retriever = store.clone() as Arc<dyn BlockRangeRootRetriever<MKTreeStoreInMemory>>;
1326
1327            let stored_ranges: Vec<(BlockRange, MKTreeNode)> = retriever
1328                .retrieve_block_range_roots(BlockNumber(u64::MAX))
1329                .await
1330                .unwrap()
1331                .collect();
1332            assert_eq!(&test_data_set(), &stored_ranges);
1333        }
1334
1335        #[tokio::test]
1336        async fn up_to_below_start_of_the_first_range_returns_nothing() {
1337            let store = Arc::new(
1338                InMemoryChainDataStore::builder()
1339                    .with_block_range_roots(&test_data_set())
1340                    .build(),
1341            );
1342            let retriever = store.clone() as Arc<dyn BlockRangeRootRetriever<MKTreeStoreInMemory>>;
1343
1344            let stored_ranges: Vec<(BlockRange, MKTreeNode)> = retriever
1345                .retrieve_block_range_roots(BlockNumber(10))
1346                .await
1347                .unwrap()
1348                .collect();
1349            assert_eq!(&Vec::<(BlockRange, MKTreeNode)>::new(), &stored_ranges);
1350        }
1351
1352        #[tokio::test]
1353        async fn up_to_right_below_start_of_the_third_range_returns_the_first_two_ranges() {
1354            let store = Arc::new(
1355                InMemoryChainDataStore::builder()
1356                    .with_block_range_roots(&test_data_set())
1357                    .build(),
1358            );
1359            let retriever = store.clone() as Arc<dyn BlockRangeRootRetriever<MKTreeStoreInMemory>>;
1360
1361            let stored_ranges: Vec<(BlockRange, MKTreeNode)> = retriever
1362                .retrieve_block_range_roots(BlockNumber(44))
1363                .await
1364                .unwrap()
1365                .collect();
1366            assert_eq!(&test_data_set()[0..2], &stored_ranges);
1367        }
1368
1369        #[tokio::test]
1370        async fn up_to_right_at_start_of_the_third_range_returns_the_first_two_ranges() {
1371            let store = Arc::new(
1372                InMemoryChainDataStore::builder()
1373                    .with_block_range_roots(&test_data_set())
1374                    .build(),
1375            );
1376            let retriever = store.clone() as Arc<dyn BlockRangeRootRetriever<MKTreeStoreInMemory>>;
1377
1378            let stored_ranges: Vec<(BlockRange, MKTreeNode)> = retriever
1379                .retrieve_block_range_roots(BlockNumber(45))
1380                .await
1381                .unwrap()
1382                .collect();
1383            assert_eq!(&test_data_set()[0..2], &stored_ranges);
1384        }
1385
1386        #[tokio::test]
1387        async fn up_to_right_after_start_of_the_third_range_returns_the_first_three_ranges() {
1388            let store = Arc::new(
1389                InMemoryChainDataStore::builder()
1390                    .with_block_range_roots(&test_data_set())
1391                    .build(),
1392            );
1393            let retriever = store.clone() as Arc<dyn BlockRangeRootRetriever<MKTreeStoreInMemory>>;
1394
1395            let stored_ranges: Vec<(BlockRange, MKTreeNode)> = retriever
1396                .retrieve_block_range_roots(BlockNumber(46))
1397                .await
1398                .unwrap()
1399                .collect();
1400            assert_eq!(&test_data_set()[0..3], &stored_ranges);
1401        }
1402    }
1403}