reth_blockchain_tree/
canonical_chain.rs

1use alloy_eips::BlockNumHash;
2use alloy_primitives::{BlockHash, BlockNumber};
3use std::collections::BTreeMap;
4
5/// This keeps track of (non-finalized) blocks of the canonical chain.
6///
7/// This is a wrapper type around an ordered set of block numbers and hashes that belong to the
8/// canonical chain that is not yet finalized.
9#[derive(Debug, Clone, Default)]
10pub(crate) struct CanonicalChain {
11    /// All blocks of the canonical chain in order of their block number.
12    chain: BTreeMap<BlockNumber, BlockHash>,
13}
14
15impl CanonicalChain {
16    pub(crate) const fn new(chain: BTreeMap<BlockNumber, BlockHash>) -> Self {
17        Self { chain }
18    }
19
20    /// Replaces the current chain with the given one.
21    #[inline]
22    pub(crate) fn replace(&mut self, chain: BTreeMap<BlockNumber, BlockHash>) {
23        self.chain = chain;
24    }
25
26    /// Returns the block hash of the (non-finalized) canonical block with the given number.
27    #[inline]
28    pub(crate) fn canonical_hash(&self, number: &BlockNumber) -> Option<BlockHash> {
29        self.chain.get(number).copied()
30    }
31
32    /// Returns the block number of the (non-finalized) canonical block with the given hash.
33    #[inline]
34    pub(crate) fn canonical_number(&self, block_hash: &BlockHash) -> Option<BlockNumber> {
35        self.chain.iter().find_map(|(number, hash)| (hash == block_hash).then_some(*number))
36    }
37
38    /// Extends all items from the given iterator to the chain.
39    #[inline]
40    pub(crate) fn extend(&mut self, blocks: impl Iterator<Item = (BlockNumber, BlockHash)>) {
41        self.chain.extend(blocks)
42    }
43
44    /// Retains only the elements specified by the predicate.
45    #[inline]
46    pub(crate) fn retain<F>(&mut self, f: F)
47    where
48        F: FnMut(&BlockNumber, &mut BlockHash) -> bool,
49    {
50        self.chain.retain(f)
51    }
52
53    #[inline]
54    pub(crate) const fn inner(&self) -> &BTreeMap<BlockNumber, BlockHash> {
55        &self.chain
56    }
57
58    #[inline]
59    pub(crate) fn tip(&self) -> BlockNumHash {
60        self.chain
61            .last_key_value()
62            .map(|(&number, &hash)| BlockNumHash { number, hash })
63            .unwrap_or_default()
64    }
65
66    #[inline]
67    pub(crate) fn iter(&self) -> impl Iterator<Item = (BlockNumber, BlockHash)> + '_ {
68        self.chain.iter().map(|(&number, &hash)| (number, hash))
69    }
70
71    #[inline]
72    pub(crate) fn into_iter(self) -> impl Iterator<Item = (BlockNumber, BlockHash)> {
73        self.chain.into_iter()
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    #[test]
82    fn test_replace_canonical_chain() {
83        // Initialize a chain with some blocks
84        let mut initial_chain = BTreeMap::new();
85        initial_chain.insert(BlockNumber::from(1u64), BlockHash::from([0x01; 32]));
86        initial_chain.insert(BlockNumber::from(2u64), BlockHash::from([0x02; 32]));
87
88        let mut canonical_chain = CanonicalChain::new(initial_chain.clone());
89
90        // Verify initial chain state
91        assert_eq!(canonical_chain.chain.len(), 2);
92        assert_eq!(
93            canonical_chain.chain.get(&BlockNumber::from(1u64)),
94            Some(&BlockHash::from([0x01; 32]))
95        );
96
97        // Replace with a new chain
98        let mut new_chain = BTreeMap::new();
99        new_chain.insert(BlockNumber::from(3u64), BlockHash::from([0x03; 32]));
100        new_chain.insert(BlockNumber::from(4u64), BlockHash::from([0x04; 32]));
101        new_chain.insert(BlockNumber::from(5u64), BlockHash::from([0x05; 32]));
102
103        canonical_chain.replace(new_chain.clone());
104
105        // Verify replaced chain state
106        assert_eq!(canonical_chain.chain.len(), 3);
107        assert!(!canonical_chain.chain.contains_key(&BlockNumber::from(1u64)));
108        assert_eq!(
109            canonical_chain.chain.get(&BlockNumber::from(3u64)),
110            Some(&BlockHash::from([0x03; 32]))
111        );
112    }
113
114    #[test]
115    fn test_canonical_hash_canonical_chain() {
116        // Initialize a chain with some blocks
117        let mut chain = BTreeMap::new();
118        chain.insert(BlockNumber::from(1u64), BlockHash::from([0x01; 32]));
119        chain.insert(BlockNumber::from(2u64), BlockHash::from([0x02; 32]));
120        chain.insert(BlockNumber::from(3u64), BlockHash::from([0x03; 32]));
121
122        // Create an instance of a canonical chain
123        let canonical_chain = CanonicalChain::new(chain.clone());
124
125        // Check that the function returns the correct hash for a given block number
126        let block_number = BlockNumber::from(2u64);
127        let expected_hash = BlockHash::from([0x02; 32]);
128        assert_eq!(canonical_chain.canonical_hash(&block_number), Some(expected_hash));
129
130        // Check that a non-existent block returns None
131        let non_existent_block = BlockNumber::from(5u64);
132        assert_eq!(canonical_chain.canonical_hash(&non_existent_block), None);
133    }
134
135    #[test]
136    fn test_canonical_number_canonical_chain() {
137        // Initialize a chain with some blocks
138        let mut chain = BTreeMap::new();
139        chain.insert(BlockNumber::from(1u64), BlockHash::from([0x01; 32]));
140        chain.insert(BlockNumber::from(2u64), BlockHash::from([0x02; 32]));
141        chain.insert(BlockNumber::from(3u64), BlockHash::from([0x03; 32]));
142
143        // Create an instance of a canonical chain
144        let canonical_chain = CanonicalChain::new(chain.clone());
145
146        // Check that the function returns the correct block number for a given block hash
147        let block_hash = BlockHash::from([0x02; 32]);
148        let expected_number = BlockNumber::from(2u64);
149        assert_eq!(canonical_chain.canonical_number(&block_hash), Some(expected_number));
150
151        // Check that a non-existent block hash returns None
152        let non_existent_hash = BlockHash::from([0x05; 32]);
153        assert_eq!(canonical_chain.canonical_number(&non_existent_hash), None);
154    }
155
156    #[test]
157    fn test_extend_canonical_chain() {
158        // Initialize an empty chain
159        let mut canonical_chain = CanonicalChain::new(BTreeMap::new());
160
161        // Create an iterator with some blocks
162        let blocks = vec![
163            (BlockNumber::from(1u64), BlockHash::from([0x01; 32])),
164            (BlockNumber::from(2u64), BlockHash::from([0x02; 32])),
165        ]
166        .into_iter();
167
168        // Extend the chain with the created blocks
169        canonical_chain.extend(blocks);
170
171        // Check if the blocks were added correctly
172        assert_eq!(canonical_chain.chain.len(), 2);
173        assert_eq!(
174            canonical_chain.chain.get(&BlockNumber::from(1u64)),
175            Some(&BlockHash::from([0x01; 32]))
176        );
177        assert_eq!(
178            canonical_chain.chain.get(&BlockNumber::from(2u64)),
179            Some(&BlockHash::from([0x02; 32]))
180        );
181
182        // Test extending with additional blocks again
183        let more_blocks = vec![(BlockNumber::from(3u64), BlockHash::from([0x03; 32]))].into_iter();
184        canonical_chain.extend(more_blocks);
185
186        assert_eq!(canonical_chain.chain.len(), 3);
187        assert_eq!(
188            canonical_chain.chain.get(&BlockNumber::from(3u64)),
189            Some(&BlockHash::from([0x03; 32]))
190        );
191    }
192
193    #[test]
194    fn test_retain_canonical_chain() {
195        // Initialize a chain with some blocks
196        let mut chain = BTreeMap::new();
197        chain.insert(BlockNumber::from(1u64), BlockHash::from([0x01; 32]));
198        chain.insert(BlockNumber::from(2u64), BlockHash::from([0x02; 32]));
199        chain.insert(BlockNumber::from(3u64), BlockHash::from([0x03; 32]));
200
201        // Create an instance of CanonicalChain
202        let mut canonical_chain = CanonicalChain::new(chain);
203
204        // Retain only blocks with even block numbers
205        canonical_chain.retain(|number, _| number % 2 == 0);
206
207        // Check if the chain only contains the block with number 2
208        assert_eq!(canonical_chain.chain.len(), 1);
209        assert_eq!(
210            canonical_chain.chain.get(&BlockNumber::from(2u64)),
211            Some(&BlockHash::from([0x02; 32]))
212        );
213
214        // Ensure that the blocks with odd numbers were removed
215        assert_eq!(canonical_chain.chain.get(&BlockNumber::from(1u64)), None);
216        assert_eq!(canonical_chain.chain.get(&BlockNumber::from(3u64)), None);
217    }
218
219    #[test]
220    fn test_tip_canonical_chain() {
221        // Initialize a chain with some blocks
222        let mut chain = BTreeMap::new();
223        chain.insert(BlockNumber::from(1u64), BlockHash::from([0x01; 32]));
224        chain.insert(BlockNumber::from(2u64), BlockHash::from([0x02; 32]));
225        chain.insert(BlockNumber::from(3u64), BlockHash::from([0x03; 32]));
226
227        // Create an instance of a canonical chain
228        let canonical_chain = CanonicalChain::new(chain);
229
230        // Call the tip method and verify the returned value
231        let tip = canonical_chain.tip();
232        assert_eq!(tip.number, BlockNumber::from(3u64));
233        assert_eq!(tip.hash, BlockHash::from([0x03; 32]));
234
235        // Test with an empty chain
236        let empty_chain = CanonicalChain::new(BTreeMap::new());
237        let empty_tip = empty_chain.tip();
238        assert_eq!(empty_tip.number, BlockNumber::default());
239        assert_eq!(empty_tip.hash, BlockHash::default());
240    }
241}
reth_blockchain_tree/canonical_chain.rs

reth_blockchain_tree/
canonical_chain.rs
