use std::cell::RefCell;
use std::ops::{Deref, DerefMut};
use std::ptr::NonNull;
pub trait MemoryPoolObject {
fn alloc() -> Option<NonNull<Self>>;
fn dealloc(obj: NonNull<Self>);
}
pub struct PooledObject<'a, T: MemoryPoolObject> {
inner: NonNull<T>,
allocator: &'a MemoryPool<T>,
}
impl<'a, T: MemoryPoolObject> Deref for PooledObject<'a, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
unsafe { self.inner.as_ref() }
}
}
impl<'a, T: MemoryPoolObject> DerefMut for PooledObject<'a, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { self.inner.as_mut() }
}
}
impl<'a, T: MemoryPoolObject> Drop for PooledObject<'a, T> {
fn drop(&mut self) {
let pool = self.allocator;
if pool.free_list.borrow().len() + 1 >= pool.free_list_limit {
T::dealloc(self.inner);
} else {
let mut free_list = pool.free_list.borrow_mut();
free_list.push(self.inner);
}
self.inner = NonNull::dangling();
}
}
pub struct MemoryPool<T: MemoryPoolObject> {
free_list: RefCell<Vec<NonNull<T>>>,
free_list_limit: usize,
}
impl<T: MemoryPoolObject> Drop for MemoryPool<T> {
fn drop(&mut self) {
for cached_obj in self.free_list.take() {
T::dealloc(cached_obj);
}
}
}
impl<T: MemoryPoolObject> MemoryPool<T> {
pub fn with_capacity(_size: usize) -> Self {
Self {
free_list: RefCell::new(Vec::new()),
free_list_limit: 0,
}
}
pub fn alloc_obj(&self) -> Option<PooledObject<T>> {
if self.free_list.borrow().len() > 0 {
let mut free_list = self.free_list.borrow_mut();
Some(PooledObject {
inner: free_list.pop().unwrap(),
allocator: self,
})
} else {
Some(PooledObject {
inner: T::alloc()?,
allocator: self,
})
}
}
}