use alloc::{collections::VecDeque, fmt, rc::Rc};
use core::{
cell::UnsafeCell,
hash::{Hash, Hasher},
marker::PhantomData,
mem::{forget, MaybeUninit},
ops::{Deref, DerefMut},
ptr,
};
use crate::PoolAllocator;
#[derive(Debug)]
pub struct LocalPool<P: PoolAllocator<T>, T> {
allocator: P,
storage: UnsafeCell<VecDeque<T>>,
_phantom: PhantomData<*mut usize>,
}
impl<P: PoolAllocator<T>, T> LocalPool<P, T> {
pub fn new_prefilled(pool_size: usize, allocator: P) -> Self {
let mut storage = VecDeque::with_capacity(pool_size);
for _ in 0..pool_size {
storage.push_back(allocator.allocate());
}
LocalPool {
allocator,
storage: UnsafeCell::new(storage),
_phantom: PhantomData,
}
}
pub fn new(pool_size: usize, allocator: P) -> Self {
LocalPool {
allocator,
storage: UnsafeCell::new(VecDeque::with_capacity(pool_size)),
_phantom: PhantomData,
}
}
#[allow(clippy::mut_from_ref)]
fn storage_mut(&self) -> &mut VecDeque<T> {
unsafe { &mut *self.storage.get() }
}
#[allow(clippy::mut_from_ref)]
fn storage_borrow(&self) -> &VecDeque<T> {
unsafe { &*self.storage.get() }
}
pub fn to_rc(self) -> Rc<Self> {
Rc::new(self)
}
pub fn try_get(&self) -> Option<RefLocalGuard<'_, P, T>> {
self.storage_mut().pop_front().map(|mut obj| {
self.allocator.reset(&mut obj);
RefLocalGuard::new(obj, self)
})
}
pub fn get(&'_ self) -> RefLocalGuard<'_, P, T> {
match self.storage_mut().pop_front() {
Some(mut obj) => {
self.allocator.reset(&mut obj);
RefLocalGuard::new(obj, self)
}
None => RefLocalGuard::new(self.allocator.allocate(), self),
}
}
pub fn try_get_rc(self: Rc<Self>) -> Option<RcLocalGuard<P, T>> {
self.storage_mut().pop_front().map(|mut obj| {
self.allocator.reset(&mut obj);
RcLocalGuard::new(obj, &self)
})
}
pub fn get_rc(self: Rc<Self>) -> RcLocalGuard<P, T> {
match self.storage_mut().pop_front() {
Some(mut obj) => {
self.allocator.reset(&mut obj);
RcLocalGuard::new(obj, &self)
}
None => RcLocalGuard::new(self.allocator.allocate(), &self),
}
}
pub fn len(&self) -> usize {
self.storage_borrow().len()
}
pub fn is_empty(&self) -> bool {
self.storage_borrow().is_empty()
}
pub fn capacity(&self) -> usize {
self.storage_borrow().capacity()
}
}
pub struct RefLocalGuard<'a, P: PoolAllocator<T>, T> {
obj: MaybeUninit<T>,
pool: &'a LocalPool<P, T>,
}
impl<'a, P: PoolAllocator<T>, T> RefLocalGuard<'a, P, T> {
fn new(obj: T, pool: &'a LocalPool<P, T>) -> Self {
RefLocalGuard {
obj: MaybeUninit::new(obj),
pool,
}
}
pub fn into_inner(self) -> T {
let obj = unsafe { self.obj.as_ptr().read() };
forget(self);
obj
}
}
impl<'a, P: PoolAllocator<T>, T> Deref for RefLocalGuard<'a, P, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
unsafe { &*self.obj.as_ptr() }
}
}
impl<'a, P: PoolAllocator<T>, T> DerefMut for RefLocalGuard<'a, P, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { &mut *self.obj.as_mut_ptr() }
}
}
impl<'a, P: PoolAllocator<T>, T> Drop for RefLocalGuard<'a, P, T> {
fn drop(&mut self) {
let storage = self.pool.storage_mut();
if self.pool.allocator.is_valid(self.deref()) && storage.len() < storage.capacity() {
storage.push_back(unsafe { ptr::read(self.obj.as_mut_ptr()) });
} else {
unsafe {
ptr::drop_in_place(self.obj.as_mut_ptr());
}
}
}
}
impl<'a, P: PoolAllocator<T>, T: Hash> Hash for RefLocalGuard<'a, P, T> {
#[inline]
fn hash<H: Hasher>(&self, state: &mut H) {
(**self).hash(state);
}
}
impl<'a, P: PoolAllocator<T>, T: fmt::Display> fmt::Display for RefLocalGuard<'a, P, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(&**self, f)
}
}
impl<'a, P: PoolAllocator<T>, T: fmt::Debug> fmt::Debug for RefLocalGuard<'a, P, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&**self, f)
}
}
impl<'a, P: PoolAllocator<T>, T> fmt::Pointer for RefLocalGuard<'a, P, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Pointer::fmt(&(&**self as *const T), f)
}
}
impl<'a, P: PoolAllocator<T>, T: PartialEq> PartialEq for RefLocalGuard<'a, P, T> {
#[inline]
fn eq(&self, other: &RefLocalGuard<'a, P, T>) -> bool {
self.deref().eq(other)
}
}
impl<'a, P: PoolAllocator<T>, T: Eq> Eq for RefLocalGuard<'a, P, T> {}
impl<'a, P: PoolAllocator<T>, T: PartialOrd> PartialOrd for RefLocalGuard<'a, P, T> {
#[inline]
fn partial_cmp(&self, other: &RefLocalGuard<'a, P, T>) -> Option<core::cmp::Ordering> {
(**self).partial_cmp(&**other)
}
#[inline]
fn lt(&self, other: &RefLocalGuard<'a, P, T>) -> bool {
**self < **other
}
#[inline]
fn le(&self, other: &RefLocalGuard<'a, P, T>) -> bool {
**self <= **other
}
#[inline]
fn gt(&self, other: &RefLocalGuard<'a, P, T>) -> bool {
**self > **other
}
#[inline]
fn ge(&self, other: &RefLocalGuard<'a, P, T>) -> bool {
**self >= **other
}
}
impl<'a, P: PoolAllocator<T>, T: Ord> Ord for RefLocalGuard<'a, P, T> {
#[inline]
fn cmp(&self, other: &RefLocalGuard<'a, P, T>) -> core::cmp::Ordering {
(**self).cmp(&**other)
}
}
impl<'a, P: PoolAllocator<T>, T> core::borrow::Borrow<T> for RefLocalGuard<'a, P, T> {
#[inline(always)]
fn borrow(&self) -> &T {
self
}
}
impl<'a, P: PoolAllocator<T>, T> AsRef<T> for RefLocalGuard<'a, P, T> {
#[inline(always)]
fn as_ref(&self) -> &T {
self
}
}
pub struct RcLocalGuard<P: PoolAllocator<T>, T> {
obj: MaybeUninit<T>,
pool: Rc<LocalPool<P, T>>,
}
impl<P: PoolAllocator<T>, T> RcLocalGuard<P, T> {
fn new(obj: T, pool: &Rc<LocalPool<P, T>>) -> Self {
Self {
obj: MaybeUninit::new(obj),
pool: pool.clone(),
}
}
pub fn into_inner(mut self) -> T {
let obj = unsafe { self.obj.as_ptr().read() };
unsafe { ptr::drop_in_place(&mut self.pool) }
forget(self);
obj
}
}
impl<P: PoolAllocator<T>, T> Deref for RcLocalGuard<P, T> {
type Target = T;
#[inline(always)]
fn deref(&self) -> &Self::Target {
unsafe { &*self.obj.as_ptr() }
}
}
impl<P: PoolAllocator<T>, T> DerefMut for RcLocalGuard<P, T> {
#[inline(always)]
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { &mut *self.obj.as_mut_ptr() }
}
}
impl<P: PoolAllocator<T>, T> Drop for RcLocalGuard<P, T> {
fn drop(&mut self) {
let storage = self.pool.storage_mut();
if self.pool.allocator.is_valid(self.deref()) && storage.len() < storage.capacity() {
storage.push_back(unsafe { ptr::read(self.obj.as_mut_ptr()) });
} else {
unsafe {
ptr::drop_in_place(self.obj.as_mut_ptr());
}
}
}
}
impl<P: PoolAllocator<T>, T: Hash> Hash for RcLocalGuard<P, T> {
#[inline]
fn hash<H: Hasher>(&self, state: &mut H) {
(**self).hash(state);
}
}
impl<P: PoolAllocator<T>, T: fmt::Display> fmt::Display for RcLocalGuard<P, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(&**self, f)
}
}
impl<P: PoolAllocator<T>, T: fmt::Debug> fmt::Debug for RcLocalGuard<P, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&**self, f)
}
}
impl<P: PoolAllocator<T>, T> fmt::Pointer for RcLocalGuard<P, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Pointer::fmt(&(&**self as *const T), f)
}
}
impl<P: PoolAllocator<T>, T: PartialEq> PartialEq for RcLocalGuard<P, T> {
#[inline]
fn eq(&self, other: &RcLocalGuard<P, T>) -> bool {
self.deref().eq(other)
}
}
impl<P: PoolAllocator<T>, T: Eq> Eq for RcLocalGuard<P, T> {}
impl<P: PoolAllocator<T>, T: PartialOrd> PartialOrd for RcLocalGuard<P, T> {
#[inline]
fn partial_cmp(&self, other: &RcLocalGuard<P, T>) -> Option<core::cmp::Ordering> {
(**self).partial_cmp(&**other)
}
#[inline]
fn lt(&self, other: &RcLocalGuard<P, T>) -> bool {
**self < **other
}
#[inline]
fn le(&self, other: &RcLocalGuard<P, T>) -> bool {
**self <= **other
}
#[inline]
fn gt(&self, other: &RcLocalGuard<P, T>) -> bool {
**self > **other
}
#[inline]
fn ge(&self, other: &RcLocalGuard<P, T>) -> bool {
**self >= **other
}
}
impl<P: PoolAllocator<T>, T: Ord> Ord for RcLocalGuard<P, T> {
#[inline]
fn cmp(&self, other: &RcLocalGuard<P, T>) -> core::cmp::Ordering {
(**self).cmp(&**other)
}
}
impl<P: PoolAllocator<T>, T> core::borrow::Borrow<T> for RcLocalGuard<P, T> {
#[inline(always)]
fn borrow(&self) -> &T {
self
}
}
impl<P: PoolAllocator<T>, T> AsRef<T> for RcLocalGuard<P, T> {
#[inline(always)]
fn as_ref(&self) -> &T {
self
}
}