atomic_pool/
lib.rs

1#![cfg_attr(not(test), no_std)]
2
3mod atomic_bitset;
4
5use core::cell::UnsafeCell;
6use core::hash::{Hash, Hasher};
7use core::mem::MaybeUninit;
8use core::ops::{Deref, DerefMut};
9use core::{cmp, mem, ptr::NonNull};
10use portable_atomic::AtomicU32;
11
12use crate::atomic_bitset::AtomicBitset;
13
14/// Implementation detail. Not covered by semver guarantees.
15#[doc(hidden)]
16pub trait PoolStorage<T> {
17    fn alloc(&self) -> Option<NonNull<T>>;
18    unsafe fn free(&self, p: NonNull<T>);
19}
20
21/// Implementation detail. Not covered by semver guarantees.
22#[doc(hidden)]
23pub struct PoolStorageImpl<T, const N: usize, const K: usize>
24where
25    [AtomicU32; K]: Sized,
26{
27    used: AtomicBitset<N, K>,
28    data: [UnsafeCell<MaybeUninit<T>>; N],
29}
30
31unsafe impl<T, const N: usize, const K: usize> Send for PoolStorageImpl<T, N, K> {}
32unsafe impl<T, const N: usize, const K: usize> Sync for PoolStorageImpl<T, N, K> {}
33
34impl<T, const N: usize, const K: usize> PoolStorageImpl<T, N, K>
35where
36    [AtomicU32; K]: Sized,
37{
38    const UNINIT: UnsafeCell<MaybeUninit<T>> = UnsafeCell::new(MaybeUninit::uninit());
39
40    pub const fn new() -> Self {
41        Self {
42            used: AtomicBitset::new(),
43            data: [Self::UNINIT; N],
44        }
45    }
46}
47
48impl<T, const N: usize, const K: usize> PoolStorage<T> for PoolStorageImpl<T, N, K>
49where
50    [AtomicU32; K]: Sized,
51{
52    fn alloc(&self) -> Option<NonNull<T>> {
53        let n = self.used.alloc()?;
54        let p = self.data[n].get() as *mut T;
55        Some(unsafe { NonNull::new_unchecked(p) })
56    }
57
58    /// safety: p must be a pointer obtained from self.alloc that hasn't been freed yet.
59    unsafe fn free(&self, p: NonNull<T>) {
60        let origin = self.data.as_ptr() as *mut T;
61        let n = p.as_ptr().offset_from(origin);
62        assert!(n >= 0);
63        assert!((n as usize) < N);
64        self.used.free(n as usize);
65    }
66}
67
68pub trait Pool: 'static {
69    type Item: 'static;
70
71    /// Implementation detail. Not covered by semver guarantees.
72    #[doc(hidden)]
73    type Storage: PoolStorage<Self::Item>;
74
75    /// Implementation detail. Not covered by semver guarantees.
76    #[doc(hidden)]
77    fn get() -> &'static Self::Storage;
78}
79
80pub struct Box<P: Pool> {
81    ptr: NonNull<P::Item>,
82}
83
84impl<P: Pool> Box<P> {
85    pub fn new(item: P::Item) -> Option<Self> {
86        let p = match P::get().alloc() {
87            Some(p) => p,
88            None => return None,
89        };
90        unsafe { p.as_ptr().write(item) };
91        Some(Self { ptr: p })
92    }
93
94    pub fn into_raw(b: Self) -> NonNull<P::Item> {
95        let res = b.ptr;
96        mem::forget(b);
97        res
98    }
99
100    pub unsafe fn from_raw(ptr: NonNull<P::Item>) -> Self {
101        Self { ptr }
102    }
103}
104
105impl<P: Pool> Drop for Box<P> {
106    fn drop(&mut self) {
107        unsafe {
108            //trace!("dropping {:u32}", self.ptr as u32);
109            self.ptr.as_ptr().drop_in_place();
110            P::get().free(self.ptr);
111        };
112    }
113}
114
115unsafe impl<P: Pool> Send for Box<P> where P::Item: Send {}
116
117unsafe impl<P: Pool> Sync for Box<P> where P::Item: Sync {}
118
119unsafe impl<P: Pool> stable_deref_trait::StableDeref for Box<P> {}
120
121impl<P: Pool> as_slice_01::AsSlice for Box<P>
122where
123    P::Item: as_slice_01::AsSlice,
124{
125    type Element = <P::Item as as_slice_01::AsSlice>::Element;
126
127    fn as_slice(&self) -> &[Self::Element] {
128        self.deref().as_slice()
129    }
130}
131
132impl<P: Pool> as_slice_01::AsMutSlice for Box<P>
133where
134    P::Item: as_slice_01::AsMutSlice,
135{
136    fn as_mut_slice(&mut self) -> &mut [Self::Element] {
137        self.deref_mut().as_mut_slice()
138    }
139}
140
141impl<P: Pool> as_slice_02::AsSlice for Box<P>
142where
143    P::Item: as_slice_02::AsSlice,
144{
145    type Element = <P::Item as as_slice_02::AsSlice>::Element;
146
147    fn as_slice(&self) -> &[Self::Element] {
148        self.deref().as_slice()
149    }
150}
151
152impl<P: Pool> as_slice_02::AsMutSlice for Box<P>
153where
154    P::Item: as_slice_02::AsMutSlice,
155{
156    fn as_mut_slice(&mut self) -> &mut [Self::Element] {
157        self.deref_mut().as_mut_slice()
158    }
159}
160
161impl<P: Pool> Deref for Box<P> {
162    type Target = P::Item;
163
164    fn deref(&self) -> &P::Item {
165        unsafe { self.ptr.as_ref() }
166    }
167}
168
169impl<P: Pool> DerefMut for Box<P> {
170    fn deref_mut(&mut self) -> &mut P::Item {
171        unsafe { self.ptr.as_mut() }
172    }
173}
174
175impl<P: Pool> core::fmt::Debug for Box<P>
176where
177    P::Item: core::fmt::Debug,
178{
179    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
180        <P::Item as core::fmt::Debug>::fmt(self, f)
181    }
182}
183
184impl<P: Pool> core::fmt::Display for Box<P>
185where
186    P::Item: core::fmt::Display,
187{
188    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
189        <P::Item as core::fmt::Display>::fmt(self, f)
190    }
191}
192
193impl<P: Pool> PartialEq for Box<P>
194where
195    P::Item: PartialEq,
196{
197    fn eq(&self, rhs: &Box<P>) -> bool {
198        <P::Item as PartialEq>::eq(self, rhs)
199    }
200}
201
202impl<P: Pool> Eq for Box<P> where P::Item: Eq {}
203
204impl<P: Pool> PartialOrd for Box<P>
205where
206    P::Item: PartialOrd,
207{
208    fn partial_cmp(&self, rhs: &Box<P>) -> Option<cmp::Ordering> {
209        <P::Item as PartialOrd>::partial_cmp(self, rhs)
210    }
211}
212
213impl<P: Pool> Ord for Box<P>
214where
215    P::Item: Ord,
216{
217    fn cmp(&self, rhs: &Box<P>) -> cmp::Ordering {
218        <P::Item as Ord>::cmp(self, rhs)
219    }
220}
221
222impl<P: Pool> Hash for Box<P>
223where
224    P::Item: Hash,
225{
226    fn hash<H>(&self, state: &mut H)
227    where
228        H: Hasher,
229    {
230        <P::Item as Hash>::hash(self, state)
231    }
232}
233
234#[macro_export]
235macro_rules! pool {
236    ($vis:vis $name:ident: [$ty:ty; $n:expr]) => {
237        $vis struct $name { _uninhabited: ::core::convert::Infallible }
238        impl $crate::Pool for $name {
239            type Item = $ty;
240            type Storage = $crate::PoolStorageImpl<$ty, {$n}, {($n+31)/32}>;
241            fn get() -> &'static Self::Storage {
242                static POOL: $crate::PoolStorageImpl<$ty, {$n}, {($n+31)/32}> = $crate::PoolStorageImpl::new();
243                &POOL
244            }
245        }
246    };
247}
248#[cfg(test)]
249mod test {
250    use super::*;
251
252    pool!(TestPool: [u32; 4]);
253
254    #[test]
255    fn test_pool() {
256        let b1 = Box::<TestPool>::new(111).unwrap();
257        let b2 = Box::<TestPool>::new(222).unwrap();
258        let b3 = Box::<TestPool>::new(333).unwrap();
259        let b4 = Box::<TestPool>::new(444).unwrap();
260        assert!(Box::<TestPool>::new(555).is_none());
261        assert_eq!(*b1, 111);
262        assert_eq!(*b2, 222);
263        assert_eq!(*b3, 333);
264        assert_eq!(*b4, 444);
265        mem::drop(b3);
266        let b5 = Box::<TestPool>::new(555).unwrap();
267        assert!(Box::<TestPool>::new(666).is_none());
268        assert_eq!(*b1, 111);
269        assert_eq!(*b2, 222);
270        assert_eq!(*b4, 444);
271        assert_eq!(*b5, 555);
272    }
273}
atomic_pool/lib.rs

atomic_pool/
lib.rs
