use super::*;
#[derive(Debug)]
pub(crate) struct CheckpointValues<T> {
values: Vec<T>,
_memory: Option<Box<dyn RuntimeMemoryPermit>>,
capacity: usize,
}
#[derive(Debug)]
pub(crate) struct CheckpointSharedValues<T>(Arc<CheckpointValues<T>>);
impl<T> Clone for CheckpointSharedValues<T> {
fn clone(&self) -> Self {
Self(Arc::clone(&self.0))
}
}
impl<T> std::ops::Deref for CheckpointSharedValues<T> {
type Target = [T];
fn deref(&self) -> &Self::Target {
&self.0.values
}
}
impl<T> CheckpointValues<T> {
pub(crate) fn as_slice(&self) -> &[T] {
&self.values
}
pub(crate) fn as_mut_slice(&mut self) -> &mut [T] {
&mut self.values
}
pub(crate) fn new(
capacity: usize,
work: &CheckpointWorkContext,
) -> Result<Self, CheckpointWorkError> {
let unit = work.start_unit()?;
let bytes = capacity
.checked_mul(std::mem::size_of::<T>())
.and_then(|bytes| bytes.checked_add(std::mem::size_of::<Self>()))
.and_then(|bytes| {
bytes.checked_add(2 * std::mem::size_of::<std::sync::atomic::AtomicUsize>())
})
.ok_or_else(|| {
work.record_failure(CheckpointWorkError::Allocation {
bytes: u64::MAX,
reason: "checkpoint value capacity overflows usize".into(),
})
})?;
let memory = work.reserve_memory(bytes)?;
let mut values = Vec::new();
values.try_reserve_exact(capacity).map_err(|error| {
work.record_failure(CheckpointWorkError::Allocation {
bytes: bytes as u64,
reason: error.to_string(),
})
})?;
if std::mem::size_of::<T>() != 0 && values.capacity() != capacity {
return Err(work.record_failure(CheckpointWorkError::Allocation {
bytes: bytes as u64,
reason: format!(
"allocator granted {} values beyond admitted capacity {capacity}",
values.capacity()
),
}));
}
let output = Self {
values,
_memory: memory,
capacity,
};
unit.finish();
work.checkpoint()?;
Ok(output)
}
pub(crate) fn push(
&mut self,
value: T,
work: &CheckpointWorkContext,
) -> Result<(), CheckpointWorkError> {
let unit = work.start_unit()?;
if self.values.len() == self.capacity {
return Err(work.record_failure(CheckpointWorkError::Allocation {
bytes: self
.values
.capacity()
.saturating_mul(std::mem::size_of::<T>()) as u64,
reason: "checkpoint value push would exceed admitted capacity".into(),
}));
}
self.values.push(value);
unit.finish();
work.checkpoint()
}
pub(crate) fn share(
self,
work: &CheckpointWorkContext,
) -> Result<CheckpointSharedValues<T>, CheckpointWorkError> {
let unit = work.start_unit()?;
let output = CheckpointSharedValues(Arc::new(self));
unit.finish();
work.checkpoint()?;
Ok(output)
}
}