Skip to main content

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}