use core::cell::Cell;
use crate::error::ErrorCode;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StackBuffer<const N: usize> {
data: [u8; N],
len: usize,
}
impl<const N: usize> StackBuffer<N> {
pub const fn new() -> Self {
Self {
data: [0u8; N],
len: 0,
}
}
pub fn from_slice(slice: &[u8]) -> Option<Self> {
if slice.len() > N {
return None;
}
let mut buffer = Self::new();
buffer.data[..slice.len()].copy_from_slice(slice);
buffer.len = slice.len();
Some(buffer)
}
pub fn as_slice(&self) -> &[u8] {
&self.data[..self.len]
}
pub fn as_mut_slice(&mut self) -> &mut [u8] {
&mut self.data[..self.len]
}
pub fn as_str(&self) -> Result<&str, core::str::Utf8Error> {
core::str::from_utf8(self.as_slice())
}
pub const fn len(&self) -> usize {
self.len
}
pub const fn is_empty(&self) -> bool {
self.len == 0
}
pub const fn capacity(&self) -> usize {
N
}
pub const fn remaining(&self) -> usize {
N.saturating_sub(self.len)
}
pub fn clear(&mut self) {
self.len = 0;
}
pub fn push(&mut self, byte: u8) -> bool {
if self.len >= N {
return false;
}
self.data[self.len] = byte;
self.len += 1;
true
}
pub fn extend_from_slice(&mut self, slice: &[u8]) -> bool {
if self.len + slice.len() > N {
return false;
}
self.data[self.len..self.len + slice.len()].copy_from_slice(slice);
self.len += slice.len();
true
}
}
impl<const N: usize> Default for StackBuffer<N> {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Default)]
pub struct MemoryTracker {
current: Cell<usize>,
peak: Cell<usize>,
limit: usize,
}
impl MemoryTracker {
pub const fn new() -> Self {
Self {
current: Cell::new(0),
peak: Cell::new(0),
limit: 0,
}
}
pub const fn with_limit(limit: usize) -> Self {
Self {
current: Cell::new(0),
peak: Cell::new(0),
limit,
}
}
pub fn allocate(&self, bytes: usize) -> Result<(), &'static str> {
let new_current = self.current.get().saturating_add(bytes);
if self.limit > 0 && new_current > self.limit {
return Err("Memory limit exceeded");
}
self.current.set(new_current);
if new_current > self.peak.get() {
self.peak.set(new_current);
}
Ok(())
}
pub fn deallocate(&self, bytes: usize) {
let current = self.current.get();
self.current.set(current.saturating_sub(bytes));
}
pub fn current(&self) -> usize {
self.current.get()
}
pub fn peak(&self) -> usize {
self.peak.get()
}
pub fn limit(&self) -> usize {
self.limit
}
pub fn reset(&self) {
self.current.set(0);
self.peak.set(0);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EmbeddedLimits {
pub max_depth: usize,
pub max_token_length: usize,
pub max_container_items: usize,
}
impl EmbeddedLimits {
pub const CONSERVATIVE: Self = Self {
max_depth: 16,
max_token_length: 256,
max_container_items: 64,
};
pub const MINIMAL: Self = Self {
max_depth: 8,
max_token_length: 64,
max_container_items: 16,
};
#[inline]
pub fn check_depth(&self, current_depth: usize) -> Result<(), ErrorCode> {
if current_depth > self.max_depth {
Err(ErrorCode::SyntaxError)
} else {
Ok(())
}
}
}
impl Default for EmbeddedLimits {
fn default() -> Self {
Self::CONSERVATIVE
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct CompactError {
pub code: ErrorCode,
pub offset: u32,
}
impl CompactError {
pub const fn new(code: ErrorCode, offset: u32) -> Self {
Self { code, offset }
}
}
impl core::fmt::Display for CompactError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "{:?} at byte offset {}", self.code, self.offset)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_stack_buffer() {
let mut buf = StackBuffer::<16>::new();
assert!(buf.is_empty());
assert_eq!(buf.capacity(), 16);
assert_eq!(buf.remaining(), 16);
assert!(buf.push(b'H'));
assert!(buf.push(b'i'));
assert_eq!(buf.len(), 2);
assert_eq!(buf.as_slice(), b"Hi");
assert_eq!(buf.as_str().unwrap(), "Hi");
assert!(buf.extend_from_slice(b" embedded"));
assert_eq!(buf.as_str().unwrap(), "Hi embedded");
assert!(!buf.extend_from_slice(b" this will overflow the buffer capacity"));
buf.clear();
assert!(buf.is_empty());
}
#[test]
fn test_memory_tracker() {
let tracker = MemoryTracker::with_limit(100);
assert_eq!(tracker.current(), 0);
assert_eq!(tracker.limit(), 100);
assert!(tracker.allocate(60).is_ok());
assert_eq!(tracker.current(), 60);
assert_eq!(tracker.peak(), 60);
assert!(tracker.allocate(50).is_err()); tracker.deallocate(20);
assert_eq!(tracker.current(), 40);
assert_eq!(tracker.peak(), 60); }
#[test]
fn test_limits() {
let limits = EmbeddedLimits::CONSERVATIVE;
assert!(limits.check_depth(10).is_ok());
assert!(limits.check_depth(17).is_err());
}
#[test]
fn test_compact_error() {
let err = CompactError::new(ErrorCode::UnexpectedEof, 42);
assert_eq!(err.code, ErrorCode::UnexpectedEof);
assert_eq!(err.offset, 42);
}
}