use std::collections::HashMap;
pub type Epochs = Vec<(u64, u64)>;
pub type AlgorithmMetas = Vec<(u64, Epochs)>;
pub trait Storage {
type Error: std::error::Error + Send + Sync + 'static;
fn store_leaf(&mut self, index: u64, data: &[u8]) -> Result<(), Self::Error>;
fn get_leaf(&self, index: u64) -> Result<Vec<u8>, Self::Error>;
fn len(&self) -> u64;
fn is_empty(&self) -> bool {
self.len() == 0
}
fn store_node(
&mut self,
alg_id: u64,
left: u64,
height: usize,
hash: &[u8],
) -> Result<(), Self::Error>;
fn get_node(
&self,
alg_id: u64,
left: u64,
height: usize,
) -> Result<Option<Vec<u8>>, Self::Error>;
fn store_algorithm_meta(
&mut self,
alg_id: u64,
epochs: &[(u64, u64)],
) -> Result<(), Self::Error>;
fn load_algorithm_metas(&self) -> Result<AlgorithmMetas, Self::Error>;
}
#[derive(Debug, Default)]
pub struct MemoryStorage {
leaves: Vec<Vec<u8>>,
nodes: HashMap<(u64, u64, usize), Vec<u8>>,
algorithm_metas: HashMap<u64, Vec<(u64, u64)>>,
}
impl MemoryStorage {
#[must_use]
pub fn new() -> Self {
Self {
leaves: Vec::new(),
nodes: HashMap::new(),
algorithm_metas: HashMap::new(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemoryStorageError {
pub index: u64,
pub stored: u64,
}
impl std::fmt::Display for MemoryStorageError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"leaf index {} not found (storage contains {} leaves)",
self.index, self.stored
)
}
}
impl std::error::Error for MemoryStorageError {}
impl Storage for MemoryStorage {
type Error = MemoryStorageError;
fn store_leaf(&mut self, index: u64, data: &[u8]) -> Result<(), Self::Error> {
debug_assert_eq!(
index,
self.leaves.len() as u64,
"store_leaf called out of order"
);
self.leaves.push(data.to_vec());
Ok(())
}
fn get_leaf(&self, index: u64) -> Result<Vec<u8>, Self::Error> {
self.leaves
.get(index as usize)
.cloned()
.ok_or(MemoryStorageError {
index,
stored: self.leaves.len() as u64,
})
}
fn len(&self) -> u64 {
self.leaves.len() as u64
}
fn store_node(
&mut self,
alg_id: u64,
left: u64,
height: usize,
hash: &[u8],
) -> Result<(), Self::Error> {
self.nodes.insert((alg_id, left, height), hash.to_vec());
Ok(())
}
fn get_node(
&self,
alg_id: u64,
left: u64,
height: usize,
) -> Result<Option<Vec<u8>>, Self::Error> {
Ok(self.nodes.get(&(alg_id, left, height)).cloned())
}
fn store_algorithm_meta(
&mut self,
alg_id: u64,
epochs: &[(u64, u64)],
) -> Result<(), Self::Error> {
self.algorithm_metas.insert(alg_id, epochs.to_vec());
Ok(())
}
fn load_algorithm_metas(&self) -> Result<AlgorithmMetas, Self::Error> {
Ok(self
.algorithm_metas
.iter()
.map(|(&id, e)| (id, e.clone()))
.collect())
}
}