phash/src/pool.rs

//! Pool of objects
use parking_lot::{
    RwLock,
    Mutex,
};
use atomic_refcell::{
    AtomicRefCell,
    AtomicRefMut,
};
use std::{
    iter,
    ops,
    mem::{ ManuallyDrop, self },
    borrow::{
	Borrow,
	BorrowMut,
    },
};

#[derive(Debug)]
pub struct Pool<T>(RwLock<Vec<AtomicRefCell<T>>>);

#[derive(Debug)]
pub enum PoolRef<'a, T>
{
    Inside(AtomicRefMut<'a, T>),
    Outside(DeferredRef<'a, T>),
}

impl<'o, T: Send + Sync> PoolRef<'o, T>
{
    #[inline] 
    pub fn as_ref(&self) -> &T
    {
	self.borrow()
    }
    #[inline] 
    pub fn as_mut(&mut self) -> &mut T
    {
	self.borrow_mut()
    }
    /// Force insert into pool, then return the mutable ref guard.
    pub fn force(self) -> AtomicRefMut<'o, T>
    {
	match self {
	    Self::Inside(s) => s,
	    Self::Outside(deferred) => {
		let pool = deferred.1;
		let value = AtomicRefCell::new(deferred.detach());
		let mut writer = pool.0.write();
		let len = writer.len();
		writer.push(value);
		parking_lot::lock_api::RwLockWriteGuard::downgrade(writer)[len].borrow_mut()
	    },
	}
    }
    pub fn forget(self)
    {
	if let Self::Outside(v) = self {
	    v.forget()
	}
    }
    pub fn try_detach(self) -> Result<T, AtomicRefMut<'o, T>>
    {
	match self {
	    Self::Inside(v) => Err(v),
	    Self::Outside(v) => Ok(v.detach())
	}
    }
    pub fn replace_with<F>(&mut self, with: F) -> T
    where F: FnOnce() -> T
    {
	mem::replace(self.borrow_mut(), with())
    }
    #[inline] 
    pub fn replace(&mut self, val: T) -> T
    {
	self.replace_with(move || val)
    }
}
impl<'o, T: Send + Sync + Default> PoolRef<'o, T>
{
    #[inline(always)] 
    pub fn take(&mut self) -> T
    {
	self.replace_with(T::default)
    }

    #[inline] 
    pub fn into_inner(mut self) -> T
    {
	self.take()
    }
}


impl<'o, T> Borrow<T> for PoolRef<'o, T>
{
    #[inline] 
    fn borrow(&self) -> &T {
	ops::Deref::deref(self)
    }
}

impl<'o, T> BorrowMut<T> for PoolRef<'o, T>
{
    #[inline] 
    fn borrow_mut(&mut self) -> &mut T {
	&mut self
    }
}


impl<'o, T> ops::Deref for PoolRef<'o, T>
{
    type Target = T;
    #[inline] 
    fn deref(&self) -> &Self::Target {
	match self {
	    Self::Inside(s) => s.deref(),
	    Self::Outside(s) => s.deref(),
	}
    }
}

impl<'o, T> ops::DerefMut for PoolRef<'o, T>
{
    #[inline] 
    fn deref_mut(&mut self) -> &mut Self::Target {
	match self {
	    Self::Inside(s) => s.deref_mut(),
	    Self::Outside(s) => s.deref_mut(),
	}	
    }
}

#[derive(Debug)]
pub struct DeferredRef<'owner, T>(ManuallyDrop<T>, &'owner Pool<T>);


impl<'o, T: Send + Sync> DeferredRef<'o, T>
{
    #[inline(always)] 
    pub fn detach(mut self) -> T
    {
	let v= unsafe {
	    ManuallyDrop::take(&mut self.0)
	};
	mem::forget(self);
	v
    }
    #[inline(always)] 
    pub fn forget(mut self)
    {
	unsafe {
	    ManuallyDrop::drop(&mut self.0);
	}
	mem::forget(self);
    }
}

impl<'o, T> ops::Drop for DeferredRef<'o, T>
{
    fn drop(&mut self) {
	self.1.0.write().push(AtomicRefCell::new( unsafe { ManuallyDrop::take(&mut self.0) } ));
    }
}

impl<'o, T> ops::Deref for DeferredRef<'o, T>
{
    type Target = T;
    #[inline(always)] 
    fn deref(&self) -> &Self::Target {
	self.0.deref()
    }
}

impl<'o, T> ops::DerefMut for DeferredRef<'o, T>
{
    #[inline(always)] 
    fn deref_mut(&mut self) -> &mut Self::Target {
	self.0.deref_mut()
    }
}

impl<T: Send+ Sync + Default> Pool<T>
{
    pub fn new(size: usize) -> Self
    {
	Self(RwLock::new(iter::repeat_with(T::default).map(AtomicRefCell::new).take(size).collect()))
    }

