mnemosyne_arena/scratch/pool/
mod.rs

1//! Scratch buffer pool: a depth-tracked set of reusable aligned scratch
2//! buffers for temporal allocations.
3//!
4//! The pool is decomposed by responsibility (each submodule is private; its
5//! methods are re-exported through the `impl` blocks below):
6//! - `borrow` — the PROVISION-generic `borrow_slot`, `with_scratch`,
7//!   `with_scratch_bounded`, and `with_scratch_uninit` borrow entry points.
8//! - `manage` — provision-aware `release`, `reset`, `prewarm`, `preload`,
9//!   and `shrink_all_slots` lifecycle methods.
10//! - `query` — read-only `borrow_depth`, `capacity`, `slot_capacity`,
11//!   `is_available`, `total_capacity_bytes` accessors.
12//!
13//! `mod.rs` keeps only the struct shape, `MAX_POOL_SLOTS`, and the
14//! `new`/`with_slot_capacity`/`Default` constructors.
15
16mod borrow;
17mod manage;
18mod query;
19
20use super::aligned_vec::AlignedVec;
21use super::element::ScratchElement;
22use core::cell::{Cell, UnsafeCell};
23
24/// Maximum concurrent borrows (recursive/nested calls) the pool supports.
25pub const MAX_POOL_SLOTS: usize = 4;
26
27/// A pool of reusable, aligned scratch buffers for a specific element type.
28///
29/// `Send` but **not** `Sync` — designed for `thread_local!` storage.
30///
31/// # Usage
32///
33/// ```rust,ignore
34/// use mnemosyne_arena::scratch::ScratchPool;
35///
36/// thread_local! {
37///     static POOL: ScratchPool<f64> = ScratchPool::new();
38/// }
39///
40/// POOL.with(|pool| {
41///     pool.with_scratch(1024, |scratch| {
42///         // scratch: &mut [f64] of exactly 1024 elements, 64-byte aligned
43///     });
44/// });
45/// ```
46pub struct ScratchPool<T: ScratchElement> {
47    pub(super) slots: [UnsafeCell<AlignedVec<T>>; MAX_POOL_SLOTS],
48    pub(super) borrow_depth: Cell<u8>,
49    /// Per-depth high-water request, recorded by
50    /// [`with_scratch_bounded`](Self::with_scratch_bounded) and honored by
51    /// [`release`](Self::release). Provisioned slots keep capacity for their
52    /// working set across a release; unprovisioned slots reclaim entirely.
53    pub(super) provisions: [Cell<usize>; MAX_POOL_SLOTS],
54    /// Slot 0's capacity, republished by the borrow that grows it.
55    ///
56    /// [`ScratchPool::capacity`] is reachable from inside a live
57    /// [`ScratchPool::with_scratch`] borrow through entirely safe code (both
58    /// take `&self`, and the pool's documented home is a `thread_local!`), so
59    /// the accessor must not derive a reference into a slot that borrow already
60    /// holds exclusively. Mirroring the figure outside the `UnsafeCell` removes
61    /// the aliasing by construction rather than forbidding the call.
62    ///
63    /// Slot 0's capacity changes only where this is written: construction, and
64    /// the grow branch of a depth-0 `with_scratch` (slot index equals borrow
65    /// depth, so only depth 0 touches slot 0). A `debug_assert!` in
66    /// `with_scratch` fails the tests if a future mutation path escapes that
67    /// set.
68    pub(super) slot_capacities: [Cell<usize>; MAX_POOL_SLOTS],
69}
70
71// SAFETY: a `ScratchPool` uniquely owns its slot buffers (each `AlignedVec` owns
72// its heap storage with no aliasing), so moving the whole pool to another thread
73// is sound. It is deliberately *not* `Sync`: the `UnsafeCell` slots and the
74// `Cell` borrow-depth and capacity fields are guarded only by single-threaded
75// `borrow_depth` tracking, which assumes one thread at a time (`thread_local!`
76// storage), so it must never be shared by reference across threads.
77unsafe impl<T: ScratchElement> Send for ScratchPool<T> {}
78
79impl<T: ScratchElement> Default for ScratchPool<T> {
80    #[inline]
81    fn default() -> Self {
82        Self::new()
83    }
84}
85
86impl<T: ScratchElement> ScratchPool<T> {
87    /// Creates a new empty scratch pool (zero allocation at construction).
88    #[inline]
89    pub const fn new() -> Self {
90        Self {
91            slots: [
92                UnsafeCell::new(AlignedVec::dangling()),
93                UnsafeCell::new(AlignedVec::dangling()),
94                UnsafeCell::new(AlignedVec::dangling()),
95                UnsafeCell::new(AlignedVec::dangling()),
96            ],
97            borrow_depth: Cell::new(0),
98            provisions: [const { Cell::new(0) }; MAX_POOL_SLOTS],
99            slot_capacities: [const { Cell::new(0) }; MAX_POOL_SLOTS],
100        }
101    }
102
103    /// Creates a new scratch pool with pre-allocated capacity per slot.
104    #[inline]
105    pub fn with_slot_capacity(capacity: usize) -> Self {
106        let mk = || {
107            if capacity == 0 {
108                AlignedVec::dangling()
109            } else {
110                AlignedVec::with_capacity(capacity)
111            }
112        };
113        Self {
114            slots: [
115                UnsafeCell::new(mk()),
116                UnsafeCell::new(mk()),
117                UnsafeCell::new(mk()),
118                UnsafeCell::new(mk()),
119            ],
120            borrow_depth: Cell::new(0),
121            provisions: [const { Cell::new(0) }; MAX_POOL_SLOTS],
122            // `mk()` gives every slot exactly `capacity` (a zero request yields
123            // the zero-capacity dangling sentinel), so every slot's mirror must
124            // begin with the same warm capacity. This keeps the public
125            // `slot_capacity`/`total_capacity_bytes` figures in sync with the
126            // actual backing allocations across all prewarmed slots.
127            slot_capacities: {
128                let mut mirrors = [const { Cell::new(0) }; MAX_POOL_SLOTS];
129                for cell in &mut mirrors {
130                    cell.set(capacity);
131                }
132                mirrors
133            },
134        }
135    }
136}
mnemosyne_arena/scratch/pool/mod.rs

mnemosyne_arena/scratch/pool/
mod.rs
