use super::{
ic_snapshot_data::{DataExtent, IcSnapshotDataReply},
ic_snapshot_metadata::IcSnapshotMetadataReply,
};
use std::fmt;
use thiserror::Error;
pub struct IcSnapshotDataCoverage<'metadata> {
metadata: &'metadata IcSnapshotMetadataReply<'metadata>,
regions: [u64; 3],
chunks: Vec<bool>,
}
impl<'metadata> IcSnapshotDataCoverage<'metadata> {
#[must_use]
pub fn new(metadata: &'metadata IcSnapshotMetadataReply<'metadata>) -> Self {
Self {
metadata,
regions: [0; 3],
chunks: vec![false; metadata.metadata().wasm_chunk_store.len()],
}
}
pub fn admit(
&mut self,
reply: &IcSnapshotDataReply<'_, '_>,
) -> Result<(), IcSnapshotDataCoverageError> {
if reply.request().metadata().digest() != self.metadata.digest() {
return Err(IcSnapshotDataCoverageError::MetadataMismatch);
}
let (index, offset, size) = match DataExtent::from(reply.request().kind()) {
DataExtent::Region {
region,
offset,
size,
} => (region.index(), offset, size),
DataExtent::Chunk { hash } => {
let index = self
.metadata
.chunk_index(hash)
.ok_or(IcSnapshotDataCoverageError::MetadataMismatch)?;
if self.chunks[index] {
return Err(IcSnapshotDataCoverageError::DuplicateChunk);
}
self.chunks[index] = true;
return Ok(());
}
};
if offset != self.regions[index] {
return Err(IcSnapshotDataCoverageError::NoncontiguousRange);
}
let end = offset
.checked_add(size)
.ok_or(IcSnapshotDataCoverageError::NoncontiguousRange)?;
self.regions[index] = end;
Ok(())
}
#[must_use]
pub const fn metadata(&self) -> &'metadata IcSnapshotMetadataReply<'metadata> {
self.metadata
}
#[must_use]
pub const fn covered_region_bytes(&self) -> [u64; 3] {
self.regions
}
#[must_use]
pub fn covered_chunks(&self) -> usize {
self.chunks.iter().filter(|covered| **covered).count()
}
#[must_use]
pub fn complete(&self) -> Option<IcSnapshotDataCoverageView<'_, 'metadata>> {
let sizes = self.metadata.region_sizes();
(self.regions == sizes && self.chunks.iter().all(|covered| *covered))
.then_some(IcSnapshotDataCoverageView { coverage: self })
}
}
impl fmt::Debug for IcSnapshotDataCoverage<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IcSnapshotDataCoverage")
.field("covered_region_bytes", &self.regions)
.field("covered_chunks", &self.covered_chunks())
.finish_non_exhaustive()
}
}
#[derive(Debug)]
pub struct IcSnapshotDataCoverageView<'coverage, 'metadata> {
coverage: &'coverage IcSnapshotDataCoverage<'metadata>,
}
impl<'metadata> IcSnapshotDataCoverageView<'_, 'metadata> {
#[must_use]
pub const fn metadata(&self) -> &'metadata IcSnapshotMetadataReply<'metadata> {
self.coverage.metadata()
}
}
#[derive(Debug, Error, Eq, PartialEq)]
pub enum IcSnapshotDataCoverageError {
#[error("snapshot data coverage metadata mismatch")]
MetadataMismatch,
#[error("snapshot data coverage requires the next contiguous range")]
NoncontiguousRange,
#[error("snapshot data coverage already contains this chunk")]
DuplicateChunk,
}
#[cfg(test)]
mod tests;