use std::alloc::{alloc, dealloc, Layout};
use std::ptr::NonNull;
pub struct FrameArena {
base: NonNull<u8>,
head: usize,
capacity: usize,
}
impl FrameArena {
pub fn new(capacity: usize) -> Self {
let layout = Layout::from_size_align(capacity, 16).expect("Invalid arena layout");
let ptr = unsafe { alloc(layout) };
let base = NonNull::new(ptr).expect("Failed to allocate frame arena");
Self {
base,
head: 0,
capacity,
}
}
pub fn alloc<T>(&mut self) -> *mut T {
self.alloc_layout(Layout::new::<T>()) as *mut T
}
pub fn alloc_layout(&mut self, layout: Layout) -> *mut u8 {
let align = layout.align();
let size = layout.size();
let aligned_head = (self.head + align - 1) & !(align - 1);
if aligned_head + size > self.capacity {
return std::ptr::null_mut();
}
let ptr = unsafe { self.base.as_ptr().add(aligned_head) };
self.head = aligned_head + size;
ptr
}
pub fn alloc_slice<T>(&mut self, count: usize) -> *mut T {
let layout = Layout::array::<T>(count).expect("Invalid array layout");
self.alloc_layout(layout) as *mut T
}
pub fn head(&self) -> usize {
self.head
}
pub fn reset(&mut self) {
self.head = 0;
#[cfg(feature = "debug")]
unsafe {
std::ptr::write_bytes(self.base.as_ptr(), 0xCD, self.capacity);
}
}
pub fn reset_to(&mut self, head: usize) {
debug_assert!(head <= self.head, "Cannot reset forward");
self.head = head;
}
pub fn remaining(&self) -> usize {
self.capacity - self.head
}
pub fn capacity(&self) -> usize {
self.capacity
}
pub fn allocated(&self) -> usize {
self.head
}
}
impl Drop for FrameArena {
fn drop(&mut self) {
let layout = Layout::from_size_align(self.capacity, 16).expect("Invalid arena layout");
unsafe {
dealloc(self.base.as_ptr(), layout);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_basic_allocation() {
let mut arena = FrameArena::new(1024);
let ptr1 = arena.alloc::<u32>();
assert!(!ptr1.is_null());
let ptr2 = arena.alloc::<u64>();
assert!(!ptr2.is_null());
assert_ne!(ptr1 as *mut u8, ptr2 as *mut u8);
}
#[test]
fn test_reset() {
let mut arena = FrameArena::new(1024);
let ptr1 = arena.alloc::<u32>();
let _head_before = arena.head();
arena.reset();
assert_eq!(arena.head(), 0);
let ptr2 = arena.alloc::<u32>();
assert_eq!(ptr1, ptr2);
}
#[test]
fn test_exhaustion() {
let mut arena = FrameArena::new(32);
let _ = arena.alloc::<[u8; 16]>();
let _ = arena.alloc::<[u8; 16]>();
let ptr = arena.alloc::<[u8; 16]>();
assert!(ptr.is_null());
}
}