use std::{
convert::Infallible,
mem::ManuallyDrop,
ops::{Deref, DerefMut},
};
use stable_deref_trait::StableDeref;
use super::{PoolError, error::catch_operation};
pub trait PoolItem: Send + 'static {
fn reset(&mut self);
}
pub trait PoolProvider<T: PoolItem>: Send + 'static {
type Entry: DerefMut<Target = T> + StableDeref + Send + 'static;
fn take<F, E>(&self, create: F) -> Result<Self::Entry, PoolError<E>>
where
F: FnOnce() -> Result<T, E>;
fn return_entry(&self, entry: Self::Entry) -> Result<(), Self::Entry>;
fn warm<F, E>(&self, count: usize, create: F) -> Result<usize, PoolError<E>>
where
F: FnMut() -> Result<T, E>;
}
#[allow(clippy::redundant_pub_crate)]
pub(super) struct Storage<E, const N: usize> {
slots: [Option<E>; N],
len: usize,
}
impl<E, const N: usize> Storage<E, N> {
pub(super) const fn new() -> Self {
Self {
slots: [const { None }; N],
len: 0,
}
}
#[inline(always)]
pub(super) const fn pop(&mut self) -> Option<E> {
if self.len == 0 {
return None;
}
self.len -= 1;
self.slots[self.len].take()
}
#[inline(always)]
pub(super) fn try_push(&mut self, entry: E) -> Result<(), E> {
if self.len < N {
self.slots[self.len] = Some(entry);
self.len += 1;
Ok(())
} else {
Err(entry)
}
}
#[inline(always)]
pub(super) const fn len(&self) -> usize {
self.len
}
pub(super) fn extend_newest_first(&mut self, entries: &mut impl Iterator<Item = E>) {
let start = self.len;
while self.len < N {
let Some(entry) = entries.next() else {
break;
};
self.slots[self.len] = Some(entry);
self.len += 1;
}
self.slots[start..self.len].reverse();
}
}
pub struct PoolGuard<T: PoolItem, P: PoolProvider<T>> {
node: ManuallyDrop<P::Entry>,
provider: P,
}
impl<T: PoolItem, P: PoolProvider<T>> PoolGuard<T, P> {
#[inline(always)]
pub fn acquire(provider: P) -> Result<Self, PoolError<Infallible>>
where
T: Default,
{
Self::acquire_with(provider, T::default)
}
#[inline(always)]
pub fn acquire_with<F>(provider: P, create: F) -> Result<Self, PoolError<Infallible>>
where
F: FnOnce() -> T,
{
Self::try_acquire_with(provider, || Ok::<T, Infallible>(create()))
}
#[inline(always)]
pub fn try_acquire_with<F, E>(provider: P, create: F) -> Result<Self, PoolError<E>>
where
F: FnOnce() -> Result<T, E>,
{
Ok(Self {
node: ManuallyDrop::new(
catch_operation(|| provider.take(create)).map_err(PoolError::Panic)??,
),
provider,
})
}
}
impl<T: PoolItem, P: PoolProvider<T>> Deref for PoolGuard<T, P> {
type Target = T;
#[inline(always)]
fn deref(&self) -> &Self::Target {
self.node.deref()
}
}
impl<T: PoolItem, P: PoolProvider<T>> DerefMut for PoolGuard<T, P> {
#[inline(always)]
fn deref_mut(&mut self) -> &mut Self::Target {
self.node.deref_mut()
}
}
unsafe impl<T: PoolItem, P: PoolProvider<T>> StableDeref for PoolGuard<T, P> {}
impl<T: PoolItem, P: PoolProvider<T>> Drop for PoolGuard<T, P> {
#[inline(always)]
fn drop(&mut self) {
let mut node = unsafe { ManuallyDrop::take(&mut self.node) };
node.reset();
if let Err(node) = self.provider.return_entry(node) {
discard_entry(node);
}
}
}
#[cold]
#[inline(never)]
fn discard_entry<E>(entry: E) {
drop(entry);
}
#[cfg(test)]
mod tests;