use crate::{Heightable, marshal::store::Blocks, types::Height};
use commonware_cryptography::Digestible;
use commonware_runtime::Handle;
use commonware_storage::archive::Identifier;
use commonware_utils::sync::Mutex;
use std::sync::Arc;
#[derive(Debug, PartialEq, Eq)]
pub enum Op {
Put(Height, usize),
Get(Option<Height>),
}
pub struct Recording<T: Blocks> {
inner: T,
ops: Arc<Mutex<Vec<Op>>>,
payload_address: fn(&T::Block) -> usize,
}
impl<T: Blocks> Recording<T> {
pub fn new(inner: T, payload_address: fn(&T::Block) -> usize) -> Self {
Self {
inner,
ops: Arc::new(Mutex::new(Vec::new())),
payload_address,
}
}
pub fn ops(&self) -> Arc<Mutex<Vec<Op>>> {
Arc::clone(&self.ops)
}
}
impl<T: Blocks> Blocks for Recording<T> {
type Block = T::Block;
type Error = T::Error;
async fn put(
mut self,
digest: <Self::Block as Digestible>::Digest,
block: &Self::Block,
) -> Result<Self, Self::Error> {
self.ops
.lock()
.push(Op::Put(block.height(), (self.payload_address)(block)));
self.inner = self.inner.put(digest, block).await?;
Ok(self)
}
async fn sync(mut self) -> Result<Self, Self::Error> {
self.inner = self.inner.sync().await?;
Ok(self)
}
async fn start_sync(mut self) -> Result<(Self, Handle<()>), Self::Error> {
let handle;
(self.inner, handle) = self.inner.start_sync().await?;
Ok((self, handle))
}
async fn get(
&self,
id: Identifier<'_, <Self::Block as Digestible>::Digest>,
) -> Result<Option<Self::Block>, Self::Error> {
let height = match &id {
Identifier::Index(index) => Some(Height::new(*index)),
Identifier::Key(_) => None,
};
self.ops.lock().push(Op::Get(height));
self.inner.get(id).await
}
async fn prune(mut self, min: Height) -> Result<Self, Self::Error> {
self.inner = self.inner.prune(min).await?;
Ok(self)
}
fn missing_items(&self, start: Height, max: usize) -> Vec<Height> {
self.inner.missing_items(start, max)
}
fn next_gap(&self, value: Height) -> (Option<Height>, Option<Height>) {
self.inner.next_gap(value)
}
fn last_index(&self) -> Option<Height> {
self.inner.last_index()
}
}