1use crate::engine::EngineApiKind;
4use alloy_eips::{eip1898::BlockWithParent, merge::EPOCH_SLOTS, BlockNumHash};
5use alloy_primitives::{
6 map::{HashMap, HashSet},
7 BlockNumber, B256,
8};
9use reth_chain_state::{EthPrimitives, ExecutedBlockWithTrieUpdates};
10use reth_primitives_traits::{AlloyBlockHeader, NodePrimitives, SealedHeader};
11use reth_trie::updates::TrieUpdates;
12use std::{
13 collections::{btree_map, hash_map, BTreeMap, VecDeque},
14 ops::Bound,
15 sync::Arc,
16};
17use tracing::debug;
18
19const DEFAULT_PERSISTED_TRIE_UPDATES_RETENTION: u64 = EPOCH_SLOTS * 2;
21
22const OPSTACK_PERSISTED_TRIE_UPDATES_RETENTION: u64 = EPOCH_SLOTS;
26
27#[derive(Debug, Default)]
34pub struct TreeState<N: NodePrimitives = EthPrimitives> {
35 pub(crate) blocks_by_hash: HashMap<B256, ExecutedBlockWithTrieUpdates<N>>,
39 pub(crate) blocks_by_number: BTreeMap<BlockNumber, Vec<ExecutedBlockWithTrieUpdates<N>>>,
45 pub(crate) parent_to_child: HashMap<B256, HashSet<B256>>,
47 pub(crate) persisted_trie_updates: HashMap<B256, (BlockNumber, Arc<TrieUpdates>)>,
51 pub(crate) current_canonical_head: BlockNumHash,
53 pub(crate) engine_kind: EngineApiKind,
55}
56
57impl<N: NodePrimitives> TreeState<N> {
58 pub(crate) fn new(current_canonical_head: BlockNumHash, engine_kind: EngineApiKind) -> Self {
60 Self {
61 blocks_by_hash: HashMap::default(),
62 blocks_by_number: BTreeMap::new(),
63 current_canonical_head,
64 parent_to_child: HashMap::default(),
65 persisted_trie_updates: HashMap::default(),
66 engine_kind,
67 }
68 }
69
70 pub(crate) fn reset(&mut self, current_canonical_head: BlockNumHash) {
72 *self = Self::new(current_canonical_head, self.engine_kind);
73 }
74
75 pub(crate) fn block_count(&self) -> usize {
77 self.blocks_by_hash.len()
78 }
79
80 pub(crate) fn executed_block_by_hash(
82 &self,
83 hash: B256,
84 ) -> Option<&ExecutedBlockWithTrieUpdates<N>> {
85 self.blocks_by_hash.get(&hash)
86 }
87
88 pub(crate) fn sealed_header_by_hash(
90 &self,
91 hash: &B256,
92 ) -> Option<SealedHeader<N::BlockHeader>> {
93 self.blocks_by_hash.get(hash).map(|b| b.sealed_block().sealed_header().clone())
94 }
95
96 pub(crate) fn blocks_by_hash(
101 &self,
102 hash: B256,
103 ) -> Option<(B256, Vec<ExecutedBlockWithTrieUpdates<N>>)> {
104 let block = self.blocks_by_hash.get(&hash).cloned()?;
105 let mut parent_hash = block.recovered_block().parent_hash();
106 let mut blocks = vec![block];
107 while let Some(executed) = self.blocks_by_hash.get(&parent_hash) {
108 parent_hash = executed.recovered_block().parent_hash();
109 blocks.push(executed.clone());
110 }
111
112 Some((parent_hash, blocks))
113 }
114
115 pub(crate) fn insert_executed(&mut self, executed: ExecutedBlockWithTrieUpdates<N>) {
117 let hash = executed.recovered_block().hash();
118 let parent_hash = executed.recovered_block().parent_hash();
119 let block_number = executed.recovered_block().number();
120
121 if self.blocks_by_hash.contains_key(&hash) {
122 return;
123 }
124
125 self.blocks_by_hash.insert(hash, executed.clone());
126
127 self.blocks_by_number.entry(block_number).or_default().push(executed);
128
129 self.parent_to_child.entry(parent_hash).or_default().insert(hash);
130
131 for children in self.parent_to_child.values_mut() {
132 children.retain(|child| self.blocks_by_hash.contains_key(child));
133 }
134 }
135
136 fn remove_by_hash(
142 &mut self,
143 hash: B256,
144 ) -> Option<(ExecutedBlockWithTrieUpdates<N>, HashSet<B256>)> {
145 let executed = self.blocks_by_hash.remove(&hash)?;
146
147 let parent_entry = self.parent_to_child.entry(executed.recovered_block().parent_hash());
149 if let hash_map::Entry::Occupied(mut entry) = parent_entry {
150 entry.get_mut().remove(&hash);
151
152 if entry.get().is_empty() {
153 entry.remove();
154 }
155 }
156
157 let children = self.parent_to_child.remove(&hash).unwrap_or_default();
159
160 let block_number_entry = self.blocks_by_number.entry(executed.recovered_block().number());
162 if let btree_map::Entry::Occupied(mut entry) = block_number_entry {
163 if let Some(index) = entry.get().iter().position(|b| b.recovered_block().hash() == hash)
165 {
166 entry.get_mut().swap_remove(index);
167
168 if entry.get().is_empty() {
170 entry.remove();
171 }
172 }
173 }
174
175 Some((executed, children))
176 }
177
178 pub(crate) fn is_canonical(&self, hash: B256) -> bool {
180 let mut current_block = self.current_canonical_head.hash;
181 if current_block == hash {
182 return true
183 }
184
185 while let Some(executed) = self.blocks_by_hash.get(¤t_block) {
186 current_block = executed.recovered_block().parent_hash();
187 if current_block == hash {
188 return true
189 }
190 }
191
192 false
193 }
194
195 pub(crate) fn remove_canonical_until(
198 &mut self,
199 upper_bound: BlockNumber,
200 last_persisted_hash: B256,
201 ) {
202 debug!(target: "engine::tree", ?upper_bound, ?last_persisted_hash, "Removing canonical blocks from the tree");
203
204 if !self.is_canonical(last_persisted_hash) {
207 return
208 }
209
210 let mut current_block = self.current_canonical_head.hash;
213 while let Some(executed) = self.blocks_by_hash.get(¤t_block) {
214 current_block = executed.recovered_block().parent_hash();
215 if executed.recovered_block().number() <= upper_bound {
216 let num_hash = executed.recovered_block().num_hash();
217 debug!(target: "engine::tree", ?num_hash, "Attempting to remove block walking back from the head");
218 if let Some((mut removed, _)) =
219 self.remove_by_hash(executed.recovered_block().hash())
220 {
221 debug!(target: "engine::tree", ?num_hash, "Removed block walking back from the head");
222 let Some(trie_updates) = removed.trie.take_present() else {
224 debug!(target: "engine::tree", ?num_hash, "No trie updates found for persisted block");
225 continue;
226 };
227 self.persisted_trie_updates.insert(
228 removed.recovered_block().hash(),
229 (removed.recovered_block().number(), trie_updates),
230 );
231 }
232 }
233 }
234 debug!(target: "engine::tree", ?upper_bound, ?last_persisted_hash, "Removed canonical blocks from the tree");
235 }
236
237 pub(crate) fn prune_persisted_trie_updates(&mut self) {
240 let retention_blocks = if self.engine_kind.is_opstack() {
241 OPSTACK_PERSISTED_TRIE_UPDATES_RETENTION
242 } else {
243 DEFAULT_PERSISTED_TRIE_UPDATES_RETENTION
244 };
245
246 let earliest_block_to_retain =
247 self.current_canonical_head.number.saturating_sub(retention_blocks);
248
249 self.persisted_trie_updates
250 .retain(|_, (block_number, _)| *block_number > earliest_block_to_retain);
251 }
252
253 pub(crate) fn prune_finalized_sidechains(&mut self, finalized_num_hash: BlockNumHash) {
256 let BlockNumHash { number: finalized_num, hash: finalized_hash } = finalized_num_hash;
257
258 let blocks_to_remove = self
267 .blocks_by_number
268 .range((Bound::Unbounded, Bound::Excluded(finalized_num)))
269 .flat_map(|(_, blocks)| blocks.iter().map(|b| b.recovered_block().hash()))
270 .collect::<Vec<_>>();
271 for hash in blocks_to_remove {
272 if let Some((removed, _)) = self.remove_by_hash(hash) {
273 debug!(target: "engine::tree", num_hash=?removed.recovered_block().num_hash(), "Removed finalized sidechain block");
274 }
275 }
276
277 self.prune_persisted_trie_updates();
278
279 let mut blocks_to_remove = self.blocks_by_number.remove(&finalized_num).unwrap_or_default();
286
287 if let Some(position) =
289 blocks_to_remove.iter().position(|b| b.recovered_block().hash() == finalized_hash)
290 {
291 let finalized_block = blocks_to_remove.swap_remove(position);
292 self.blocks_by_number.insert(finalized_num, vec![finalized_block]);
293 }
294
295 let mut blocks_to_remove = blocks_to_remove
296 .into_iter()
297 .map(|e| e.recovered_block().hash())
298 .collect::<VecDeque<_>>();
299 while let Some(block) = blocks_to_remove.pop_front() {
300 if let Some((removed, children)) = self.remove_by_hash(block) {
301 debug!(target: "engine::tree", num_hash=?removed.recovered_block().num_hash(), "Removed finalized sidechain child block");
302 blocks_to_remove.extend(children);
303 }
304 }
305 }
306
307 pub(crate) fn remove_until(
318 &mut self,
319 upper_bound: BlockNumHash,
320 last_persisted_hash: B256,
321 finalized_num_hash: Option<BlockNumHash>,
322 ) {
323 debug!(target: "engine::tree", ?upper_bound, ?finalized_num_hash, "Removing blocks from the tree");
324
325 let finalized_num_hash = finalized_num_hash.map(|mut finalized| {
328 if upper_bound.number < finalized.number {
329 finalized = upper_bound;
330 debug!(target: "engine::tree", ?finalized, "Adjusted upper bound");
331 }
332 finalized
333 });
334
335 self.remove_canonical_until(upper_bound.number, last_persisted_hash);
343
344 if let Some(finalized_num_hash) = finalized_num_hash {
347 self.prune_finalized_sidechains(finalized_num_hash);
348 }
349 }
350
351 pub(crate) fn is_descendant(&self, first: BlockNumHash, second: BlockWithParent) -> bool {
355 if second.parent == first.hash {
358 return true
359 }
360
361 if second.block.number <= first.number {
364 return false
365 }
366
367 let Some(mut current_block) = self.blocks_by_hash.get(&second.parent) else {
369 return false
371 };
372
373 while current_block.recovered_block().number() > first.number + 1 {
374 let Some(block) =
375 self.blocks_by_hash.get(¤t_block.recovered_block().parent_hash())
376 else {
377 return false
379 };
380
381 current_block = block;
382 }
383
384 current_block.recovered_block().parent_hash() == first.hash
386 }
387
388 pub(crate) const fn set_canonical_head(&mut self, new_head: BlockNumHash) {
390 self.current_canonical_head = new_head;
391 }
392
393 pub(crate) const fn canonical_head(&self) -> &BlockNumHash {
395 &self.current_canonical_head
396 }
397
398 pub(crate) const fn canonical_block_hash(&self) -> B256 {
400 self.canonical_head().hash
401 }
402
403 pub(crate) const fn canonical_block_number(&self) -> BlockNumber {
405 self.canonical_head().number
406 }
407}
408
409#[cfg(test)]
410mod tests {
411 use super::*;
412 use reth_chain_state::test_utils::TestBlockBuilder;
413
414 #[test]
415 fn test_tree_state_normal_descendant() {
416 let mut tree_state = TreeState::new(BlockNumHash::default(), EngineApiKind::Ethereum);
417 let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..4).collect();
418
419 tree_state.insert_executed(blocks[0].clone());
420 assert!(tree_state.is_descendant(
421 blocks[0].recovered_block().num_hash(),
422 blocks[1].recovered_block().block_with_parent()
423 ));
424
425 tree_state.insert_executed(blocks[1].clone());
426
427 assert!(tree_state.is_descendant(
428 blocks[0].recovered_block().num_hash(),
429 blocks[2].recovered_block().block_with_parent()
430 ));
431 assert!(tree_state.is_descendant(
432 blocks[1].recovered_block().num_hash(),
433 blocks[2].recovered_block().block_with_parent()
434 ));
435 }
436
437 #[tokio::test]
438 async fn test_tree_state_insert_executed() {
439 let mut tree_state = TreeState::new(BlockNumHash::default(), EngineApiKind::Ethereum);
440 let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..4).collect();
441
442 tree_state.insert_executed(blocks[0].clone());
443 tree_state.insert_executed(blocks[1].clone());
444
445 assert_eq!(
446 tree_state.parent_to_child.get(&blocks[0].recovered_block().hash()),
447 Some(&HashSet::from_iter([blocks[1].recovered_block().hash()]))
448 );
449
450 assert!(!tree_state.parent_to_child.contains_key(&blocks[1].recovered_block().hash()));
451
452 tree_state.insert_executed(blocks[2].clone());
453
454 assert_eq!(
455 tree_state.parent_to_child.get(&blocks[1].recovered_block().hash()),
456 Some(&HashSet::from_iter([blocks[2].recovered_block().hash()]))
457 );
458 assert!(tree_state.parent_to_child.contains_key(&blocks[1].recovered_block().hash()));
459
460 assert!(!tree_state.parent_to_child.contains_key(&blocks[2].recovered_block().hash()));
461 }
462
463 #[tokio::test]
464 async fn test_tree_state_insert_executed_with_reorg() {
465 let mut tree_state = TreeState::new(BlockNumHash::default(), EngineApiKind::Ethereum);
466 let mut test_block_builder = TestBlockBuilder::eth();
467 let blocks: Vec<_> = test_block_builder.get_executed_blocks(1..6).collect();
468
469 for block in &blocks {
470 tree_state.insert_executed(block.clone());
471 }
472 assert_eq!(tree_state.blocks_by_hash.len(), 5);
473
474 let fork_block_3 = test_block_builder
475 .get_executed_block_with_number(3, blocks[1].recovered_block().hash());
476 let fork_block_4 = test_block_builder
477 .get_executed_block_with_number(4, fork_block_3.recovered_block().hash());
478 let fork_block_5 = test_block_builder
479 .get_executed_block_with_number(5, fork_block_4.recovered_block().hash());
480
481 tree_state.insert_executed(fork_block_3.clone());
482 tree_state.insert_executed(fork_block_4.clone());
483 tree_state.insert_executed(fork_block_5.clone());
484
485 assert_eq!(tree_state.blocks_by_hash.len(), 8);
486 assert_eq!(tree_state.blocks_by_number[&3].len(), 2); assert_eq!(tree_state.parent_to_child[&blocks[1].recovered_block().hash()].len(), 2); tree_state.insert_executed(fork_block_4.clone());
491 assert_eq!(tree_state.blocks_by_hash.len(), 8);
492
493 assert!(tree_state.parent_to_child[&fork_block_3.recovered_block().hash()]
494 .contains(&fork_block_4.recovered_block().hash()));
495 assert!(tree_state.parent_to_child[&fork_block_4.recovered_block().hash()]
496 .contains(&fork_block_5.recovered_block().hash()));
497
498 assert_eq!(tree_state.blocks_by_number[&4].len(), 2);
499 assert_eq!(tree_state.blocks_by_number[&5].len(), 2);
500 }
501
502 #[tokio::test]
503 async fn test_tree_state_remove_before() {
504 let start_num_hash = BlockNumHash::default();
505 let mut tree_state = TreeState::new(start_num_hash, EngineApiKind::Ethereum);
506 let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..6).collect();
507
508 for block in &blocks {
509 tree_state.insert_executed(block.clone());
510 }
511
512 let last = blocks.last().unwrap();
513
514 tree_state.set_canonical_head(last.recovered_block().num_hash());
516
517 tree_state.remove_until(
519 BlockNumHash::new(2, blocks[1].recovered_block().hash()),
520 start_num_hash.hash,
521 Some(blocks[1].recovered_block().num_hash()),
522 );
523
524 assert!(!tree_state.blocks_by_hash.contains_key(&blocks[0].recovered_block().hash()));
525 assert!(!tree_state.blocks_by_hash.contains_key(&blocks[1].recovered_block().hash()));
526 assert!(!tree_state.blocks_by_number.contains_key(&1));
527 assert!(!tree_state.blocks_by_number.contains_key(&2));
528
529 assert!(tree_state.blocks_by_hash.contains_key(&blocks[2].recovered_block().hash()));
530 assert!(tree_state.blocks_by_hash.contains_key(&blocks[3].recovered_block().hash()));
531 assert!(tree_state.blocks_by_hash.contains_key(&blocks[4].recovered_block().hash()));
532 assert!(tree_state.blocks_by_number.contains_key(&3));
533 assert!(tree_state.blocks_by_number.contains_key(&4));
534 assert!(tree_state.blocks_by_number.contains_key(&5));
535
536 assert!(!tree_state.parent_to_child.contains_key(&blocks[0].recovered_block().hash()));
537 assert!(!tree_state.parent_to_child.contains_key(&blocks[1].recovered_block().hash()));
538 assert!(tree_state.parent_to_child.contains_key(&blocks[2].recovered_block().hash()));
539 assert!(tree_state.parent_to_child.contains_key(&blocks[3].recovered_block().hash()));
540 assert!(!tree_state.parent_to_child.contains_key(&blocks[4].recovered_block().hash()));
541
542 assert_eq!(
543 tree_state.parent_to_child.get(&blocks[2].recovered_block().hash()),
544 Some(&HashSet::from_iter([blocks[3].recovered_block().hash()]))
545 );
546 assert_eq!(
547 tree_state.parent_to_child.get(&blocks[3].recovered_block().hash()),
548 Some(&HashSet::from_iter([blocks[4].recovered_block().hash()]))
549 );
550 }
551
552 #[tokio::test]
553 async fn test_tree_state_remove_before_finalized() {
554 let start_num_hash = BlockNumHash::default();
555 let mut tree_state = TreeState::new(start_num_hash, EngineApiKind::Ethereum);
556 let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..6).collect();
557
558 for block in &blocks {
559 tree_state.insert_executed(block.clone());
560 }
561
562 let last = blocks.last().unwrap();
563
564 tree_state.set_canonical_head(last.recovered_block().num_hash());
566
567 tree_state.remove_until(
569 BlockNumHash::new(2, blocks[1].recovered_block().hash()),
570 start_num_hash.hash,
571 None,
572 );
573
574 assert!(!tree_state.blocks_by_hash.contains_key(&blocks[0].recovered_block().hash()));
575 assert!(!tree_state.blocks_by_hash.contains_key(&blocks[1].recovered_block().hash()));
576 assert!(!tree_state.blocks_by_number.contains_key(&1));
577 assert!(!tree_state.blocks_by_number.contains_key(&2));
578
579 assert!(tree_state.blocks_by_hash.contains_key(&blocks[2].recovered_block().hash()));
580 assert!(tree_state.blocks_by_hash.contains_key(&blocks[3].recovered_block().hash()));
581 assert!(tree_state.blocks_by_hash.contains_key(&blocks[4].recovered_block().hash()));
582 assert!(tree_state.blocks_by_number.contains_key(&3));
583 assert!(tree_state.blocks_by_number.contains_key(&4));
584 assert!(tree_state.blocks_by_number.contains_key(&5));
585
586 assert!(!tree_state.parent_to_child.contains_key(&blocks[0].recovered_block().hash()));
587 assert!(!tree_state.parent_to_child.contains_key(&blocks[1].recovered_block().hash()));
588 assert!(tree_state.parent_to_child.contains_key(&blocks[2].recovered_block().hash()));
589 assert!(tree_state.parent_to_child.contains_key(&blocks[3].recovered_block().hash()));
590 assert!(!tree_state.parent_to_child.contains_key(&blocks[4].recovered_block().hash()));
591
592 assert_eq!(
593 tree_state.parent_to_child.get(&blocks[2].recovered_block().hash()),
594 Some(&HashSet::from_iter([blocks[3].recovered_block().hash()]))
595 );
596 assert_eq!(
597 tree_state.parent_to_child.get(&blocks[3].recovered_block().hash()),
598 Some(&HashSet::from_iter([blocks[4].recovered_block().hash()]))
599 );
600 }
601
602 #[tokio::test]
603 async fn test_tree_state_remove_before_lower_finalized() {
604 let start_num_hash = BlockNumHash::default();
605 let mut tree_state = TreeState::new(start_num_hash, EngineApiKind::Ethereum);
606 let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..6).collect();
607
608 for block in &blocks {
609 tree_state.insert_executed(block.clone());
610 }
611
612 let last = blocks.last().unwrap();
613
614 tree_state.set_canonical_head(last.recovered_block().num_hash());
616
617 tree_state.remove_until(
619 BlockNumHash::new(2, blocks[1].recovered_block().hash()),
620 start_num_hash.hash,
621 Some(blocks[0].recovered_block().num_hash()),
622 );
623
624 assert!(!tree_state.blocks_by_hash.contains_key(&blocks[0].recovered_block().hash()));
625 assert!(!tree_state.blocks_by_hash.contains_key(&blocks[1].recovered_block().hash()));
626 assert!(!tree_state.blocks_by_number.contains_key(&1));
627 assert!(!tree_state.blocks_by_number.contains_key(&2));
628
629 assert!(tree_state.blocks_by_hash.contains_key(&blocks[2].recovered_block().hash()));
630 assert!(tree_state.blocks_by_hash.contains_key(&blocks[3].recovered_block().hash()));
631 assert!(tree_state.blocks_by_hash.contains_key(&blocks[4].recovered_block().hash()));
632 assert!(tree_state.blocks_by_number.contains_key(&3));
633 assert!(tree_state.blocks_by_number.contains_key(&4));
634 assert!(tree_state.blocks_by_number.contains_key(&5));
635
636 assert!(!tree_state.parent_to_child.contains_key(&blocks[0].recovered_block().hash()));
637 assert!(!tree_state.parent_to_child.contains_key(&blocks[1].recovered_block().hash()));
638 assert!(tree_state.parent_to_child.contains_key(&blocks[2].recovered_block().hash()));
639 assert!(tree_state.parent_to_child.contains_key(&blocks[3].recovered_block().hash()));
640 assert!(!tree_state.parent_to_child.contains_key(&blocks[4].recovered_block().hash()));
641
642 assert_eq!(
643 tree_state.parent_to_child.get(&blocks[2].recovered_block().hash()),
644 Some(&HashSet::from_iter([blocks[3].recovered_block().hash()]))
645 );
646 assert_eq!(
647 tree_state.parent_to_child.get(&blocks[3].recovered_block().hash()),
648 Some(&HashSet::from_iter([blocks[4].recovered_block().hash()]))
649 );
650 }
651}