1pub struct Scratch<'a> {
2 buf: &'a mut [u8],
3 offset: usize,
4}
5
6impl<'a> Scratch<'a> {
7 pub fn new(buf: &'a mut [u8]) -> Self { Scratch { buf, offset: 0 } }
8
9 pub fn reset(&mut self) { self.offset = 0; }
10
11 pub fn remaining(&self) -> usize { self.buf.len().saturating_sub(self.offset) }
12
13 pub fn as_slice(&self) -> &[u8] { &self.buf[..self.offset] }
14
15 pub fn base_ptr(&self) -> *const u8 { self.buf.as_ptr() }
16
17 pub fn alloc_bytes(&mut self, len: usize) -> Option<&mut [u8]> {
18 self.alloc_align(len, core::mem::align_of::<u8>())
19 }
20
21 pub fn alloc_align(&mut self, len: usize, align: usize) -> Option<&mut [u8]> {
22 if len == 0 { return Some(&mut []); }
23 let align = align.max(1);
24 let base_ptr = self.buf.as_ptr() as usize;
25 let mut start = base_ptr.checked_add(self.offset)?;
26 let rem = start % align;
27 if rem != 0 { start = start.checked_add(align - rem)?; }
28 let rel = start - base_ptr;
29 if rel.saturating_add(len) > self.buf.len() { return None; }
30 self.offset = rel.saturating_add(len);
31 let s = &mut self.buf[rel..rel+len];
32 Some(s)
33 }
34
35 pub fn alloc_zeroed_f32_slice(&mut self, count: usize) -> Option<&mut [f32]> {
36 let bytes = self.alloc_align(
37 count.saturating_mul(core::mem::size_of::<f32>()),
38 core::mem::align_of::<f32>(),
39 )?;
40 let ptr = bytes.as_mut_ptr() as *mut f32;
41 let len = count;
42 unsafe { for i in 0..len { ptr.add(i).write(0f32); } }
43 Some(unsafe { core::slice::from_raw_parts_mut(ptr, len) })
44 }
45}