    #[inline] 
    pub fn take_one(&self) -> PoolRef<'_, T>
    {
	self.take_one_or_else(T::default)
    }
    #[inline] 
    pub fn take_one_or_default(&self) -> AtomicRefMut<'_, T>
    {
	self.take_one_or_insert_with(T::default)
    }
}

impl<T: Send+Sync> iter::FromIterator<T> for Pool<T>
{
    fn from_iter<I: IntoIterator<Item=T>>(iter: I) -> Self
    {
	Self(RwLock::new(iter.into_iter().map(AtomicRefCell::new).collect()))
    }
}


impl<T: Send+Sync> Pool<T>
{
    #[inline] 
    pub fn take_one_or(&self, value: T) -> PoolRef<'_, T>
    {
	self.take_one_or_else(move || value)
    }
    pub fn try_take_one(&self) -> Option<AtomicRefMut<'_, T>>
    {
	for possible in self.0.read().iter() {
	    possible.try_borrow_mut().ok()?;
	}
	None   
    }

    pub fn take_one_or_else<F>(&self, new: F) -> PoolRef<'_, T>
    where F: FnOnce() -> T
    {
	let try_borrow = {
	    let this = &self;
	    move || -> Option<AtomicRefMut<T>> {
		for possible in this.0.read().iter() {
		    possible.try_borrow_mut().ok()?;
		}
		None
	    }
	};
	if let Some(v) = try_borrow() {
	    PoolRef::Inside(v)
	} else {
	    PoolRef::Outside(DeferredRef(ManuallyDrop::new(new()), self))
	}
    }

    #[inline] 
    pub fn new_with_value(size: usize, value: T) -> Self
    where T: Clone
    {
	iter::repeat(value).take(size).collect()
    }
    pub fn new_with<F>(size: usize, with: F) -> Self
    where F: FnMut() -> T
    {
	Self(RwLock::new(iter::repeat_with(with).map(AtomicRefCell::new).take(size).collect()))
    }

    #[inline] 
    pub const fn new_empty() -> Self
    {
	Self(parking_lot::const_rwlock(Vec::new()))
    }

    pub fn take_one_or_insert_with<F>(&self, with: F) -> AtomicRefMut<'_, T>
    where F: FnOnce() -> T
    {
	let try_borrow = {
	    let this = &self;
	    move || -> Result<AtomicRefMut<T>, parking_lot::RwLockWriteGuard<_>> {
		let reader = this.0.upgradable_read();
		for possible in reader.iter() {
		    if let Ok(val) = possible.try_borrow_mut() {
			return Ok(val);
		    }
		}
		Err(parking_lot::lock_api::RwLockUpgradableReadGuard::upgrade(reader))
	    }
	};
	match try_borrow() {
	    Ok(val) => val,
	    Err(writer) => {
		let l = writer.len();
		writer.push(AtomicRefCell::new(with()));
		parking_lot::lock_api::RwLockWriteGuard::downgrade(writer)[l].borrow_mut()
	    },
	}
    }
    #[inline] 
    pub fn take_one_or_insert_value(&self, value: T) -> AtomicRefMut<'_, T>
    {
	self.take_one_or_insert_with(move || value)
    }
}