//! 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(RwLock>>); #[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> { match self { Self::Inside(v) => Err(v), Self::Outside(v) => Ok(v.detach()) } } pub fn replace_with(&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 for PoolRef<'o, T> { #[inline] fn borrow(&self) -> &T { ops::Deref::deref(self) } } impl<'o, T> BorrowMut 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, &'owner Pool); 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 Pool { 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 iter::FromIterator for Pool { fn from_iter>(iter: I) -> Self { Self(RwLock::new(iter.into_iter().map(AtomicRefCell::new).collect())) } } impl Pool { #[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> { for possible in self.0.read().iter() { possible.try_borrow_mut().ok()?; } None } pub fn take_one_or_else(&self, new: F) -> PoolRef<'_, T> where F: FnOnce() -> T { let try_borrow = { let this = &self; move || -> Option> { 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(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(&self, with: F) -> AtomicRefMut<'_, T> where F: FnOnce() -> T { let try_borrow = { let this = &self; move || -> Result, 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) } }