pub mod sync;
use std::alloc::{alloc_zeroed, dealloc, Layout};
use std::cell::Cell;
use std::mem::{align_of, size_of};
use std::ptr::{self, NonNull};
use std::rc::Rc;
use std::sync::{Arc, Mutex, Weak};
use crate::error::{Error, Result};
use crate::format::PixelFormat;
pub(crate) const MAX_ALIGN: usize = 64;
pub(crate) fn buffer_layout(cap: usize) -> Option<Layout> {
if cap == 0 {
None
} else {
Layout::from_size_align(cap, MAX_ALIGN).ok()
}
}
pub(crate) struct Buffer {
pub(crate) ptr: NonNull<u8>,
pub(crate) cap: usize,
}
unsafe impl Send for Buffer {}
unsafe impl Sync for Buffer {}
impl Buffer {
pub(crate) fn new_zeroed(cap: usize) -> Self {
match buffer_layout(cap) {
None => {
Buffer {
ptr: unsafe { NonNull::new_unchecked(MAX_ALIGN as *mut u8) },
cap: 0,
}
}
Some(layout) => {
let raw = unsafe { alloc_zeroed(layout) };
let ptr =
NonNull::new(raw).unwrap_or_else(|| std::alloc::handle_alloc_error(layout));
Buffer { ptr, cap }
}
}
}
pub(crate) fn zero(&mut self) {
if self.cap > 0 {
unsafe { ptr::write_bytes(self.ptr.as_ptr(), 0, self.cap) };
}
}
}
impl Drop for Buffer {
fn drop(&mut self) {
if let Some(layout) = buffer_layout(self.cap) {
unsafe { dealloc(self.ptr.as_ptr(), layout) };
}
}
}
pub struct ArenaPool {
inner: Mutex<PoolInner>,
cap_per_arena: usize,
max_arenas: usize,
max_alloc_count_per_arena: u32,
}
struct PoolInner {
idle: Vec<Buffer>,
total_allocated: usize,
}
impl ArenaPool {
pub fn new(max_arenas: usize, cap_per_arena: usize) -> Arc<Self> {
Self::with_alloc_count_cap(max_arenas, cap_per_arena, 1_000_000)
}
pub fn with_alloc_count_cap(
max_arenas: usize,
cap_per_arena: usize,
max_alloc_count_per_arena: u32,
) -> Arc<Self> {
Arc::new(Self {
inner: Mutex::new(PoolInner {
idle: Vec::with_capacity(max_arenas),
total_allocated: 0,
}),
cap_per_arena,
max_arenas,
max_alloc_count_per_arena,
})
}
pub fn cap_per_arena(&self) -> usize {
self.cap_per_arena
}
pub fn max_arenas(&self) -> usize {
self.max_arenas
}
pub fn lease(self: &Arc<Self>) -> Result<Arena> {
let buffer = {
let mut inner = self.inner.lock().expect("ArenaPool mutex poisoned");
if let Some(buf) = inner.idle.pop() {
buf
} else if inner.total_allocated < self.max_arenas {
inner.total_allocated += 1;
Buffer::new_zeroed(self.cap_per_arena)
} else {
return Err(Error::resource_exhausted(format!(
"ArenaPool exhausted: all {} arenas checked out",
self.max_arenas
)));
}
};
let base = buffer.ptr;
Ok(Arena {
buffer: Cell::new(Some(buffer)),
base,
cursor: Cell::new(0),
alloc_count: Cell::new(0),
cap: self.cap_per_arena,
alloc_count_cap: self.max_alloc_count_per_arena,
pool: Arc::downgrade(self),
})
}
fn release(&self, mut buffer: Buffer) {
buffer.zero();
if let Ok(mut inner) = self.inner.lock() {
inner.idle.push(buffer);
}
}
}
pub struct Arena {
buffer: Cell<Option<Buffer>>,
base: NonNull<u8>,
cursor: Cell<usize>,
alloc_count: Cell<u32>,
cap: usize,
alloc_count_cap: u32,
pool: Weak<ArenaPool>,
}
impl Arena {
pub fn capacity(&self) -> usize {
self.cap
}
pub fn used(&self) -> usize {
self.cursor.get()
}
pub fn alloc_count(&self) -> u32 {
self.alloc_count.get()
}
pub fn alloc_count_exceeded(&self) -> bool {
self.alloc_count.get() >= self.alloc_count_cap
}
#[allow(clippy::mut_from_ref)] pub fn alloc<T>(&self, count: usize) -> Result<&mut [T]>
where
T: bytemuck::Zeroable,
{
const fn assert_align<T>() {
assert!(
align_of::<T>() <= MAX_ALIGN,
"Arena::alloc<T>: align_of::<T>() exceeds MAX_ALIGN; \
increase MAX_ALIGN in arena/mod.rs"
);
}
const { assert_align::<T>() };
let next_count =
self.alloc_count.get().checked_add(1).ok_or_else(|| {
Error::resource_exhausted("Arena alloc_count overflow".to_string())
})?;
if next_count > self.alloc_count_cap {
return Err(Error::resource_exhausted(format!(
"Arena alloc-count cap of {} exceeded",
self.alloc_count_cap
)));
}
let elem_size = size_of::<T>();
let elem_align = align_of::<T>();
let bytes = elem_size
.checked_mul(count)
.ok_or_else(|| Error::resource_exhausted("Arena alloc size overflow".to_string()))?;
let cursor = self.cursor.get();
let aligned = align_up(cursor, elem_align).ok_or_else(|| {
Error::resource_exhausted("Arena cursor alignment overflow".to_string())
})?;
let new_cursor = aligned.checked_add(bytes).ok_or_else(|| {
Error::resource_exhausted("Arena cursor advance overflow".to_string())
})?;
if new_cursor > self.cap {
return Err(Error::resource_exhausted(format!(
"Arena cap of {} bytes exceeded (would consume {} bytes)",
self.cap, new_cursor
)));
}
let slice: &mut [T] = unsafe {
let elem_ptr = self.base.as_ptr().add(aligned).cast::<T>();
std::slice::from_raw_parts_mut(elem_ptr, count)
};
self.cursor.set(new_cursor);
self.alloc_count.set(next_count);
Ok(slice)
}
pub fn reset(&mut self) {
self.cursor.set(0);
self.alloc_count.set(0);
}
}
impl Drop for Arena {
fn drop(&mut self) {
if let Some(buffer) = self.buffer.take() {
if let Some(pool) = self.pool.upgrade() {
pool.release(buffer);
} else {
drop(buffer);
}
}
}
}
fn align_up(n: usize, align: usize) -> Option<usize> {
debug_assert!(align.is_power_of_two(), "alignment must be a power of two");
let mask = align - 1;
n.checked_add(mask).map(|m| m & !mask)
}
#[non_exhaustive]
#[derive(Copy, Clone, Debug)]
pub struct FrameHeader {
pub width: u32,
pub height: u32,
pub pixel_format: PixelFormat,
pub presentation_timestamp: Option<i64>,
}
impl FrameHeader {
pub fn new(
width: u32,
height: u32,
pixel_format: PixelFormat,
presentation_timestamp: Option<i64>,
) -> Self {
Self {
width,
height,
pixel_format,
presentation_timestamp,
}
}
}
pub const MAX_PLANES: usize = 4;
pub struct FrameInner {
arena: Arena,
plane_offsets: [(usize, usize); MAX_PLANES],
plane_count: u8,
header: FrameHeader,
}
pub type Frame = Rc<FrameInner>;
impl FrameInner {
pub fn new(arena: Arena, planes: &[(usize, usize)], header: FrameHeader) -> Result<Frame> {
if planes.len() > MAX_PLANES {
return Err(Error::invalid(format!(
"FrameInner supports at most {} planes (got {})",
MAX_PLANES,
planes.len()
)));
}
let used = arena.used();
for (i, (off, len)) in planes.iter().enumerate() {
let end = off
.checked_add(*len)
.ok_or_else(|| Error::invalid(format!("plane {i}: offset+len overflow")))?;
if end > used {
return Err(Error::invalid(format!(
"plane {i}: range {off}..{end} exceeds arena used={used}"
)));
}
}
let mut plane_offsets = [(0usize, 0usize); MAX_PLANES];
for (i, p) in planes.iter().enumerate() {
plane_offsets[i] = *p;
}
Ok(Rc::new(FrameInner {
arena,
plane_offsets,
plane_count: planes.len() as u8,
header,
}))
}
pub fn plane_count(&self) -> usize {
self.plane_count as usize
}
pub fn plane(&self, i: usize) -> Option<&[u8]> {
if i >= self.plane_count as usize {
return None;
}
let (off, len) = self.plane_offsets[i];
let buf: &[u8] = unsafe {
let elem_ptr = self.arena.base.as_ptr().add(off);
std::slice::from_raw_parts(elem_ptr, len)
};
Some(buf)
}
pub fn header(&self) -> &FrameHeader {
&self.header
}
}
#[cfg(test)]
mod tests {
use super::*;
fn small_pool(slots: usize, cap: usize) -> Arc<ArenaPool> {
ArenaPool::new(slots, cap)
}
#[test]
fn pool_lease_returns_err_when_exhausted() {
let pool = small_pool(2, 1024);
let a = pool.lease().expect("first lease");
let b = pool.lease().expect("second lease");
let third = pool.lease();
assert!(matches!(third, Err(Error::ResourceExhausted(_))));
drop((a, b));
}
#[test]
fn arena_alloc_caps_at_size_limit() {
let pool = small_pool(1, 64);
let arena = pool.lease().unwrap();
let _: &mut [u8] = arena.alloc::<u8>(32).unwrap();
let _: &mut [u8] = arena.alloc::<u8>(32).unwrap();
let third = arena.alloc::<u8>(1);
assert!(matches!(third, Err(Error::ResourceExhausted(_))));
}
#[test]
fn arena_alloc_count_cap_fires() {
let pool = ArenaPool::with_alloc_count_cap(1, 1024, 3);
let arena = pool.lease().unwrap();
let _: &mut [u8] = arena.alloc::<u8>(1).unwrap();
let _: &mut [u8] = arena.alloc::<u8>(1).unwrap();
let _: &mut [u8] = arena.alloc::<u8>(1).unwrap();
assert!(arena.alloc_count_exceeded());
let fourth = arena.alloc::<u8>(1);
assert!(matches!(fourth, Err(Error::ResourceExhausted(_))));
}
#[test]
fn arena_returns_to_pool_on_drop() {
let pool = small_pool(1, 256);
{
let arena = pool.lease().expect("first lease");
assert!(matches!(pool.lease(), Err(Error::ResourceExhausted(_))));
drop(arena);
}
let _again = pool.lease().expect("re-lease after drop");
}
#[test]
fn arena_alignment_is_respected() {
let pool = small_pool(1, 64);
let arena = pool.lease().unwrap();
let _: &mut [u8] = arena.alloc::<u8>(1).unwrap();
let s: &mut [u32] = arena.alloc::<u32>(4).unwrap();
let addr = s.as_ptr() as usize;
assert_eq!(addr % align_of::<u32>(), 0);
assert_eq!(s.len(), 4);
}
#[cfg(miri)]
#[test]
fn arena_alloc_can_return_misaligned_typed_slice() {
let pool = small_pool(1, 0);
let arena = pool.lease().unwrap();
let _s: &mut [u32] = arena.alloc::<u32>(0).unwrap();
}
#[cfg(miri)]
#[test]
fn arena_alloc_allows_invalid_bit_patterns_for_copy_types() {
fn requires_zeroable<T: bytemuck::Zeroable>() {}
requires_zeroable::<u8>();
}
#[cfg(miri)]
#[test]
fn arena_alloc_second_slice_invalidates_first_mut_reference() {
let pool = small_pool(1, 2);
let arena = pool.lease().unwrap();
let first = arena.alloc::<u8>(1).unwrap();
let second = arena.alloc::<u8>(1).unwrap();
first[0] = 1;
second[0] = 2;
}
fn build_simple_frame(pool: &Arc<ArenaPool>) -> Frame {
let arena = pool.lease().unwrap();
let plane0: &mut [u8] = arena.alloc::<u8>(16).unwrap();
for (i, b) in plane0.iter_mut().enumerate() {
*b = i as u8;
}
let header = FrameHeader::new(4, 4, PixelFormat::Gray8, Some(42));
FrameInner::new(arena, &[(0, 16)], header).unwrap()
}
#[test]
fn frame_refcount_keeps_arena_alive() {
let pool = small_pool(1, 256);
let frame = build_simple_frame(&pool);
let clone = Rc::clone(&frame);
drop(frame);
let plane = clone.plane(0).expect("plane 0");
assert_eq!(plane.len(), 16);
for (i, b) in plane.iter().enumerate() {
assert_eq!(*b, i as u8);
}
assert_eq!(clone.header().width, 4);
assert_eq!(clone.header().height, 4);
assert_eq!(clone.header().presentation_timestamp, Some(42));
assert!(matches!(pool.lease(), Err(Error::ResourceExhausted(_))));
}
#[test]
fn last_drop_returns_arena_to_pool() {
let pool = small_pool(1, 256);
let frame = build_simple_frame(&pool);
let clone = Rc::clone(&frame);
drop(frame);
drop(clone);
let _again = pool.lease().expect("lease after last drop");
}
#[test]
fn frame_rejects_too_many_planes() {
let pool = small_pool(1, 256);
let arena = pool.lease().unwrap();
let header = FrameHeader::new(1, 1, PixelFormat::Gray8, None);
let too_many = vec![(0usize, 0usize); MAX_PLANES + 1];
let r = FrameInner::new(arena, &too_many, header);
assert!(matches!(r, Err(Error::InvalidData(_))));
}
#[test]
fn frame_rejects_plane_outside_arena() {
let pool = small_pool(1, 64);
let arena = pool.lease().unwrap();
let header = FrameHeader::new(1, 1, PixelFormat::Gray8, None);
let r = FrameInner::new(arena, &[(0, 16)], header);
assert!(matches!(r, Err(Error::InvalidData(_))));
}
#[test]
fn pool_outlives_buffer_drop_when_pool_dropped_first() {
let pool = small_pool(1, 64);
let arena = pool.lease().unwrap();
drop(pool);
drop(arena);
}
#[test]
fn arena_reset_clears_allocations() {
let pool = small_pool(1, 32);
let mut arena = pool.lease().unwrap();
let _: &mut [u8] = arena.alloc::<u8>(32).unwrap();
assert!(matches!(
arena.alloc::<u8>(1),
Err(Error::ResourceExhausted(_))
));
arena.reset();
let _: &mut [u8] = arena.alloc::<u8>(32).unwrap();
}
}