Zebra/zebra-state/src/in_memory/block_index.rs

68 lines
1.7 KiB
Rust

use std::{
collections::{btree_map::Entry, BTreeMap, HashMap},
error::Error,
sync::Arc,
};
use zebra_chain::{
block::{Block, BlockHeaderHash},
types::BlockHeight,
};
#[derive(Default)]
pub(super) struct BlockIndex {
by_hash: HashMap<BlockHeaderHash, Arc<Block>>,
by_height: BTreeMap<BlockHeight, Arc<Block>>,
}
impl BlockIndex {
pub(super) fn insert(
&mut self,
block: impl Into<Arc<Block>>,
) -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
let block = block.into();
let hash = block.as_ref().into();
let height = block.coinbase_height().unwrap();
match self.by_height.entry(height) {
Entry::Vacant(entry) => {
let _ = entry.insert(block.clone());
let _ = self.by_hash.insert(hash, block);
Ok(())
}
Entry::Occupied(_) => Err("forks in the chain aren't supported yet")?,
}
}
pub(super) fn get(&mut self, query: impl Into<BlockQuery>) -> Option<Arc<Block>> {
match query.into() {
BlockQuery::ByHash(hash) => self.by_hash.get(&hash),
BlockQuery::ByHeight(height) => self.by_height.get(&height),
}
.cloned()
}
pub(super) fn get_tip(&self) -> Option<Arc<Block>> {
self.by_height
.iter()
.next_back()
.map(|(_key, value)| value)
.cloned()
}
}
pub(super) enum BlockQuery {
ByHash(BlockHeaderHash),
ByHeight(BlockHeight),
}
impl From<BlockHeaderHash> for BlockQuery {
fn from(hash: BlockHeaderHash) -> Self {
Self::ByHash(hash)
}
}
impl From<BlockHeight> for BlockQuery {
fn from(height: BlockHeight) -> Self {
Self::ByHeight(height)
}
}