use alloc::vec::Vec;
use core::error::Error;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AllocationContext {
ProgramCodeLine,
ProgramPayload,
ProgramRuleTable,
CanonicalSource,
RuntimeInputValidation,
RuntimeOnceRuleState,
RuntimeRewriteState,
PayloadView,
RuntimeStateView,
FinalOutput,
ReturnOutput,
TraceSnapshot,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AllocationError {
context: AllocationContext,
kind: AllocationErrorKind,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AllocationErrorKind {
CapacityOverflow,
ReservationFailed {
requested_capacity: RequestedCapacity,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct RequestedCapacity {
value: usize,
}
impl RequestedCapacity {
pub(crate) const fn new(value: usize) -> Self {
Self { value }
}
#[must_use]
pub const fn get(self) -> usize {
self.value
}
}
impl core::fmt::Display for RequestedCapacity {
fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
self.value.fmt(formatter)
}
}
impl AllocationError {
pub(crate) const fn capacity_overflow(context: AllocationContext) -> Self {
Self {
context,
kind: AllocationErrorKind::CapacityOverflow,
}
}
pub(crate) const fn reservation_failed(
context: AllocationContext,
requested_capacity: RequestedCapacity,
) -> Self {
Self {
context,
kind: AllocationErrorKind::ReservationFailed { requested_capacity },
}
}
#[must_use]
pub const fn context(&self) -> AllocationContext {
self.context
}
#[must_use]
pub const fn kind(&self) -> AllocationErrorKind {
self.kind
}
}
impl Error for AllocationError {}
pub(crate) fn try_reserve_total_exact<T>(
vec: &mut Vec<T>,
total_capacity: RequestedCapacity,
context: AllocationContext,
) -> Result<(), AllocationError> {
if vec.capacity() >= total_capacity.get() {
return Ok(());
}
let additional = total_capacity
.get()
.checked_sub(vec.len())
.ok_or_else(|| AllocationError::capacity_overflow(context))?;
vec.try_reserve_exact(additional)
.map_err(|_| AllocationError::reservation_failed(context, total_capacity))
}
pub(crate) fn try_push<T>(
vec: &mut Vec<T>,
value: T,
context: AllocationContext,
) -> Result<(), AllocationError> {
if vec.len() == vec.capacity() {
let minimum_capacity = vec
.len()
.checked_add(1)
.ok_or_else(|| AllocationError::capacity_overflow(context))?;
let doubled_capacity = vec.capacity().checked_mul(2).unwrap_or(minimum_capacity);
let requested_capacity = RequestedCapacity::new(core::cmp::max(
minimum_capacity,
core::cmp::max(4, doubled_capacity),
));
try_reserve_total_exact(vec, requested_capacity, context)?;
}
vec.push(value);
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_support::{TestResult, ensure_eq};
use alloc::string::ToString;
#[test]
fn allocation_contexts_are_publicly_inspectable() -> TestResult {
let error = AllocationError::reservation_failed(
AllocationContext::TraceSnapshot,
RequestedCapacity::new(123),
);
ensure_eq!(error.context(), AllocationContext::TraceSnapshot)?;
ensure_eq!(
error.kind(),
AllocationErrorKind::ReservationFailed {
requested_capacity: RequestedCapacity::new(123),
},
)?;
let error = AllocationError::capacity_overflow(AllocationContext::CanonicalSource);
ensure_eq!(error.context(), AllocationContext::CanonicalSource)?;
ensure_eq!(error.kind(), AllocationErrorKind::CapacityOverflow)?;
let error = AllocationError::reservation_failed(
AllocationContext::RuntimeInputValidation,
RequestedCapacity::new(4),
);
ensure_eq!(error.context(), AllocationContext::RuntimeInputValidation)?;
Ok(())
}
#[test]
fn allocation_display_names_the_failed_context_and_capacity() -> TestResult {
let error = AllocationError::reservation_failed(
AllocationContext::TraceSnapshot,
RequestedCapacity::new(123),
);
ensure_eq!(
error.to_string(),
"allocation reservation failure while building trace snapshot; requested capacity: 123"
)?;
let error = AllocationError::reservation_failed(
AllocationContext::RuntimeStateView,
RequestedCapacity::new(456),
);
ensure_eq!(
error.to_string(),
"allocation reservation failure while building runtime state view; requested capacity: 456",
)?;
let error = AllocationError::reservation_failed(
AllocationContext::RuntimeInputValidation,
RequestedCapacity::new(789),
);
ensure_eq!(
error.to_string(),
"allocation reservation failure while building runtime input validation; requested capacity: 789",
)?;
let error = AllocationError::capacity_overflow(AllocationContext::CanonicalSource);
ensure_eq!(
error.to_string(),
"allocation capacity overflow while building canonical source bytes",
)
}
}