use core::cell::Cell;
use core::mem;
use core::ptr::NonNull;
use core::slice;
const BLOCK_ALIGN: usize = 64;
#[derive(Clone, Copy)]
#[repr(C, align(64))]
struct Block([u8; BLOCK_ALIGN]);
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub struct ArenaStats {
pub used: usize,
pub capacity: usize,
pub peak: usize,
}
impl ArenaStats {
#[must_use]
pub fn peak_fraction(&self) -> f32 {
if self.capacity == 0 {
0.0
} else {
self.peak as f32 / self.capacity as f32
}
}
}
pub struct FrameArena {
_storage: Vec<Block>,
base: NonNull<u8>,
capacity: usize,
offset: Cell<usize>,
peak: Cell<usize>,
}
#[allow(clippy::mut_from_ref)]
impl FrameArena {
#[must_use]
pub fn with_capacity(capacity: usize) -> Self {
let blocks = capacity.div_ceil(BLOCK_ALIGN);
let mut storage = vec![Block([0; BLOCK_ALIGN]); blocks];
let base =
NonNull::new(storage.as_mut_ptr().cast::<u8>()).expect("Vec::as_mut_ptr is never null");
Self {
_storage: storage,
base,
capacity: blocks * BLOCK_ALIGN,
offset: Cell::new(0),
peak: Cell::new(0),
}
}
#[inline]
#[must_use]
pub fn used(&self) -> usize {
self.offset.get()
}
#[inline]
#[must_use]
pub const fn capacity(&self) -> usize {
self.capacity
}
#[inline]
#[must_use]
pub fn remaining(&self) -> usize {
self.capacity - self.offset.get()
}
#[inline]
#[must_use]
pub fn stats(&self) -> ArenaStats {
ArenaStats {
used: self.offset.get(),
capacity: self.capacity,
peak: self.peak.get(),
}
}
#[inline]
pub fn reset(&mut self) {
self.offset.set(0);
}
#[inline]
pub fn reset_peak(&mut self) {
self.peak.set(self.offset.get());
}
#[inline]
pub fn alloc_slice<T: Copy>(&self, len: usize, value: T) -> &mut [T] {
let (capacity, free) = (self.capacity, self.remaining());
match self.try_alloc_slice(len, value) {
Some(slice) => slice,
None => exhausted::<T>(len, free, capacity),
}
}
pub fn try_alloc_slice<T: Copy>(&self, len: usize, value: T) -> Option<&mut [T]> {
let ptr = self.reserve::<T>(len)?;
unsafe {
for index in 0..len {
ptr.add(index).write(value);
}
Some(slice::from_raw_parts_mut(ptr, len))
}
}
#[inline]
pub fn alloc_slice_copy<T: Copy>(&self, src: &[T]) -> &mut [T] {
let (capacity, free) = (self.capacity, self.remaining());
match self.try_alloc_slice_copy(src) {
Some(slice) => slice,
None => exhausted::<T>(src.len(), free, capacity),
}
}
pub fn try_alloc_slice_copy<T: Copy>(&self, src: &[T]) -> Option<&mut [T]> {
let ptr = self.reserve::<T>(src.len())?;
unsafe {
ptr.copy_from_nonoverlapping(src.as_ptr(), src.len());
Some(slice::from_raw_parts_mut(ptr, src.len()))
}
}
#[inline]
pub fn alloc<T: Copy>(&self, value: T) -> &mut T {
&mut self.alloc_slice(1, value)[0]
}
fn reserve<T: Copy>(&self, len: usize) -> Option<*mut T> {
let align = mem::align_of::<T>();
assert!(
align <= BLOCK_ALIGN,
"FrameArena supports alignments up to {BLOCK_ALIGN}, {} needs {align}",
core::any::type_name::<T>(),
);
let size = mem::size_of::<T>().checked_mul(len)?;
let start = self.offset.get().next_multiple_of(align);
let end = start.checked_add(size)?;
if end > self.capacity {
return None;
}
self.offset.set(end);
self.peak.set(self.peak.get().max(end));
if size == 0 {
return Some(NonNull::<T>::dangling().as_ptr());
}
Some(unsafe { self.base.as_ptr().add(start).cast::<T>() })
}
}
#[cold]
#[inline(never)]
fn exhausted<T>(len: usize, free: usize, capacity: usize) -> ! {
panic!(
"FrameArena exhausted: wanted {} bytes for [{}; {len}], {free} of {capacity} free",
mem::size_of::<T>().saturating_mul(len),
core::any::type_name::<T>(),
)
}
impl core::fmt::Debug for FrameArena {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("FrameArena")
.field("used", &self.offset.get())
.field("capacity", &self.capacity)
.field("peak", &self.peak.get())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn capacity_rounds_up_to_a_whole_block() {
let arena = FrameArena::with_capacity(1);
assert_eq!(arena.capacity(), BLOCK_ALIGN);
let arena = FrameArena::with_capacity(BLOCK_ALIGN + 1);
assert_eq!(arena.capacity(), BLOCK_ALIGN * 2);
}
#[test]
fn allocations_are_aligned_for_their_type() {
#[derive(Clone, Copy, Debug, PartialEq)]
#[repr(align(32))]
struct Wide(u64);
let arena = FrameArena::with_capacity(4096);
let bytes = arena.alloc_slice(3, 0u8);
assert_eq!(bytes.len(), 3);
let wide = arena.alloc_slice(2, Wide(7));
assert_eq!(wide.as_ptr() as usize % 32, 0);
assert_eq!(wide, [Wide(7), Wide(7)]);
let floats = arena.alloc_slice(5, 1.5f64);
assert_eq!(floats.as_ptr() as usize % mem::align_of::<f64>(), 0);
assert!(floats.iter().all(|f| *f == 1.5));
}
#[test]
fn allocations_do_not_overlap() {
let arena = FrameArena::with_capacity(4096);
let first = arena.alloc_slice(8, 1u32).as_ptr() as usize;
let second = arena.alloc_slice(8, 2u32);
assert!(second.iter().all(|v| *v == 2));
assert!(second.as_ptr() as usize >= first + 8 * 4);
}
#[test]
fn many_slices_stay_live_and_independent() {
let mut arena = FrameArena::with_capacity(4096);
let instances = arena.alloc_slice(4, [0.0f32; 4]);
let draw_args = arena.alloc_slice(3, 0u32);
let indices = arena.alloc_slice_copy(&[10u16, 11, 12, 13, 14]);
let flag = arena.alloc(false);
let wide = arena.alloc_slice(2, 0u64);
for round in 0..8u32 {
for (index, instance) in instances.iter_mut().enumerate() {
instance[index % 4] = (round * 10 + index as u32) as f32;
}
for (index, arg) in draw_args.iter_mut().enumerate() {
*arg = round * 100 + index as u32;
}
for (index, wide) in wide.iter_mut().enumerate() {
*wide = u64::from(round) << 32 | index as u64;
}
*flag = round.is_multiple_of(2);
}
assert_eq!(instances[0][0], 70.0);
assert_eq!(instances[3][3], 73.0);
assert_eq!(draw_args, [700, 701, 702]);
assert_eq!(
indices,
[10, 11, 12, 13, 14],
"untouched slice must be intact"
);
assert_eq!(wide, [7u64 << 32, 7u64 << 32 | 1]);
assert!(!*flag);
let spans = [
(instances.as_ptr() as usize, size_of_val(&*instances)),
(draw_args.as_ptr() as usize, size_of_val(&*draw_args)),
(indices.as_ptr() as usize, size_of_val(&*indices)),
(wide.as_ptr() as usize, size_of_val(&*wide)),
(std::ptr::from_ref(&*flag) as usize, size_of_val(&*flag)),
];
for (i, (start, len)) in spans.iter().enumerate() {
for (other_start, other_len) in spans.iter().skip(i + 1) {
assert!(
start + len <= *other_start || other_start + other_len <= *start,
"allocations overlap: {start}+{len} vs {other_start}+{other_len}"
);
}
}
arena.reset();
assert_eq!(arena.used(), 0);
}
#[test]
fn allocation_works_through_a_shared_reference() {
let arena = FrameArena::with_capacity(256);
fn fill(arena: &FrameArena, value: u8) -> &mut [u8] {
arena.alloc_slice(8, value)
}
let a = fill(&arena, 1);
let b = fill(&arena, 2);
a[0] = 9;
assert_eq!(a, [9, 1, 1, 1, 1, 1, 1, 1]);
assert_eq!(b, [2; 8]);
assert_eq!(arena.used(), 16);
}
#[test]
fn slices_can_be_copied_in() {
let arena = FrameArena::with_capacity(256);
let src = [1.0f32, 2.0, 3.0, 4.0];
let copied = arena.alloc_slice_copy(&src);
assert_eq!(copied, src);
copied[0] = 9.0;
assert_eq!(src[0], 1.0, "the arena must hold a copy, not a view");
assert_eq!(arena.used(), 16);
assert!(arena.try_alloc_slice_copy(&[0u8; 512]).is_none());
}
#[test]
fn an_arena_allocation_hands_back_storage_the_caller_can_write_through() {
let arena = FrameArena::with_capacity(256);
let value = arena.alloc([1u16, 2, 3]);
value[1] = 9;
assert_eq!(*value, [1, 9, 3]);
}
#[test]
fn zero_sized_and_empty_allocations_cost_nothing() {
let arena = FrameArena::with_capacity(128);
let empty = arena.alloc_slice(0, 1u64);
assert!(empty.is_empty());
let units = arena.alloc_slice(1000, ());
assert_eq!(units.len(), 1000);
assert_eq!(arena.used(), 0);
}
#[test]
fn reset_reuses_the_block_and_keeps_the_peak() {
let mut arena = FrameArena::with_capacity(1024);
let first = arena.alloc_slice(64, 0u8).as_ptr() as usize;
assert_eq!(arena.used(), 64);
arena.reset();
assert_eq!(arena.used(), 0);
assert_eq!(arena.stats().peak, 64);
let second = arena.alloc_slice(16, 0u8).as_ptr() as usize;
assert_eq!(second, first, "reset must hand back the same memory");
assert_eq!(
arena.stats(),
ArenaStats {
used: 16,
capacity: 1024,
peak: 64
}
);
arena.reset_peak();
assert_eq!(arena.stats().peak, 16);
}
#[test]
fn peak_tracks_the_worst_frame() {
let mut arena = FrameArena::with_capacity(1024);
for len in [10usize, 700, 3, 200] {
arena.reset();
let _ = arena.alloc_slice(len, 0u8);
}
assert_eq!(arena.stats().peak, 700);
assert!((arena.stats().peak_fraction() - 700.0 / 1024.0).abs() < 1e-6);
}
#[test]
fn exhaustion_is_reported_not_grown_into() {
let arena = FrameArena::with_capacity(128);
assert!(arena.try_alloc_slice(100, 0u8).is_some());
assert_eq!(arena.remaining(), 28);
assert!(arena.try_alloc_slice(29, 0u8).is_none());
assert_eq!(arena.capacity(), 128, "capacity must not have grown");
assert_eq!(arena.used(), 100);
assert!(arena.try_alloc_slice(28, 0u8).is_some());
}
#[test]
#[should_panic(expected = "FrameArena exhausted")]
fn alloc_panics_when_exhausted() {
let arena = FrameArena::with_capacity(64);
let _ = arena.alloc_slice(65, 0u8);
}
#[test]
fn a_fresh_arena_reports_default_stats_and_a_zero_peak_fraction() {
let arena = FrameArena::with_capacity(0);
assert_eq!(arena.stats(), ArenaStats::default());
assert_eq!(arena.stats().peak_fraction(), 0.0);
assert!(format!("{arena:?}").contains("FrameArena"));
}
}