use core::{
cell::UnsafeCell,
marker::PhantomData,
mem::MaybeUninit,
ops::{Deref, DerefMut},
panic::{RefUnwindSafe, UnwindSafe},
ptr::null_mut,
sync::atomic::AtomicPtr,
};
use crate::{
bindings::{rt_mutex, rt_pool, rt_pool_get, rt_pool_put},
cell::SyncUnsafeCell,
list::list_init,
ptr_macros::{ptr_to_field, ptr_to_field_mut},
sync::condvar::c_cond_init,
};
pub struct Pool<T: ?Sized> {
pool: UnsafeCell<rt_pool>,
_phantom_data: PhantomData<T>,
}
unsafe impl<T: ?Sized + Send> Send for Pool<T> {}
unsafe impl<T: ?Sized + Send> Sync for Pool<T> {}
impl<T: ?Sized> UnwindSafe for Pool<T> {}
impl<T: ?Sized> RefUnwindSafe for Pool<T> {}
impl<T> Pool<T> {
#[must_use]
pub const unsafe fn init<const N: usize>(
this: *const Self,
ptrs: &'static [AtomicPtr<T>; N],
) -> Pool<T> {
let pool = UnsafeCell::raw_get(ptr_to_field!(this, pool));
Pool {
pool: UnsafeCell::new(rt_pool {
mutex: rt_mutex {
holder: 0,
wait_list: list_init(ptr_to_field_mut!(pool, mutex, wait_list)),
list: list_init(ptr_to_field_mut!(pool, mutex, list)),
level: -1,
},
cond: c_cond_init(ptr_to_field_mut!(pool, cond)),
ptrs: ptrs.as_ptr().cast_mut().cast(),
avail: N,
}),
_phantom_data: PhantomData,
}
}
#[must_use]
pub const fn make_ptrs<const N: usize>(
data: &'static [SyncUnsafeCell<MaybeUninit<T>>; N],
) -> [AtomicPtr<T>; N] {
let mut ptrs = [const { AtomicPtr::new(null_mut()) }; N];
let mut i = 0;
while i < N {
ptrs[i] = AtomicPtr::new(data[i].get().cast());
i += 1;
}
ptrs
}
#[inline]
pub fn alloc(&self, val: T) -> PoolBox<'_, T> {
let pool_ptr = self.pool.get();
let ptr: *mut T = unsafe { rt_pool_get(pool_ptr) }.cast();
unsafe { ptr.write(val) }
PoolBox { ptr, pool: self }
}
}
pub struct PoolBox<'a, T: ?Sized> {
ptr: *mut T,
pool: &'a Pool<T>,
}
impl<T: ?Sized> Deref for PoolBox<'_, T> {
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.ptr }
}
}
impl<T: ?Sized> DerefMut for PoolBox<'_, T> {
fn deref_mut(&mut self) -> &mut T {
unsafe { &mut *self.ptr }
}
}
impl<T: ?Sized> Drop for PoolBox<'_, T> {
#[inline]
fn drop(&mut self) {
unsafe { self.ptr.drop_in_place() }
let pool_ptr = self.pool.pool.get();
let ptr = self.ptr.cast();
unsafe { rt_pool_put(pool_ptr, ptr) }
}
}
unsafe impl<T: Send> Send for PoolBox<'_, T> {}
unsafe impl<T: Sync> Sync for PoolBox<'_, T> {}
impl<T> UnwindSafe for PoolBox<'_, T> {}
impl<T> RefUnwindSafe for PoolBox<'_, T> {}
#[macro_export]
macro_rules! pool {
($name: ident, $type: ty, $num: expr) => {
static $name: $crate::sync::Pool<$type> = {
use core::{mem::MaybeUninit, sync::atomic::AtomicPtr};
use $crate::cell::SyncUnsafeCell;
let ptr = &raw const $name;
static DATA: [SyncUnsafeCell<MaybeUninit<$type>>; $num] =
[const { SyncUnsafeCell::new(MaybeUninit::zeroed()) }; $num];
static PTRS: [AtomicPtr<$type>; $num] = $crate::sync::Pool::make_ptrs(&DATA);
unsafe { $crate::sync::Pool::init(ptr, &PTRS) }
};
};
}
#[cfg(test)]
mod tests {
crate::pool!(POOL, i32, 10);
#[test]
fn fast_path() {
let mut x = POOL.alloc(0);
*x += 1;
}
}