oxideav_core/arena/mod.rs
1//! Refcounted arena pool for decoder frame allocations.
2//!
3//! This module is the runtime half of the DoS-protection framework
4//! described in [`crate::limits`]. It provides three types:
5//!
6//! - [`ArenaPool`] — a pool of reusable raw byte buffers (allocated
7//! via [`std::alloc::alloc`] with a fixed [`MAX_ALIGN`] alignment so
8//! each buffer's base pointer is suitable for any `T` whose
9//! alignment is `<= MAX_ALIGN`) that a decoder leases from. Pool
10//! size and per-buffer capacity are fixed at construction; together
11//! they bound peak RSS by construction
12//! (`max_arenas × cap_per_arena`).
13//!
14//! - [`Arena`] — a single buffer leased from the pool. Allocations are
15//! bump-pointer (no per-alloc bookkeeping, no fragmentation). When
16//! the `Arena` is dropped, its buffer is returned to the pool, *not*
17//! freed — this is what makes the pool memory-reusing rather than
18//! memory-leaking. If the pool has been dropped before the arena
19//! (last-arena-outlives-pool), the arena's buffer is freed normally.
20//!
21//! - [`Frame`] / [`FrameInner`] — a refcounted (`Rc<FrameInner>`)
22//! handle that holds an `Arena` plus per-plane offset/length pairs
23//! and a small [`FrameHeader`]. As long as any clone of a `Frame`
24//! exists, its arena (and therefore its buffer) stays out of the
25//! pool. The last `Drop` returns the buffer.
26//!
27//! ## Design choices for round 1
28//!
29//! - **Hand-rolled bump allocator** over a raw `NonNull<u8>` from
30//! [`std::alloc::alloc`]. We deliberately do not depend on the
31//! `bumpalo` crate yet — the logic is twenty lines and avoids
32//! pulling in a dependency before profiling justifies it. The
33//! signature is intentionally compatible with what a
34//! `bumpalo`-backed implementation would look like, so swapping
35//! later is a contained refactor.
36//!
37//! - **`Rc` for `Frame`, not `Arc`.** This module targets the
38//! single-threaded decode path (one decoder, one consumer thread).
39//! The bump-pointer cursor is `Cell<usize>` for the same reason
40//! (no atomics on the hot path). For the cross-thread decode path
41//! — where a decoder produces frames on one thread and a consumer
42//! reads them on another — see the sibling [`sync`] module, which
43//! mirrors this API 1:1 with `Arc<FrameInner>` / atomic cursor so
44//! `Frame: Send + Sync`.
45//!
46//! - **`Arena::alloc<T>` returns `&mut [T]` borrowed from the arena.**
47//! The borrow is bounded by the lifetime of the `&Arena` reference,
48//! not the lifetime of the arena itself; the arena holds the
49//! buffer's base address as a raw [`NonNull<u8>`] so multiple
50//! calls to `alloc` against the same `&Arena` can each carve out
51//! non-overlapping sub-slices without ever materialising a
52//! whole-buffer mutable borrow (which would invalidate previously
53//! returned slices under stacked borrows). This matches
54//! `bumpalo::Bump::alloc_slice_*` semantics.
55//!
56//! ## Soundness notes
57//!
58//! Three issues called out by an external Miri audit (PR #12, May
59//! 2026) shaped the current implementation; they are noted here so
60//! future refactors don't reintroduce them:
61//!
62//! 1. **Base-pointer alignment.** A `Box<[u8]>` is byte-aligned only,
63//! so even an empty `&mut [u32]` carved out of one would have an
64//! unaligned pointer (UB). Each pool buffer is now allocated
65//! directly via [`std::alloc::alloc`] with [`MAX_ALIGN`] (= 64 B,
66//! enough for AVX-512), so the base pointer is suitable for any
67//! type the arena will hand out. `alloc::<T>` rejects types whose
68//! alignment exceeds [`MAX_ALIGN`] at compile time via a
69//! `const`-evaluated assertion.
70//!
71//! 2. **Invalid bit patterns.** Pool buffers are zero-filled, but
72//! zero is not a valid bit pattern for every `Copy` type
73//! (`NonZeroU8`, references, function pointers, niche-optimised
74//! enums, …). `alloc<T>` is therefore bounded on
75//! `bytemuck::Zeroable` rather than just `Copy`, so the safe API
76//! cannot hand out `&mut [NonZeroU8]` over zero bytes.
77//!
78//! 3. **Stacked-borrows retag.** Each `alloc` previously took
79//! `[u8]::as_mut_ptr` of the whole backing slice, which retagged
80//! the whole buffer and popped the borrow stacks of every
81//! previously returned `&mut [T]`. The fix is the raw `NonNull<u8>`
82//! base pointer above: each `alloc` does
83//! `base.as_ptr().add(offset).cast::<T>()` and never re-borrows
84//! the whole buffer.
85
86pub mod sync;
87
88use std::alloc::{alloc_zeroed, dealloc, Layout};
89use std::cell::Cell;
90use std::mem::{align_of, size_of};
91use std::ptr::{self, NonNull};
92use std::rc::Rc;
93use std::sync::{Arc, Mutex, Weak};
94
95use crate::error::{Error, Result};
96use crate::format::PixelFormat;
97
98/// Alignment used for every pool buffer's base pointer. 64 bytes
99/// covers the alignment requirements of every primitive type and of
100/// AVX-512 SIMD loads (`__m512` is 64-byte aligned). [`Arena::alloc`]
101/// statically rejects any type with a stricter alignment requirement.
102pub(crate) const MAX_ALIGN: usize = 64;
103
104/// Layout used to allocate (and deallocate) pool buffers. `cap` is the
105/// per-arena byte capacity; alignment is fixed at [`MAX_ALIGN`].
106///
107/// Returns `None` for `cap == 0` — `Layout::from_size_align` rejects
108/// zero-sized layouts and we can't pass a zero-sized layout to
109/// `std::alloc::alloc`. Callers must special-case the empty arena.
110pub(crate) fn buffer_layout(cap: usize) -> Option<Layout> {
111 if cap == 0 {
112 None
113 } else {
114 Layout::from_size_align(cap, MAX_ALIGN).ok()
115 }
116}
117
118/// Backing storage for one pool buffer — a raw aligned byte buffer
119/// produced by [`std::alloc::alloc_zeroed`] (or a sentinel for the
120/// `cap == 0` case, which doesn't allocate). Owns the allocation;
121/// frees it in `Drop`. Used by both [`crate::arena::ArenaPool`] and
122/// [`crate::arena::sync::ArenaPool`].
123pub(crate) struct Buffer {
124 /// Base pointer. For `cap > 0` this points at a live allocation
125 /// of `cap` bytes aligned to [`MAX_ALIGN`]. For `cap == 0` this is
126 /// a [`MAX_ALIGN`]-aligned dangling pointer (no backing storage).
127 pub(crate) ptr: NonNull<u8>,
128 /// Capacity of the allocation in bytes (also the layout `size`).
129 pub(crate) cap: usize,
130}
131
132// SAFETY: `Buffer` owns its allocation outright (no aliasing) and
133// `NonNull<u8>` is `!Send + !Sync` only out of caution; sending the
134// owning handle to another thread is sound.
135unsafe impl Send for Buffer {}
136unsafe impl Sync for Buffer {}
137
138impl Buffer {
139 /// Allocate a buffer of `cap` bytes aligned to [`MAX_ALIGN`],
140 /// zero-filled. For `cap == 0` returns a dangling-but-aligned
141 /// sentinel (matching `NonNull::dangling()` semantics for an
142 /// arbitrary-alignment pointer) without touching the global
143 /// allocator.
144 pub(crate) fn new_zeroed(cap: usize) -> Self {
145 match buffer_layout(cap) {
146 None => {
147 // Produce a `MAX_ALIGN`-aligned dangling pointer that
148 // is never dereferenced (cap == 0 means no allocation
149 // accesses go through it). We synthesise it from the
150 // integer constant rather than using a strict-
151 // provenance helper because `ptr::without_provenance_mut`
152 // is only stable since Rust 1.84 and the crate's
153 // declared MSRV is 1.80. Miri under default
154 // (permissive) provenance accepts this; for strict
155 // provenance Miri it will warn but not error, and we
156 // never actually load through this pointer.
157 Buffer {
158 // SAFETY: `MAX_ALIGN` is a non-zero usize, so
159 // casting it to `*mut u8` produces a non-null
160 // pointer that we will never dereference.
161 ptr: unsafe { NonNull::new_unchecked(MAX_ALIGN as *mut u8) },
162 cap: 0,
163 }
164 }
165 Some(layout) => {
166 // SAFETY: layout has non-zero size (we just checked).
167 let raw = unsafe { alloc_zeroed(layout) };
168 let ptr =
169 NonNull::new(raw).unwrap_or_else(|| std::alloc::handle_alloc_error(layout));
170 Buffer { ptr, cap }
171 }
172 }
173 }
174
175 /// Zero the entire buffer. Called when a buffer is returned to the
176 /// pool so a subsequent lease starts from a clean (and therefore
177 /// `Zeroable`-valid) state.
178 pub(crate) fn zero(&mut self) {
179 if self.cap > 0 {
180 // SAFETY: ptr points to `cap` bytes of writable storage we
181 // own exclusively (`&mut self`).
182 unsafe { ptr::write_bytes(self.ptr.as_ptr(), 0, self.cap) };
183 }
184 }
185}
186
187impl Drop for Buffer {
188 fn drop(&mut self) {
189 if let Some(layout) = buffer_layout(self.cap) {
190 // SAFETY: ptr was returned by `alloc_zeroed(layout)` and
191 // we have not freed it yet.
192 unsafe { dealloc(self.ptr.as_ptr(), layout) };
193 }
194 }
195}
196
197/// Pool of reusable byte buffers for arena-backed frame allocations.
198///
199/// Construct one per decoder via [`ArenaPool::new`]. Lease an
200/// [`Arena`] per frame via [`ArenaPool::lease`]; drop the arena (or
201/// drop the last clone of a [`Frame`] holding it) to return its
202/// buffer to the pool.
203///
204/// **Backpressure:** when all `max_arenas` slots are checked out the
205/// next [`ArenaPool::lease`] returns
206/// [`Error::ResourceExhausted`]. A decoder that hits this should
207/// surface the error to its caller rather than busy-loop — the
208/// upstream pipeline is supposed to drop frames it no longer needs,
209/// which returns a buffer to the pool.
210///
211/// `ArenaPool` is `Send + Sync` (the inner `Mutex<Vec<…>>` makes it
212/// safe to share across threads even though [`Arena`] / [`Frame`]
213/// themselves are `!Send` due to their `Rc`/`Cell` contents). This
214/// asymmetry is intentional: a parallel-decoder thread can share a
215/// single pool while each thread owns its own arenas — see also the
216/// sibling [`sync::ArenaPool`] whose leases are themselves `Send + Sync`.
217pub struct ArenaPool {
218 inner: Mutex<PoolInner>,
219 cap_per_arena: usize,
220 max_arenas: usize,
221 max_alloc_count_per_arena: u32,
222}
223
224struct PoolInner {
225 /// Buffers currently sitting idle in the pool (ready to lease).
226 idle: Vec<Buffer>,
227 /// Total buffers ever allocated by this pool (idle + in-flight).
228 /// Caps lazy growth at `max_arenas`.
229 total_allocated: usize,
230}
231
232impl ArenaPool {
233 /// Construct a new pool with `max_arenas` buffer slots, each of
234 /// `cap_per_arena` bytes. Buffers are allocated lazily on first
235 /// lease — a freshly constructed pool holds no memory.
236 ///
237 /// Per-arena allocation count is capped at `max_alloc_count` (use
238 /// [`ArenaPool::new`] which defaults to a generous 1M, or
239 /// [`ArenaPool::with_alloc_count_cap`] to tighten further).
240 pub fn new(max_arenas: usize, cap_per_arena: usize) -> Arc<Self> {
241 Self::with_alloc_count_cap(max_arenas, cap_per_arena, 1_000_000)
242 }
243
244 /// Like [`ArenaPool::new`] but lets the caller set the per-arena
245 /// allocation-count cap. Useful when the caller is plumbing
246 /// [`crate::DecoderLimits`] through.
247 pub fn with_alloc_count_cap(
248 max_arenas: usize,
249 cap_per_arena: usize,
250 max_alloc_count_per_arena: u32,
251 ) -> Arc<Self> {
252 Arc::new(Self {
253 inner: Mutex::new(PoolInner {
254 idle: Vec::with_capacity(max_arenas),
255 total_allocated: 0,
256 }),
257 cap_per_arena,
258 max_arenas,
259 max_alloc_count_per_arena,
260 })
261 }
262
263 /// Capacity of each arena buffer this pool hands out, in bytes.
264 pub fn cap_per_arena(&self) -> usize {
265 self.cap_per_arena
266 }
267
268 /// Maximum number of arenas that may be checked out at once.
269 pub fn max_arenas(&self) -> usize {
270 self.max_arenas
271 }
272
273 /// Lease one arena from the pool. Returns
274 /// [`Error::ResourceExhausted`] if every arena slot is already
275 /// checked out by an [`Arena`] (or a [`Frame`] holding one).
276 pub fn lease(self: &Arc<Self>) -> Result<Arena> {
277 let buffer = {
278 let mut inner = self.inner.lock().expect("ArenaPool mutex poisoned");
279 if let Some(buf) = inner.idle.pop() {
280 buf
281 } else if inner.total_allocated < self.max_arenas {
282 inner.total_allocated += 1;
283 Buffer::new_zeroed(self.cap_per_arena)
284 } else {
285 return Err(Error::resource_exhausted(format!(
286 "ArenaPool exhausted: all {} arenas checked out",
287 self.max_arenas
288 )));
289 }
290 };
291
292 let base = buffer.ptr;
293 Ok(Arena {
294 buffer: Cell::new(Some(buffer)),
295 base,
296 cursor: Cell::new(0),
297 alloc_count: Cell::new(0),
298 cap: self.cap_per_arena,
299 alloc_count_cap: self.max_alloc_count_per_arena,
300 pool: Arc::downgrade(self),
301 })
302 }
303
304 /// Return a buffer to the idle list. Called from `Arena::Drop`;
305 /// not part of the public API. The buffer is zeroed before being
306 /// returned so the next lease starts from a clean state — this is
307 /// what makes `Zeroable` a sufficient bound on `Arena::alloc<T>`
308 /// across pool reuse cycles.
309 fn release(&self, mut buffer: Buffer) {
310 buffer.zero();
311 if let Ok(mut inner) = self.inner.lock() {
312 inner.idle.push(buffer);
313 }
314 // If the lock is poisoned, drop the buffer normally — the
315 // pool is in an unusable state already.
316 }
317}
318
319/// One leased buffer from an [`ArenaPool`].
320///
321/// Allocations are bump-pointer: each call to [`Arena::alloc`] carves
322/// out a fresh aligned slice from the head of the buffer. There is no
323/// per-allocation header and no individual free — the entire arena
324/// is reset (returned to the pool) only when the `Arena` is dropped.
325///
326/// `Arena` is `!Send + !Sync` because its bump cursor is a `Cell` and
327/// its buffer cell is `Cell<Option<Buffer>>` (not synchronised). This
328/// is fine for the round-1 single-threaded decoder path. The sibling
329/// [`sync::Arena`] uses `AtomicUsize` for the cursor and a `Mutex`
330/// around the buffer slot to regain `Send + Sync`.
331pub struct Arena {
332 /// Backing buffer leased from the pool. `Cell<Option<Buffer>>` so
333 /// `Drop` can `take()` the buffer and hand it back to the pool
334 /// without needing `&mut self`. Outside of `Drop` this is always
335 /// `Some`.
336 ///
337 /// We never re-borrow this buffer mutably while handing out
338 /// slices from it — the typed pointers returned by `alloc` are
339 /// derived from the cached raw `base` pointer below, never from
340 /// `(*buffer).as_mut_ptr()`. This avoids the stacked-borrows
341 /// "whole-buffer retag invalidates previously returned slices"
342 /// problem.
343 buffer: Cell<Option<Buffer>>,
344 /// Cached base pointer of `buffer` (a [`MAX_ALIGN`]-aligned
345 /// allocation owned by `buffer`). Stable for the lifetime of the
346 /// arena: `Buffer` does not move its allocation, and we only take
347 /// `buffer` out of the cell during `Drop` after no allocator
348 /// activity remains. All `alloc` calls derive their typed
349 /// pointers from `base.as_ptr().add(offset)`.
350 base: NonNull<u8>,
351 /// Bump cursor: the next free byte offset within the buffer.
352 cursor: Cell<usize>,
353 /// Number of allocations performed so far.
354 alloc_count: Cell<u32>,
355 /// Cached cap (== `pool.cap_per_arena` at lease time).
356 cap: usize,
357 /// Cached cap (== `pool.max_alloc_count_per_arena` at lease time).
358 alloc_count_cap: u32,
359 /// Weak handle back to the pool so `Drop` can return the buffer.
360 pool: Weak<ArenaPool>,
361}
362
363impl Arena {
364 /// Capacity of this arena in bytes.
365 pub fn capacity(&self) -> usize {
366 self.cap
367 }
368
369 /// Bytes consumed by allocations so far.
370 pub fn used(&self) -> usize {
371 self.cursor.get()
372 }
373
374 /// Number of allocations performed so far.
375 pub fn alloc_count(&self) -> u32 {
376 self.alloc_count.get()
377 }
378
379 /// `true` once the per-arena allocation-count cap has been
380 /// reached. Decoders that produce many small allocations should
381 /// poll this and bail with [`Error::ResourceExhausted`] when it
382 /// flips, instead of waiting for the next [`Arena::alloc`] call
383 /// to fail.
384 pub fn alloc_count_exceeded(&self) -> bool {
385 self.alloc_count.get() >= self.alloc_count_cap
386 }
387
388 /// Allocate `count` `T`s out of this arena. Returns a borrowed
389 /// `&mut [T]` (lifetime bounded by the borrow of `self`).
390 ///
391 /// The returned slice points at zero-filled bytes (the pool
392 /// zero-fills on initial allocation and again whenever a buffer
393 /// is returned). The `Zeroable` bound on `T` guarantees that an
394 /// all-zero bit pattern is a valid value for `T`, so reading the
395 /// slice without first writing it is sound. **The intended
396 /// pattern is still "decoder fills the slice, then reads back
397 /// what it wrote" — but unwritten bytes will read back as
398 /// `T::zeroed()` rather than as UB.**
399 ///
400 /// Returns [`Error::ResourceExhausted`] if either the per-arena
401 /// byte cap or the per-arena allocation-count cap would be
402 /// exceeded.
403 ///
404 /// # Type bounds
405 ///
406 /// - `T: bytemuck::Zeroable` — pool buffers are zero-filled, so
407 /// handing back `&mut [T]` over those bytes is only sound when
408 /// the all-zero bit pattern is valid for `T`. This rules out
409 /// `NonZeroU8`/`NonZeroU16`/…/references/function pointers/
410 /// niche-optimised enums (anything where the optimizer relies
411 /// on a forbidden-bit-pattern invariant).
412 /// - `align_of::<T>() <= MAX_ALIGN` — checked at compile time via
413 /// a `const` assertion. The pool buffer's base pointer is
414 /// aligned to [`MAX_ALIGN`] (= 64 bytes); per-`T` alignment is
415 /// then a relative-offset adjustment of the bump cursor.
416 /// - The arena does not run destructors on allocated values, so
417 /// `T` should not have meaningful `Drop` glue. `Zeroable` is
418 /// automatically implemented only for types where this is the
419 /// case (primitives, `[T; N]` of zeroable, `#[derive(Zeroable)]`
420 /// on POD structs).
421 ///
422 /// **Aliasing model:** the bump cursor is monotonically
423 /// non-decreasing, so successive `alloc` calls return slices
424 /// covering disjoint regions of the underlying buffer. The
425 /// returned typed pointer is derived from the arena's cached raw
426 /// base pointer (`base.as_ptr().add(offset)`), never from a
427 /// re-borrow of the whole buffer — that's what keeps previously
428 /// returned `&mut [T]` slices valid under stacked borrows. This
429 /// is the standard arena-allocator pattern (cf.
430 /// `bumpalo::Bump::alloc_slice_*`) and is the reason this method
431 /// takes `&self` rather than `&mut self`.
432 #[allow(clippy::mut_from_ref)] // see "Aliasing model" doc above.
433 pub fn alloc<T>(&self, count: usize) -> Result<&mut [T]>
434 where
435 T: bytemuck::Zeroable,
436 {
437 // Compile-time check: T's alignment must not exceed the
438 // pool buffer's base alignment. Doing this as a const-eval'd
439 // assert means a violating monomorphisation fails the build.
440 const fn assert_align<T>() {
441 assert!(
442 align_of::<T>() <= MAX_ALIGN,
443 "Arena::alloc<T>: align_of::<T>() exceeds MAX_ALIGN; \
444 increase MAX_ALIGN in arena/mod.rs"
445 );
446 }
447 const { assert_align::<T>() };
448
449 // Allocation-count cap.
450 let next_count =
451 self.alloc_count.get().checked_add(1).ok_or_else(|| {
452 Error::resource_exhausted("Arena alloc_count overflow".to_string())
453 })?;
454 if next_count > self.alloc_count_cap {
455 return Err(Error::resource_exhausted(format!(
456 "Arena alloc-count cap of {} exceeded",
457 self.alloc_count_cap
458 )));
459 }
460
461 let elem_size = size_of::<T>();
462 let elem_align = align_of::<T>();
463 // Bytes requested.
464 let bytes = elem_size
465 .checked_mul(count)
466 .ok_or_else(|| Error::resource_exhausted("Arena alloc size overflow".to_string()))?;
467
468 // Align cursor up to T's alignment.
469 let cursor = self.cursor.get();
470 let aligned = align_up(cursor, elem_align).ok_or_else(|| {
471 Error::resource_exhausted("Arena cursor alignment overflow".to_string())
472 })?;
473 let new_cursor = aligned.checked_add(bytes).ok_or_else(|| {
474 Error::resource_exhausted("Arena cursor advance overflow".to_string())
475 })?;
476
477 if new_cursor > self.cap {
478 return Err(Error::resource_exhausted(format!(
479 "Arena cap of {} bytes exceeded (would consume {} bytes)",
480 self.cap, new_cursor
481 )));
482 }
483
484 // SAFETY:
485 //
486 // - `self.base` points to a `MAX_ALIGN`-aligned allocation of
487 // `self.cap` bytes owned by the `Buffer` inside `self.buffer`,
488 // which lives at least as long as `&self`.
489 // - `aligned + count*size_of::<T>() <= self.cap` (just checked
490 // above), so the byte range we slice is in-bounds.
491 // - `aligned` is a multiple of `align_of::<T>()` (computed via
492 // `align_up`), and `MAX_ALIGN >= align_of::<T>()` (compile-
493 // time assert above), so `base + aligned` is `T`-aligned.
494 // This holds even for `count == 0` (the slice still has an
495 // aligned dangling pointer, which is what an empty `&mut [T]`
496 // requires).
497 // - The cursor is monotonically non-decreasing, so the byte
498 // range `aligned..new_cursor` does not overlap any byte
499 // range previously returned by `alloc`. We never re-borrow
500 // the whole buffer — the typed pointer is derived from the
501 // raw base pointer — so the new `&mut [T]` does not invalidate
502 // any previously returned slice under stacked borrows.
503 // - `T: Zeroable` and the buffer bytes are zero, so the
504 // `&mut [T]` references valid `T` values (the safe API
505 // contract).
506 let slice: &mut [T] = unsafe {
507 let elem_ptr = self.base.as_ptr().add(aligned).cast::<T>();
508 std::slice::from_raw_parts_mut(elem_ptr, count)
509 };
510
511 self.cursor.set(new_cursor);
512 self.alloc_count.set(next_count);
513 Ok(slice)
514 }
515
516 /// Reset the arena to empty without releasing its buffer to the
517 /// pool. Useful for a decoder that wants to reuse the same arena
518 /// across several intermediate stages of the same frame. Callers
519 /// must ensure no slice previously returned from [`Arena::alloc`]
520 /// is still in use — Rust's borrow checker enforces this, since
521 /// `reset` takes `&mut self`.
522 pub fn reset(&mut self) {
523 self.cursor.set(0);
524 self.alloc_count.set(0);
525 }
526}
527
528impl Drop for Arena {
529 fn drop(&mut self) {
530 // Take the buffer out of the cell. We're in Drop with `&mut
531 // self`, so no `alloc`-returned slices can still be borrowing
532 // from `base`.
533 if let Some(buffer) = self.buffer.take() {
534 if let Some(pool) = self.pool.upgrade() {
535 pool.release(buffer);
536 } else {
537 // Pool was dropped before us — buffer drops here and
538 // its allocation is freed via `Buffer::Drop`.
539 drop(buffer);
540 }
541 }
542 }
543}
544
545/// Round `n` up to the next multiple of `align`. `align` must be a
546/// power of two. Returns `None` on overflow.
547fn align_up(n: usize, align: usize) -> Option<usize> {
548 debug_assert!(align.is_power_of_two(), "alignment must be a power of two");
549 let mask = align - 1;
550 n.checked_add(mask).map(|m| m & !mask)
551}
552
553/// Per-frame metadata carried alongside an [`Arena`] inside a
554/// [`Frame`]. Kept minimal in round 1; round 2 will extend with
555/// stride/colorspace/HDR fields as decoders need them.
556///
557/// `Copy` so it travels through the hot path with no allocation.
558#[non_exhaustive]
559#[derive(Copy, Clone, Debug)]
560pub struct FrameHeader {
561 pub width: u32,
562 pub height: u32,
563 pub pixel_format: PixelFormat,
564 /// Presentation timestamp in stream time-base units. `None` when
565 /// the codec did not surface one (e.g. a still image).
566 pub presentation_timestamp: Option<i64>,
567}
568
569impl FrameHeader {
570 /// Construct a header with all four mandatory fields set. Use
571 /// functional-update syntax (`FrameHeader { ..header }`) to add
572 /// future fields safely.
573 pub fn new(
574 width: u32,
575 height: u32,
576 pixel_format: PixelFormat,
577 presentation_timestamp: Option<i64>,
578 ) -> Self {
579 Self {
580 width,
581 height,
582 pixel_format,
583 presentation_timestamp,
584 }
585 }
586}
587
588/// Maximum number of planes a [`FrameInner`] can describe in round 1.
589/// Covers every real-world video pixel format (1 plane for packed
590/// RGB/YUV 4:2:2, 3 planes for I420/YV12/I444, 4 planes for YUVA / RGBA
591/// planar). Audio is handled by a separate sibling type in a future
592/// round; this module is video-only for now.
593pub const MAX_PLANES: usize = 4;
594
595/// The owned body of a refcounted [`Frame`].
596///
597/// Holds an [`Arena`] (the bytes), a fixed-size table of
598/// `(offset_in_arena, length_in_bytes)` pairs (one per plane), and a
599/// [`FrameHeader`]. The `plane_count` field tracks how many entries of
600/// `plane_offsets` are actually populated. Up to [`MAX_PLANES`] planes
601/// are supported.
602///
603/// **Lifetime:** an `Arena` returns its buffer to the pool when
604/// dropped. A `Rc<FrameInner>` keeps the arena alive via its single
605/// owned field, so as long as any clone of a [`Frame`] exists the
606/// underlying buffer stays out of the pool.
607pub struct FrameInner {
608 arena: Arena,
609 plane_offsets: [(usize, usize); MAX_PLANES],
610 plane_count: u8,
611 header: FrameHeader,
612}
613
614/// Refcounted handle to a decoded video frame. Construct via
615/// [`Frame::new`]; clone freely (each clone bumps the refcount by 1).
616/// The arena and its buffer are released back to the pool when the
617/// last clone is dropped.
618///
619/// `Frame` is `Rc<FrameInner>` (single-threaded decoder path). For the
620/// cross-thread decode path — where the consumer runs on a different
621/// thread from the decoder — use the sibling [`sync::Frame`] which is
622/// `Arc<sync::FrameInner>` and is `Send + Sync`.
623pub type Frame = Rc<FrameInner>;
624
625impl FrameInner {
626 /// Construct a `Frame` (refcounted `Rc<FrameInner>`) from an arena,
627 /// a slice of `(offset, length)` plane descriptors, and a header.
628 /// Returns [`Error::InvalidData`] if more than [`MAX_PLANES`]
629 /// planes are supplied or if any plane range falls outside the
630 /// arena's used region.
631 pub fn new(arena: Arena, planes: &[(usize, usize)], header: FrameHeader) -> Result<Frame> {
632 if planes.len() > MAX_PLANES {
633 return Err(Error::invalid(format!(
634 "FrameInner supports at most {} planes (got {})",
635 MAX_PLANES,
636 planes.len()
637 )));
638 }
639 let used = arena.used();
640 for (i, (off, len)) in planes.iter().enumerate() {
641 let end = off
642 .checked_add(*len)
643 .ok_or_else(|| Error::invalid(format!("plane {i}: offset+len overflow")))?;
644 if end > used {
645 return Err(Error::invalid(format!(
646 "plane {i}: range {off}..{end} exceeds arena used={used}"
647 )));
648 }
649 }
650 let mut plane_offsets = [(0usize, 0usize); MAX_PLANES];
651 for (i, p) in planes.iter().enumerate() {
652 plane_offsets[i] = *p;
653 }
654 Ok(Rc::new(FrameInner {
655 arena,
656 plane_offsets,
657 plane_count: planes.len() as u8,
658 header,
659 }))
660 }
661
662 /// Number of planes this frame holds.
663 pub fn plane_count(&self) -> usize {
664 self.plane_count as usize
665 }
666
667 /// Read-only access to plane `i`. Returns `None` if `i` is out of
668 /// range.
669 pub fn plane(&self, i: usize) -> Option<&[u8]> {
670 if i >= self.plane_count as usize {
671 return None;
672 }
673 let (off, len) = self.plane_offsets[i];
674 // SAFETY:
675 // - plane ranges were validated against `arena.used()` at
676 // construction (`off + len <= arena.cursor`), and the
677 // cursor is monotonically non-decreasing, so the byte
678 // range is still in-bounds.
679 // - The bytes were written by `alloc` and never moved (the
680 // buffer's allocation is stable for the arena's lifetime).
681 // - We derive the slice from the raw base pointer, never via
682 // a re-borrow of the whole buffer, so this `&[u8]` does not
683 // invalidate any other slice the caller is holding.
684 // - The borrow lifetime is bounded by `&self`.
685 let buf: &[u8] = unsafe {
686 let elem_ptr = self.arena.base.as_ptr().add(off);
687 std::slice::from_raw_parts(elem_ptr, len)
688 };
689 Some(buf)
690 }
691
692 /// Frame header (width / height / pixel format / pts).
693 pub fn header(&self) -> &FrameHeader {
694 &self.header
695 }
696}
697
698#[cfg(test)]
699mod tests {
700 use super::*;
701
702 fn small_pool(slots: usize, cap: usize) -> Arc<ArenaPool> {
703 ArenaPool::new(slots, cap)
704 }
705
706 #[test]
707 fn pool_lease_returns_err_when_exhausted() {
708 let pool = small_pool(2, 1024);
709 let a = pool.lease().expect("first lease");
710 let b = pool.lease().expect("second lease");
711 let third = pool.lease();
712 assert!(matches!(third, Err(Error::ResourceExhausted(_))));
713 // Keep a and b alive past the assertion so they aren't dropped
714 // before the failing lease.
715 drop((a, b));
716 }
717
718 #[test]
719 fn arena_alloc_caps_at_size_limit() {
720 let pool = small_pool(1, 64);
721 let arena = pool.lease().unwrap();
722 // 64 bytes capacity. Allocate 32 u8s — fits.
723 let _: &mut [u8] = arena.alloc::<u8>(32).unwrap();
724 // Allocate another 32 u8s — exactly fills.
725 let _: &mut [u8] = arena.alloc::<u8>(32).unwrap();
726 // Any further allocation fails.
727 let third = arena.alloc::<u8>(1);
728 assert!(matches!(third, Err(Error::ResourceExhausted(_))));
729 }
730
731 #[test]
732 fn arena_alloc_count_cap_fires() {
733 let pool = ArenaPool::with_alloc_count_cap(1, 1024, 3);
734 let arena = pool.lease().unwrap();
735 let _: &mut [u8] = arena.alloc::<u8>(1).unwrap();
736 let _: &mut [u8] = arena.alloc::<u8>(1).unwrap();
737 let _: &mut [u8] = arena.alloc::<u8>(1).unwrap();
738 assert!(arena.alloc_count_exceeded());
739 let fourth = arena.alloc::<u8>(1);
740 assert!(matches!(fourth, Err(Error::ResourceExhausted(_))));
741 }
742
743 #[test]
744 fn arena_returns_to_pool_on_drop() {
745 let pool = small_pool(1, 256);
746 {
747 let arena = pool.lease().expect("first lease");
748 // Sanity: arena is leased; further leases would fail.
749 assert!(matches!(pool.lease(), Err(Error::ResourceExhausted(_))));
750 drop(arena);
751 }
752 // Arena dropped — pool slot must be free again.
753 let _again = pool.lease().expect("re-lease after drop");
754 }
755
756 #[test]
757 fn arena_alignment_is_respected() {
758 let pool = small_pool(1, 64);
759 let arena = pool.lease().unwrap();
760 // Allocate a single u8 to misalign the cursor.
761 let _: &mut [u8] = arena.alloc::<u8>(1).unwrap();
762 // Now allocate u32s; expect cursor to be aligned to 4.
763 let s: &mut [u32] = arena.alloc::<u32>(4).unwrap();
764 let addr = s.as_ptr() as usize;
765 assert_eq!(addr % align_of::<u32>(), 0);
766 assert_eq!(s.len(), 4);
767 }
768
769 #[cfg(miri)]
770 #[test]
771 fn arena_alloc_can_return_misaligned_typed_slice() {
772 let pool = small_pool(1, 0);
773 let arena = pool.lease().unwrap();
774
775 // Memory-safety issue: the arena's backing allocation is a
776 // `Box<[u8]>`, so its base pointer is only guaranteed to be
777 // byte-aligned. `alloc::<T>` aligns only the byte offset, not
778 // the absolute address, and then constructs `&mut [T]`. The
779 // empty-buffer case makes this deterministic: even an empty
780 // `&mut [u32]` must have an aligned pointer, but `Box<[u8]>`
781 // uses an alignment-1 dangling pointer when its length is 0.
782 let _s: &mut [u32] = arena.alloc::<u32>(0).unwrap();
783 }
784
785 // Pre-fix this test was:
786 //
787 // let values = arena.alloc::<std::num::NonZeroU8>(1).unwrap();
788 // let _ = values[0].get();
789 //
790 // and failed under Miri because pool buffers are zero-filled and
791 // zero is not a valid `NonZeroU8`. Post-fix, the `Zeroable` bound
792 // on `Arena::alloc` makes that call a hard *compile* error — the
793 // strongest possible enforcement. The test below is a regression
794 // assertion that the bound stays as-or-stricter than `Zeroable`:
795 // if a future refactor weakened it back to `Copy`, the
796 // commented-out call site would start compiling again and Miri
797 // would once again accept the invalid bit pattern.
798 #[cfg(miri)]
799 #[test]
800 fn arena_alloc_allows_invalid_bit_patterns_for_copy_types() {
801 // `requires_zeroable::<NonZeroU8>()` would not compile —
802 // `NonZeroU8: !Zeroable`. Sanity-check the helper itself with
803 // a known zeroable type so the test is an actual exercise.
804 fn requires_zeroable<T: bytemuck::Zeroable>() {}
805 requires_zeroable::<u8>();
806 // Uncommenting the next line must fail to compile:
807 // requires_zeroable::<std::num::NonZeroU8>();
808 }
809
810 #[cfg(miri)]
811 #[test]
812 fn arena_alloc_second_slice_invalidates_first_mut_reference() {
813 let pool = small_pool(1, 2);
814 let arena = pool.lease().unwrap();
815
816 // Memory-safety issue: each `alloc` calls `[u8]::as_mut_ptr` on
817 // the whole backing slice before carving out the requested
818 // subslice. That materializes a new mutable borrow of the whole
819 // buffer and invalidates previously returned `&mut` slices, even
820 // when the byte ranges are disjoint.
821 let first = arena.alloc::<u8>(1).unwrap();
822 let second = arena.alloc::<u8>(1).unwrap();
823 first[0] = 1;
824 second[0] = 2;
825 }
826
827 fn build_simple_frame(pool: &Arc<ArenaPool>) -> Frame {
828 let arena = pool.lease().unwrap();
829 // Allocate 16 bytes for plane 0.
830 let plane0: &mut [u8] = arena.alloc::<u8>(16).unwrap();
831 for (i, b) in plane0.iter_mut().enumerate() {
832 *b = i as u8;
833 }
834 // The slice borrowed from arena ends here.
835 let header = FrameHeader::new(4, 4, PixelFormat::Gray8, Some(42));
836 FrameInner::new(arena, &[(0, 16)], header).unwrap()
837 }
838
839 #[test]
840 fn frame_refcount_keeps_arena_alive() {
841 let pool = small_pool(1, 256);
842 let frame = build_simple_frame(&pool);
843 let clone = Rc::clone(&frame);
844 drop(frame);
845 // Clone is still valid; arena still leased.
846 let plane = clone.plane(0).expect("plane 0");
847 assert_eq!(plane.len(), 16);
848 for (i, b) in plane.iter().enumerate() {
849 assert_eq!(*b, i as u8);
850 }
851 assert_eq!(clone.header().width, 4);
852 assert_eq!(clone.header().height, 4);
853 assert_eq!(clone.header().presentation_timestamp, Some(42));
854 // Pool still exhausted because clone holds the arena.
855 assert!(matches!(pool.lease(), Err(Error::ResourceExhausted(_))));
856 }
857
858 #[test]
859 fn last_drop_returns_arena_to_pool() {
860 let pool = small_pool(1, 256);
861 let frame = build_simple_frame(&pool);
862 let clone = Rc::clone(&frame);
863 drop(frame);
864 drop(clone);
865 // All clones gone — buffer must be back in the pool.
866 let _again = pool.lease().expect("lease after last drop");
867 }
868
869 #[test]
870 fn frame_rejects_too_many_planes() {
871 let pool = small_pool(1, 256);
872 let arena = pool.lease().unwrap();
873 let header = FrameHeader::new(1, 1, PixelFormat::Gray8, None);
874 let too_many = vec![(0usize, 0usize); MAX_PLANES + 1];
875 let r = FrameInner::new(arena, &too_many, header);
876 assert!(matches!(r, Err(Error::InvalidData(_))));
877 }
878
879 #[test]
880 fn frame_rejects_plane_outside_arena() {
881 let pool = small_pool(1, 64);
882 let arena = pool.lease().unwrap();
883 // arena.used() == 0; any non-empty plane is out of range.
884 let header = FrameHeader::new(1, 1, PixelFormat::Gray8, None);
885 let r = FrameInner::new(arena, &[(0, 16)], header);
886 assert!(matches!(r, Err(Error::InvalidData(_))));
887 }
888
889 #[test]
890 fn pool_outlives_buffer_drop_when_pool_dropped_first() {
891 // Exotic: arena outlives its pool. Buffer just frees normally.
892 let pool = small_pool(1, 64);
893 let arena = pool.lease().unwrap();
894 drop(pool);
895 // Drop arena — must not panic. The Weak handle won't upgrade.
896 drop(arena);
897 }
898
899 #[test]
900 fn arena_reset_clears_allocations() {
901 let pool = small_pool(1, 32);
902 let mut arena = pool.lease().unwrap();
903 let _: &mut [u8] = arena.alloc::<u8>(32).unwrap();
904 // Cap reached.
905 assert!(matches!(
906 arena.alloc::<u8>(1),
907 Err(Error::ResourceExhausted(_))
908 ));
909 arena.reset();
910 // After reset we can allocate again.
911 let _: &mut [u8] = arena.alloc::<u8>(32).unwrap();
912 }
913}