1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
//! Support for maintaining the blob pool.

use alloy_primitives::{BlockNumber, B256};
use reth_execution_types::ChainBlocks;
use std::collections::BTreeMap;

/// The type that is used to track canonical blob transactions.
#[derive(Debug, Default, Eq, PartialEq)]
pub struct BlobStoreCanonTracker {
    /// Keeps track of the blob transactions included in blocks.
    blob_txs_in_blocks: BTreeMap<BlockNumber, Vec<B256>>,
}

impl BlobStoreCanonTracker {
    /// Adds a block to the blob store maintenance.
    pub fn add_block(
        &mut self,
        block_number: BlockNumber,
        blob_txs: impl IntoIterator<Item = B256>,
    ) {
        self.blob_txs_in_blocks.insert(block_number, blob_txs.into_iter().collect());
    }

    /// Adds all blocks to the tracked list of blocks.
    ///
    /// Replaces any previously tracked blocks with the set of transactions.
    pub fn add_blocks(
        &mut self,
        blocks: impl IntoIterator<Item = (BlockNumber, impl IntoIterator<Item = B256>)>,
    ) {
        for (block_number, blob_txs) in blocks {
            self.add_block(block_number, blob_txs);
        }
    }

    /// Adds all blob transactions from the given chain to the tracker.
    ///
    /// Note: In case this is a chain that's part of a reorg, this replaces previously tracked
    /// blocks.
    pub fn add_new_chain_blocks(&mut self, blocks: &ChainBlocks<'_>) {
        let blob_txs = blocks.iter().map(|(num, blocks)| {
            let iter =
                blocks.body.iter().filter(|tx| tx.transaction.is_eip4844()).map(|tx| tx.hash);
            (*num, iter)
        });
        self.add_blocks(blob_txs);
    }

    /// Invoked when a block is finalized.
    ///
    /// This returns all blob transactions that were included in blocks that are now finalized.
    pub fn on_finalized_block(&mut self, finalized_block: BlockNumber) -> BlobStoreUpdates {
        let mut finalized = Vec::new();
        while let Some(entry) = self.blob_txs_in_blocks.first_entry() {
            if *entry.key() <= finalized_block {
                finalized.extend(entry.remove_entry().1);
            } else {
                break
            }
        }

        if finalized.is_empty() {
            BlobStoreUpdates::None
        } else {
            BlobStoreUpdates::Finalized(finalized)
        }
    }
}

/// Updates that should be applied to the blob store.
#[derive(Debug, Eq, PartialEq)]
pub enum BlobStoreUpdates {
    /// No updates.
    None,
    /// Delete the given finalized transactions from the blob store.
    Finalized(Vec<B256>),
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_finalized_tracker() {
        let mut tracker = BlobStoreCanonTracker::default();

        let block1 = vec![B256::random()];
        let block2 = vec![B256::random()];
        let block3 = vec![B256::random()];
        tracker.add_block(1, block1.clone());
        tracker.add_block(2, block2.clone());
        tracker.add_block(3, block3.clone());

        assert_eq!(tracker.on_finalized_block(0), BlobStoreUpdates::None);
        assert_eq!(tracker.on_finalized_block(1), BlobStoreUpdates::Finalized(block1));
        assert_eq!(
            tracker.on_finalized_block(3),
            BlobStoreUpdates::Finalized(block2.into_iter().chain(block3).collect::<Vec<_>>())
        );
    }
}