use alloc::{collections::BTreeMap, vec::Vec};
use super::super::storage::StorageRestoreError;
use super::{
BoundParticipantCursor, ClaimFrontiers, ClosureDebt, CursorProgressFact, CursorProgressKey,
FrontierBinding, NonzeroDebtCursorEpisode,
};
use crate::{
algebra::WideResourceVector,
wire::{BindingEpoch, ConnectionIncarnation, Generation},
};
struct Reader<'a> {
bytes: &'a [u8],
}
impl Reader<'_> {
fn array<const N: usize>(&mut self) -> Result<[u8; N], StorageRestoreError> {
let bytes = self
.bytes
.get(..N)
.ok_or(StorageRestoreError::CursorEpisode)?;
let array = bytes
.try_into()
.map_err(|_| StorageRestoreError::CursorEpisode)?;
self.bytes = self
.bytes
.get(N..)
.ok_or(StorageRestoreError::CursorEpisode)?;
Ok(array)
}
fn u64(&mut self) -> Result<u64, StorageRestoreError> {
Ok(u64::from_be_bytes(self.array()?))
}
fn u128(&mut self) -> Result<u128, StorageRestoreError> {
Ok(u128::from_be_bytes(self.array()?))
}
fn binding(&mut self, tag: u8) -> Result<(BindingEpoch, FrontierBinding), StorageRestoreError> {
let server_incarnation = self.u64()?;
let connection_ordinal = self.u64()?;
let capability_generation =
Generation::new(self.u64()?).ok_or(StorageRestoreError::CursorEpisode)?;
let epoch = BindingEpoch {
connection_incarnation: ConnectionIncarnation {
server_incarnation,
connection_ordinal,
},
capability_generation,
};
let binding = match tag {
0 => FrontierBinding::Bound(epoch),
1 => FrontierBinding::Detached(epoch),
_ => return Err(StorageRestoreError::CursorEpisode),
};
Ok((epoch, binding))
}
fn count(&mut self, width: usize) -> Result<usize, StorageRestoreError> {
let count = usize::try_from(u32::from_be_bytes(self.array()?))
.map_err(|_| StorageRestoreError::CursorEpisode)?;
if width == 0 || count > self.bytes.len() / width {
return Err(StorageRestoreError::CursorEpisode);
}
Ok(count)
}
fn progress_facts(
&mut self,
) -> Result<Vec<(CursorProgressKey, CursorProgressFact)>, StorageRestoreError> {
let count = self.count(17)?;
let mut facts = Vec::new();
let mut previous = None;
for _ in 0..count {
let key = CursorProgressKey {
participant_index: self.u64()?,
boundary: self.u64()?,
};
if previous.is_some_and(|last| last >= key) {
return Err(StorageRestoreError::CursorEpisode);
}
previous = Some(key);
let fact = match self.array::<1>()?[0] {
0 => CursorProgressFact::Pending,
1 => CursorProgressFact::Consumed,
_ => return Err(StorageRestoreError::CursorEpisode),
};
facts.push((key, fact));
}
Ok(facts)
}
}
impl NonzeroDebtCursorEpisode {
fn install_checkpoint_bindings(
&mut self,
bindings: BTreeMap<u64, FrontierBinding>,
) -> Result<(), StorageRestoreError> {
for (participant_id, binding) in bindings {
let participant = self
.participants
.get_mut(&participant_id)
.ok_or(StorageRestoreError::CursorEpisode)?;
participant.binding = binding;
}
Ok(())
}
pub(in crate::lifecycle) fn decode_checkpoint(
bytes: &[u8],
frontiers: &ClaimFrontiers,
debt: ClosureDebt,
expected_observer_progress: u64,
) -> Result<Self, StorageRestoreError> {
let mut reader = Reader { bytes };
let conversation_id = reader.u64()?;
let saved_debt = WideResourceVector {
entries: reader.u128()?,
bytes: reader.u128()?,
};
let observer_progress = reader.u64()?;
let candidate_high_watermark = reader.u64()?;
let current_floor = reader.u128()?;
let cap_floor = reader.u128()?;
if conversation_id != frontiers.conversation_id()
|| saved_debt != debt.value()
|| observer_progress != expected_observer_progress
|| candidate_high_watermark != frontiers.sequence().ledger().high_watermark()
|| current_floor != frontiers.retained_floor()
{
return Err(StorageRestoreError::CursorEpisode);
}
let count = reader.count(41)?;
let expected_participants = frontiers.active_identities().participants();
if count != expected_participants.len() {
return Err(StorageRestoreError::CursorEpisode);
}
let mut participants = Vec::new();
let mut bindings = BTreeMap::new();
for expected in expected_participants {
let tag = reader.array::<1>()?[0];
let participant_id = reader.u64()?;
let (epoch, binding) = reader.binding(tag)?;
let cursor = reader.u64()?;
if participant_id != expected.participant_index()
|| binding != expected.binding()
|| cursor != expected.cursor()
{
return Err(StorageRestoreError::CursorEpisode);
}
participants.push(BoundParticipantCursor::new(participant_id, epoch, cursor));
if bindings.insert(participant_id, binding).is_some() {
return Err(StorageRestoreError::CursorEpisode);
}
}
let facts = reader.progress_facts()?;
if !reader.bytes.is_empty() {
return Err(StorageRestoreError::CursorEpisode);
}
let mut episode = Self::restore(
conversation_id,
debt,
observer_progress,
candidate_high_watermark,
current_floor,
cap_floor,
participants,
facts,
)
.ok_or(StorageRestoreError::CursorEpisode)?;
episode.install_checkpoint_bindings(bindings)?;
if episode
.encode()
.map_err(|_| StorageRestoreError::CursorEpisode)?
!= bytes
{
return Err(StorageRestoreError::CursorEpisode);
}
Ok(episode)
}
}