pub struct Scratch<'a> {
buf: &'a mut [u8],
offset: usize,
}
impl<'a> Scratch<'a> {
pub fn new(buf: &'a mut [u8]) -> Self {
Scratch { buf, offset: 0 }
}
pub fn reset(&mut self) {
self.offset = 0;
}
pub fn remaining(&self) -> usize {
self.buf.len().saturating_sub(self.offset)
}
pub fn as_slice(&self) -> &[u8] {
&self.buf[..self.offset]
}
pub fn base_ptr(&self) -> *const u8 {
self.buf.as_ptr()
}
pub fn alloc_bytes(&mut self, len: usize) -> Option<&mut [u8]> {
self.alloc_align(len, core::mem::align_of::<u8>())
}
pub fn alloc_align(&mut self, len: usize, align: usize) -> Option<&mut [u8]> {
if len == 0 {
return Some(&mut []);
}
let align = align.max(1);
let base_ptr = self.buf.as_ptr() as usize;
let mut start = base_ptr.checked_add(self.offset)?;
let rem = start % align;
if rem != 0 {
start = start.checked_add(align - rem)?;
}
let rel = start - base_ptr;
if rel.saturating_add(len) > self.buf.len() {
return None;
}
self.offset = rel.saturating_add(len);
let s = &mut self.buf[rel..rel + len];
Some(s)
}
pub fn alloc_zeroed_f32_slice(&mut self, count: usize) -> Option<&mut [f32]> {
let bytes = self.alloc_align(
count.saturating_mul(core::mem::size_of::<f32>()),
core::mem::align_of::<f32>(),
)?;
let ptr = bytes.as_mut_ptr() as *mut f32;
let len = count;
unsafe {
for i in 0..len {
ptr.add(i).write(0f32);
}
}
Some(unsafe { core::slice::from_raw_parts_mut(ptr, len) })
}
pub fn alloc_zeroed_f64_slice(&mut self, count: usize) -> Option<&mut [f64]> {
let bytes = self.alloc_align(
count.saturating_mul(core::mem::size_of::<f64>()),
core::mem::align_of::<f64>(),
)?;
let ptr = bytes.as_mut_ptr() as *mut f64;
let len = count;
unsafe {
for i in 0..len {
ptr.add(i).write(0f64);
}
}
Some(unsafe { core::slice::from_raw_parts_mut(ptr, len) })
}
}