use crate::{AppendableChain, BlockBuffer, BlockIndices};
use alloy_primitives::{BlockHash, BlockNumber};
use reth_primitives::{Receipt, SealedBlock, SealedBlockWithSenders};
use std::collections::{BTreeMap, HashMap};
#[derive(Debug)]
pub(crate) struct TreeState {
block_chain_id_generator: u64,
pub(crate) chains: HashMap<SidechainId, AppendableChain>,
pub(crate) block_indices: BlockIndices,
pub(crate) buffered_blocks: BlockBuffer,
}
impl TreeState {
pub(crate) fn new(
last_finalized_block_number: BlockNumber,
last_canonical_hashes: impl IntoIterator<Item = (BlockNumber, BlockHash)>,
buffer_limit: u32,
) -> Self {
Self {
block_chain_id_generator: 0,
chains: Default::default(),
block_indices: BlockIndices::new(
last_finalized_block_number,
BTreeMap::from_iter(last_canonical_hashes),
),
buffered_blocks: BlockBuffer::new(buffer_limit),
}
}
#[inline]
fn next_id(&mut self) -> SidechainId {
let id = self.block_chain_id_generator;
self.block_chain_id_generator += 1;
SidechainId(id)
}
#[inline]
pub(crate) const fn block_indices(&self) -> &BlockIndices {
&self.block_indices
}
#[inline]
pub(crate) fn block_by_hash(&self, block_hash: BlockHash) -> Option<&SealedBlock> {
self.block_with_senders_by_hash(block_hash).map(|block| &block.block)
}
#[inline]
pub(crate) fn block_with_senders_by_hash(
&self,
block_hash: BlockHash,
) -> Option<&SealedBlockWithSenders> {
let id = self.block_indices.get_side_chain_id(&block_hash)?;
let chain = self.chains.get(&id)?;
chain.block_with_senders(block_hash)
}
pub(crate) fn receipts_by_block_hash(&self, block_hash: BlockHash) -> Option<Vec<&Receipt>> {
let id = self.block_indices.get_side_chain_id(&block_hash)?;
let chain = self.chains.get(&id)?;
chain.receipts_by_block_hash(block_hash)
}
pub(crate) fn insert_chain(&mut self, chain: AppendableChain) -> Option<SidechainId> {
if chain.is_empty() {
return None
}
let chain_id = self.next_id();
self.block_indices.insert_chain(chain_id, &chain);
self.chains.insert(chain_id, chain);
Some(chain_id)
}
pub(crate) fn get_buffered_block(&self, hash: &BlockHash) -> Option<&SealedBlockWithSenders> {
self.buffered_blocks.block(hash)
}
pub(crate) fn lowest_buffered_ancestor(
&self,
hash: &BlockHash,
) -> Option<&SealedBlockWithSenders> {
self.buffered_blocks.lowest_ancestor(hash)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
pub(crate) struct SidechainId(u64);
impl From<SidechainId> for u64 {
fn from(value: SidechainId) -> Self {
value.0
}
}
#[cfg(test)]
impl From<u64> for SidechainId {
fn from(value: u64) -> Self {
Self(value)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::canonical_chain::CanonicalChain;
use alloy_primitives::B256;
use reth_execution_types::Chain;
use reth_provider::ExecutionOutcome;
#[test]
fn test_tree_state_initialization() {
let last_finalized_block_number = 10u64;
let last_canonical_hashes = vec![(9u64, B256::random()), (10u64, B256::random())];
let buffer_limit = 5;
let tree_state = TreeState::new(
last_finalized_block_number,
last_canonical_hashes.clone(),
buffer_limit,
);
assert_eq!(tree_state.block_chain_id_generator, 0);
assert_eq!(tree_state.block_indices().last_finalized_block(), last_finalized_block_number);
assert_eq!(
*tree_state.block_indices.canonical_chain().inner(),
*CanonicalChain::new(last_canonical_hashes.into_iter().collect()).inner()
);
assert!(tree_state.chains.is_empty());
assert!(tree_state.buffered_blocks.lru.is_empty());
}
#[test]
fn test_tree_state_next_id() {
let mut tree_state = TreeState::new(0, vec![], 5);
let first_id = tree_state.next_id();
let second_id = tree_state.next_id();
assert_eq!(first_id, SidechainId(0));
assert_eq!(second_id, SidechainId(1));
assert_eq!(tree_state.block_chain_id_generator, 2);
}
#[test]
fn test_tree_state_insert_chain() {
let mut tree_state = TreeState::new(0, vec![], 5);
let block: SealedBlockWithSenders = Default::default();
let block1_hash = B256::random();
let block2_hash = B256::random();
let mut block1 = block.clone();
let mut block2 = block;
block1.block.header.set_hash(block1_hash);
block1.block.header.set_block_number(9);
block2.block.header.set_hash(block2_hash);
block2.block.header.set_block_number(10);
let chain = AppendableChain::new(Chain::new(
[block1, block2],
Default::default(),
Default::default(),
));
let chain_id = tree_state.insert_chain(chain).unwrap();
assert_eq!(chain_id, SidechainId(0));
assert!(tree_state.chains.contains_key(&chain_id));
assert_eq!(
tree_state.block_indices.get_side_chain_id(&block1_hash).unwrap(),
SidechainId(0)
);
assert_eq!(
tree_state.block_indices.get_side_chain_id(&block2_hash).unwrap(),
SidechainId(0)
);
assert_eq!(tree_state.block_chain_id_generator, 1);
let chain_empty = AppendableChain::new(Chain::default());
let chain_id = tree_state.insert_chain(chain_empty);
assert!(chain_id.is_none());
assert!(tree_state.chains.contains_key(&SidechainId(0)));
assert!(!tree_state.chains.contains_key(&SidechainId(1)));
assert_eq!(
tree_state.block_indices.get_side_chain_id(&block1_hash).unwrap(),
SidechainId(0)
);
assert_eq!(
tree_state.block_indices.get_side_chain_id(&block2_hash).unwrap(),
SidechainId(0)
);
assert_eq!(tree_state.block_chain_id_generator, 1);
}
#[test]
fn test_block_by_hash_side_chain() {
let mut tree_state = TreeState::new(0, vec![], 5);
let block1_hash = B256::random();
let block2_hash = B256::random();
let mut block1: SealedBlockWithSenders = Default::default();
let mut block2: SealedBlockWithSenders = Default::default();
block1.block.header.set_hash(block1_hash);
block1.block.header.set_block_number(9);
block2.block.header.set_hash(block2_hash);
block2.block.header.set_block_number(10);
let chain = AppendableChain::new(Chain::new(
vec![block1.clone(), block2.clone()],
Default::default(),
Default::default(),
));
tree_state.insert_chain(chain).unwrap();
let retrieved_block1 = tree_state.block_by_hash(block1_hash);
assert_eq!(*retrieved_block1.unwrap(), block1.block);
let retrieved_block2 = tree_state.block_by_hash(block2_hash);
assert_eq!(*retrieved_block2.unwrap(), block2.block);
let non_existent_hash = B256::random();
let result = tree_state.block_by_hash(non_existent_hash);
assert!(result.is_none());
}
#[test]
fn test_block_with_senders_by_hash() {
let mut tree_state = TreeState::new(0, vec![], 5);
let block1_hash = B256::random();
let block2_hash = B256::random();
let mut block1: SealedBlockWithSenders = Default::default();
let mut block2: SealedBlockWithSenders = Default::default();
block1.block.header.set_hash(block1_hash);
block1.block.header.set_block_number(9);
block2.block.header.set_hash(block2_hash);
block2.block.header.set_block_number(10);
let chain = AppendableChain::new(Chain::new(
vec![block1.clone(), block2.clone()],
Default::default(),
Default::default(),
));
tree_state.insert_chain(chain).unwrap();
let retrieved_block1 = tree_state.block_with_senders_by_hash(block1_hash);
assert_eq!(*retrieved_block1.unwrap(), block1);
let retrieved_block2 = tree_state.block_with_senders_by_hash(block2_hash);
assert_eq!(*retrieved_block2.unwrap(), block2);
let non_existent_hash = B256::random();
let result = tree_state.block_with_senders_by_hash(non_existent_hash);
assert!(result.is_none());
}
#[test]
fn test_get_buffered_block() {
let mut tree_state = TreeState::new(0, vec![], 5);
let block_hash = B256::random();
let mut block: SealedBlockWithSenders = Default::default();
block.block.header.set_hash(block_hash);
tree_state.buffered_blocks.insert_block(block.clone());
let retrieved_block = tree_state.get_buffered_block(&block_hash);
assert_eq!(*retrieved_block.unwrap(), block);
let non_existent_hash = B256::random();
let result = tree_state.get_buffered_block(&non_existent_hash);
assert!(result.is_none());
}
#[test]
fn test_lowest_buffered_ancestor() {
let mut tree_state = TreeState::new(0, vec![], 5);
let ancestor_hash = B256::random();
let descendant_hash = B256::random();
let mut ancestor_block: SealedBlockWithSenders = Default::default();
let mut descendant_block: SealedBlockWithSenders = Default::default();
ancestor_block.block.header.set_hash(ancestor_hash);
descendant_block.block.header.set_hash(descendant_hash);
descendant_block.block.header.set_parent_hash(ancestor_hash);
tree_state.buffered_blocks.insert_block(ancestor_block.clone());
tree_state.buffered_blocks.insert_block(descendant_block.clone());
let lowest_ancestor = tree_state.lowest_buffered_ancestor(&descendant_hash);
assert!(lowest_ancestor.is_some());
assert_eq!(lowest_ancestor.unwrap().block.hash(), ancestor_hash);
let non_existent_hash = B256::random();
let result = tree_state.lowest_buffered_ancestor(&non_existent_hash);
assert!(result.is_none());
}
#[test]
fn test_receipts_by_block_hash() {
let mut tree_state = TreeState::new(0, vec![], 5);
let block_hash = B256::random();
let receipt1 = Receipt::default();
let receipt2 = Receipt::default();
let mut block: SealedBlockWithSenders = Default::default();
block.block.header.set_hash(block_hash);
let receipts = vec![receipt1, receipt2];
let chain = AppendableChain::new(Chain::new(
vec![block.clone()],
ExecutionOutcome { receipts: receipts.clone().into(), ..Default::default() },
Default::default(),
));
tree_state.insert_chain(chain).unwrap();
let retrieved_receipts = tree_state.receipts_by_block_hash(block_hash);
assert!(retrieved_receipts.is_some());
let receipts_ref: Vec<&Receipt> = receipts.iter().collect();
assert_eq!(retrieved_receipts.unwrap(), receipts_ref);
let non_existent_hash = B256::random();
let result = tree_state.receipts_by_block_hash(non_existent_hash);
assert!(result.is_none());
}
}