osiris/types/
pool.rs

1//! This module provides pool allocator implementations.
2#![allow(dead_code)]
3
4use core::{
5    cell::UnsafeCell,
6    marker::PhantomData,
7    mem::MaybeUninit,
8    num::NonZeroUsize,
9    ops::{Deref, DerefMut, Range},
10    ptr::{self, write},
11};
12
13use crate::{sync::spinlock::SpinLocked, types::bitset::BitAlloc};
14
15pub struct FixedPoolRef<'a, T, const N: usize, const WORDS: usize> {
16    idx: usize,
17    pool: &'a FixedPool<T, N, WORDS>,
18    _marker: PhantomData<T>,
19}
20
21impl<'a, T, const N: usize, const WORDS: usize> Deref for FixedPoolRef<'a, T, N, WORDS> {
22    type Target = T;
23
24    fn deref(&self) -> &Self::Target {
25        // Safety: ptr does always point to a valid block in the pool.
26        // Only one Ref can exist for a block at a time, so there are no mutable references to the same block.
27        unsafe { &*self.pool.access(self.idx) }
28    }
29}
30
31impl<'a, T, const N: usize, const WORDS: usize> DerefMut for FixedPoolRef<'a, T, N, WORDS> {
32    fn deref_mut(&mut self) -> &mut Self::Target {
33        // Safety: ptr does always point to a valid block in the pool.
34        // Only one Ref can exist for a block at a time, so there are no mutable references to the same block.
35        unsafe { &mut *self.pool.access(self.idx) }
36    }
37}
38
39impl<T, const N: usize, const WORDS: usize> Drop for FixedPoolRef<'_, T, N, WORDS> {
40    fn drop(&mut self) {
41        // Safety: ptr does always point to a valid block in the pool.
42        unsafe { ptr::drop_in_place(self.pool.access(self.idx)) };
43        self.pool.free(self.idx);
44    }
45}
46
47pub struct FixedPool<T, const N: usize, const WORDS: usize> {
48    free: SpinLocked<BitAlloc<WORDS>>,
49    blocks: [UnsafeCell<MaybeUninit<T>>; N],
50}
51
52impl<T, const N: usize, const WORDS: usize> FixedPool<T, N, WORDS> {
53    pub const fn new() -> Self {
54        Self {
55            free: SpinLocked::new(BitAlloc::from_array([!0usize; WORDS])),
56            blocks: [const { UnsafeCell::new(MaybeUninit::uninit()) }; N],
57        }
58    }
59
60    pub fn alloc(&self, new: T) -> Option<FixedPoolRef<'_, T, N, WORDS>> {
61        // Safety: Alloc ensures that the index cannot be allocated until the next free.
62        // A free can only happen when the Ref is dropped, as the function is not publicly accessible.
63        // This guarantees that only one Ref can exist for a block at a time.
64        let mut alloc = self.free.lock();
65        let idx = alloc.alloc(1)?;
66        // BitAlloc<WORDS> exposes WORDS * BITS_PER_WORD bits, which can exceed N.
67        // Release any out-of-range index so the pool reports exhaustion instead of panicking.
68        if idx >= N {
69            alloc.free(idx, 1);
70            return None;
71        }
72        drop(alloc);
73        let ptr = self.blocks[idx].get();
74        // Safety: A block can only be allocated once.
75        unsafe { ptr.write(MaybeUninit::new(new)) };
76        Some(FixedPoolRef {
77            idx,
78            pool: self,
79            _marker: PhantomData,
80        })
81    }
82
83    fn access(&self, idx: usize) -> *mut T {
84        self.blocks[idx].get() as *mut T
85    }
86
87    fn free(&self, idx: usize) {
88        self.free.lock().free(idx, 1);
89    }
90}
91
92/// Meta information for a block in the pool.
93struct SizedPoolMeta {
94    _size: usize,
95    next: Option<NonZeroUsize>,
96}
97
98/// A pool allocator that allocates fixed-size blocks.
99#[deprecated(note = "Will be removed soon. Do not use!")]
100pub struct SizedPool<T: Default> {
101    head: Option<NonZeroUsize>,
102    _marker: PhantomData<T>,
103}
104
105#[allow(deprecated)]
106impl<T: Default> Default for SizedPool<T> {
107    fn default() -> Self {
108        Self::new()
109    }
110}
111
112#[allow(deprecated)]
113impl<T: Default> SizedPool<T> {
114    /// Create a new empty pool.
115    pub const fn new() -> Self {
116        Self {
117            head: None,
118            _marker: PhantomData,
119        }
120    }
121
122    /// Calculate the padding required to align the block to `align_of::<T>()`.
123    const fn align_up() -> usize {
124        let meta = size_of::<SizedPoolMeta>();
125        let align = align_of::<T>();
126        // Calculate the padding required to align the block.
127        (align - (meta % align)) % align
128    }
129
130    /// Add a range of blocks to the pool.
131    ///
132    /// `range` - The range of blocks to add.
133    ///
134    /// # Safety
135    ///
136    /// The caller must ensure that the range is valid and that the blocks are at least the size of `T` + `SizedPoolMeta` + Padding for `T`.
137    pub unsafe fn add_range(&mut self, range: Range<usize>) {
138        let mut ptr = range.start;
139
140        while ptr < range.end {
141            unsafe {
142                self.add_block(ptr);
143            }
144
145            ptr += Self::align_up() + size_of::<SizedPoolMeta>() + size_of::<T>();
146        }
147    }
148
149    /// Add a block to the pool.
150    ///
151    /// `ptr` - The pointer to the block to add.
152    ///
153    /// # Safety
154    ///
155    /// The caller must ensure that the pointer is valid and that the block is at least the size of `T` + `SizedPoolMeta` + Padding for `T`.
156    unsafe fn add_block(&mut self, ptr: usize) {
157        let meta = SizedPoolMeta {
158            _size: size_of::<T>(),
159            next: self.head,
160        };
161
162        unsafe {
163            write(ptr as *mut SizedPoolMeta, meta);
164        }
165
166        self.head = Some(unsafe { NonZeroUsize::new_unchecked(ptr) });
167    }
168
169    /// Allocate a block from the pool.
170    ///
171    /// Returns `Some(Owned<T>)` if a block was successfully allocated, otherwise `None`.
172    pub fn alloc(&mut self) -> Option<Owned<T>> {
173        let head = self.head.take();
174
175        head.map(|head| {
176            let meta = unsafe { &*(head.get() as *const SizedPoolMeta) };
177            self.head = meta.next;
178
179            let ptr = head.get() + size_of::<SizedPoolMeta>() + Self::align_up();
180            unsafe { write(ptr as *mut T, T::default()) };
181
182            Owned { ptr: ptr as *mut T }
183        })
184    }
185
186    /// Deallocate a block back to the pool.
187    ///
188    /// `block` - The block to deallocate.
189    pub fn dealloc(&mut self, block: Owned<T>) {
190        let ptr = block.ptr as usize - size_of::<SizedPoolMeta>() - Self::align_up();
191
192        // Append ptr to the front of the list.
193        let head = self
194            .head
195            .replace(unsafe { NonZeroUsize::new_unchecked(ptr) });
196
197        // Update the next pointer to the previous head.
198        let meta = unsafe { &mut *(ptr as *mut SizedPoolMeta) };
199        meta.next = head;
200    }
201}
202
203/// An owned block from a pool.
204pub struct Owned<T> {
205    ptr: *mut T,
206}
207
208impl<T: Default> Deref for Owned<T> {
209    type Target = T;
210
211    fn deref(&self) -> &Self::Target {
212        unsafe { &*self.ptr }
213    }
214}
215
216impl<T: Default> DerefMut for Owned<T> {
217    fn deref_mut(&mut self) -> &mut Self::Target {
218        unsafe { &mut *self.ptr }
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225
226    #[test]
227    fn fixed_pool_alloc_beyond_n_returns_none() {
228        let pool: FixedPool<u32, 4, 1> = FixedPool::new();
229        let _r0 = pool.alloc(0).unwrap();
230        let _r1 = pool.alloc(1).unwrap();
231        let _r2 = pool.alloc(2).unwrap();
232        let _r3 = pool.alloc(3).unwrap();
233        assert!(pool.alloc(4).is_none());
234    }
235}
osiris/types/pool.rs

osiris/types/
pool.rs
