parapop/src/lib.rs

#![allow(dead_code)]

use std::sync::atomic::{
    self,
    AtomicBool,
    AtomicUsize,
};
use std::sync::Arc;
use std::cmp::Ordering;
use std::mem::{self, MaybeUninit};
use std::cell::UnsafeCell;
use std::ops::Drop;

mod private
{
    pub(crate) trait Sealed{}
}

mod ext; use ext::*;

pub mod iter;
mod refer; pub use refer::*;

/// A parallel, atomic populator of items
#[derive(Debug)]
pub struct Populator<'a, T: 'a>
{
    values: UnsafeCell<Box<[MaybeUninit<T>]>>,
    pub(crate) populates: Box<[AtomicBool]>, // 
    populated: AtomicUsize, // number of populated items

    _lt: PhantomLifetime<'a>,
}

impl<'a, T: 'a> Populator<'a, T>
{
    #[inline(always)]
    fn values_mut(&mut self) -> &mut [MaybeUninit<T>]
    {
	self.values.get_mut()
    }
    
    #[inline(always)]
    fn values_ref(&self) -> &[MaybeUninit<T>]
    {
	let ptr = self.values.get() as *const Box<[_]>;
	unsafe {
	    &(*ptr)[..]
	}
    }

    #[inline(always)]
    fn get_mut_ptr(&self, idx: usize) -> *mut MaybeUninit<T>
    {
	let ptr = self.values.get();
	unsafe {
	    &mut (*ptr)[idx] as *mut _
	}
    }
}

impl<'a, T> Drop for Populator<'a, T>
{

    fn drop(&mut self)
    {
	if mem::needs_drop::<T>() {
	    let len = self.values_ref().len();
	    if *self.populated.get_mut() == len {
		// Fully populated, drop whole slice in place
		unsafe {
		    std::ptr::drop_in_place( self.values_mut() as *mut [MaybeUninit<T>] as *mut [T])
		}
	    } else if len > 0 { // If values is 0, then that means `[try_]complete()` has been called.
		// Partially populated, drop individual parts
		for value in self.values.get_mut().iter_mut()
		    .zip(self.populates.iter()
			 .map(|x| x.load(atomic::Ordering::Acquire)))
		    .filter_map(|(v, prod)|
				prod.then(move ||
					  v.as_mut_ptr()))
		{
		    unsafe {
			std::ptr::drop_in_place(value)
		    }
		}
	    }
	}
	// Both boxes will be dealloced after this, the values are dropped.
    }
}

unsafe impl<'a, T: 'a> Send for Populator<'a, T> where Box<T>: Send {}
unsafe impl<'a, T: 'a> Sync for Populator<'a, T>{} // Populator is always sync

//TODO: Maybe add methods with Arc<Self> receivors?
impl<'a, T> Populator<'a, T>
{
    /// Checks if an item exists at this index exclusively.
    ///
    /// Since this is an exclusive reference, no atomic operations are performed.
    #[inline]
    pub fn exists_exclusive(&mut self, idx: usize) -> bool
    {
	*self.populates[idx].get_mut()
    }
    /// Checks if an item exists currently at this index.
    #[inline]
    pub fn exists(&self, idx: usize) -> bool
    {
	self.populates[idx].load(atomic::Ordering::SeqCst)
    }
    /// How many items are populated
    ///
    /// Faster access as this is an exclusive reference and no atomic operations are needed
    #[inline]
    pub fn populated_exclusive(&mut self) -> usize
    {
	*self.populated.get_mut()
    }
    #[inline]
    /// How many items are populated
    pub fn populated(&self) -> usize
    {
	self.populated.load(atomic::Ordering::Acquire)
    }
    
    /// Faster fullness check for when this instance has no other references
    #[inline]
    pub fn is_full_exclusive(&mut self) -> bool {
	*self.populated.get_mut() == self.len()
    }
    /// Is the populator full?
    #[inline]
    pub fn is_full(&self) -> bool
    {
	self.populated() == self.len()
    }
    /// Number of items held by the populator
    #[inline]
    pub fn len(&self) -> usize
    {
	self.populates.len()
    }
    
    /// Create a new, empty populator with this size
    pub fn new(size: usize) -> Self
    {
	Self {
	    // SAFETY: MaybeUninit is not Copy, so instead we allocate the space for uninitialised memory and then .set_len(). 
	    values: UnsafeCell::new(unsafe {
		let mut uninit = Vec::with_capacity(size);
		uninit.set_len(size);
		uninit
	    }.into_boxed_slice()),
	    populates: std::iter::repeat_with(|| false.into()).take(size).collect(),
	    populated: 0usize.into(),

	    _lt: PhantomLifetime::new(),
	}
    }
    
    /// Try to insert `value` at `idx` from an exclusive reference.
    ///
    /// If `idx` already has a value, then `Err(value)` is returned, otherwise, `value` is inserted into the table and a mutable reference to it is returned.
    ///
    /// # Exclusive accesses
    /// Since this is an `&mut self` receiver, no atomic operations are required since there are no other threads with a reference to the current populator.
    pub fn try_insert_exclusive(&mut self, idx: usize, value: T) -> Result<&mut T, T>
    {
	match self.populates[idx].get_mut() {
	    &mut true => return Err(value),
	    none => *none = true,
	};
	if cfg!(debug_assertions) {
	    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
		let mref = &mut self.values.get_mut()[idx];
		*mref = MaybeUninit::new(value);
		unsafe {
		    mref.assume_init_mut()
		}
	    })) {
		Err(p) => {
		    *self.populates[idx].get_mut() = false;
		    std::panic::resume_unwind(p)
		},
		Ok(mref) => Ok(mref),
	    }
	} else {
	    let mref = &mut self.values.get_mut()[idx];
	    *mref = MaybeUninit::new(value);
	    Ok(unsafe {
		mref.assume_init_mut()
	    })
	}
    }
    
    /// Try to insert `value` at `idx`.
    ///
    /// If `idx` already has a value, then `Err(value)` is returned, otherwise, `value` is inserted into the table and the number of items now populated is returned.
    pub fn try_insert(&self, idx: usize, value: T) -> Result<usize, T>
    {
	//TODO: XXX: Should we use SeqCst -> Acquire, or Acquire -> Relaxed?
	if let Ok(false) = self.populates[idx].compare_exchange(false, true, atomic::Ordering::SeqCst, atomic::Ordering::Acquire) {
	    // The value at idx hasn't been set
	    
	    if cfg!(debug_assertions) {
		match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
		    let ptr = self.get_mut_ptr(idx);  //self.values[idx].get();
		    unsafe {
			*ptr = MaybeUninit::new(value);
		    }
		})) {
		    Err(p) => {
			self.populates[idx].store(false, atomic::Ordering::SeqCst); // Re-set to unoccupied
			std::panic::resume_unwind(p)
		    },
		    Ok(_) => (),
		}
	    } else {
		// SAFETY: This operation will never panic, since `values` and `populates` are always the same size
		// SAFETY: We have already ensured that `values[idx]` does not contain a value.
		unsafe {
		    *self.get_mut_ptr(idx) = MaybeUninit::new(value);
		}
	    }
	    // Value is inserted, increment `populated`
	    Ok(self.populated.fetch_add(1, atomic::Ordering::SeqCst) + 1)
	} else {
	    Err(value)
	}
    }
    
    #[inline(always)]
    fn bounds_okay(&self, idx: usize) -> bool
    {
	idx < self.populates.len()
    }
    
    #[inline(always)]
    fn bounds_check(&self, idx: usize)
    {
	#[inline(never)]
	#[cold]
	fn _panic_oob(idx: usize, len: usize) -> !
	{
	    panic!("Cannot reference slot at index {} (length is {})", idx,len)
	}
	if idx >= self.populates.len() {
	    _panic_oob(idx, self.populates.len())
	}
    }
    
    /// Get a reference to an item at `idx` whether it exists or not.
    ///
    /// # Panics
    /// If the index is out of bounds
    #[inline]
    pub fn get_ref(&self, idx: usize) -> Ref<'_, 'a, T>
    {
	self.bounds_check(idx);
	Ref {
	    pop: self,
	    inserted: &self.populates[idx],
	    idx
	}
    }

    /// Consume an atomic reference-counted instance of a populator array into a reference to a specific index in the populator, whether it exists or not.
    ///
    /// # Panics
    /// If the index is out of bounds
    #[inline]
    pub fn into_ref(self: Arc<Self>, idx: usize) -> OwnedRef<'a, T>
    {
	self.bounds_check(idx);
	OwnedRef {
	    pop: self,
	    idx
	}
    }

    /// Get an exclusive reference for the item at `idx` whether it exists or not.
    ///
    /// # Panics
    /// If the index is out of bounds
    ///
    /// # Exclusivity
    /// No atomic operations are performed on the returned reference, since as long as it exists, no other reference to this instance can exist.
    #[inline] 
    pub fn get_ref_exclusive(&mut self, idx: usize) -> RefEx<'_, 'a, T>
    {
	self.bounds_check(idx);
	RefEx {
	    pop: self,
	    idx,
	}
    }

    /// Try to get an exclusive, mutable reference to an item at `idx` if an item exists there.
    ///
    /// No atomic operations are performed since this is an exclusive reference.
    #[inline]
    pub fn try_get_exclusive_mut(&mut self, idx: usize) -> Option<&mut T>
    {
	if *self.populates[idx].get_mut() {
	    Some(unsafe{ self.values.get_mut()[idx].assume_init_mut() })
	} else {
	    None
	}
    }

    /// Try to get an exclusive, mutable reference to an item at `idx` if an item exists there.
    ///
    /// No atomic operations are performed since this is an exclusive reference.
    #[inline]
    pub fn try_get_exclusive(&mut self, idx: usize) -> Option<&T>
    {
	self.try_get_exclusive_mut(idx).map(|&mut ref a| a)
    }

    /// Insert `value` into `idx`.
    ///
    /// # Panics
    /// If `idx` already has a value inserted.
    #[inline]
    pub fn insert(&self, idx: usize, value: T) -> usize
    {
	#[inline(never)]
	#[cold]
	fn panic_inserted(i: usize) -> !
	{
	    panic!("There is already a value at {}", i)
	}

	match self.try_insert(idx, value) {
	    Ok(v) => v,
	    Err(_) => panic_inserted(idx),
	}
    }
    
    /// Insert `value` into `idx` through an exclusive reference.
    ///
    /// # Panics
    /// If `idx` already has a value inserted.
    ///
    /// # Exclusivity
    /// No atomic operations are required since the `&mut self` receiver guarantees no other references to this instance currently exist.
    #[inline]
    pub fn insert_exclusive(&mut self, idx: usize, value: T) -> &mut T
    {
	#[inline(never)]
	#[cold]
	fn panic_inserted(i: usize) -> !
	{
	    panic!("There is already a value at {}", i)
	}

	match self.try_insert_exclusive(idx, value) {
	    Ok(v) => v,
	    Err(_) => panic_inserted(idx),
	}
    }
    
    #[inline(always)]
    fn take_all(&mut self) -> (Box<[MaybeUninit<T>]>, Box<[AtomicBool]>)
    {
	let inner = self.values.get_mut();
	(mem::replace(inner, vec![].into_boxed_slice()),
	 mem::replace(&mut self.populates, vec![].into_boxed_slice()))
    }

    #[inline(always)]
    fn take_values(&mut self) -> Box<[MaybeUninit<T>]>
    {
	let inner = self.values.get_mut();
	mem::replace(inner, vec![].into_boxed_slice())
    }

    
    
    /// If all values are populated, then convert it into a boxed slice and return it.
    pub fn try_complete(mut self) -> Result<Box<[T]>, Self>
    {
	if *self.populated.get_mut() == self.len() {
	    //let ptr = Box::into_raw(std::mem::replace(&mut self.values, UnsafeCell::new(vec![].into_boxed_slice())).into_inner());
	    let ptr = {
		let inner = self.values.get_mut();
		Box::into_raw(mem::replace(inner, vec![].into_boxed_slice()))
	    };
	    Ok(unsafe {
		Box::from_raw(ptr as *mut [T])
	    })
	} else {
	    Err(self)
	}
    }

    
    /// If all values are populated and the `Arc` has no other strong references, then convert it into a boxed slice and return it.
    ///
    /// Otherwise, return the arc.
    ///
    /// # If the values are not all populated
    /// But the `Arc` is empty, then a new, single strong reference arc is constructed around the value and returned as its `Err`.
    pub fn try_complete_owned(self: Arc<Self>) -> Result<Box<[T]>, Arc<Self>>
    {
	match Arc::try_unwrap(self) {
	    Ok(extracted) => extracted.try_complete().map_err(Arc::new),
	    Err(e) => Err(e),
	}
    }

    /// If all values are populated, returns a slice of all the elements.
    ///
    /// Performs a single atomic load of the number of currently inserted elements to check for completion
    pub fn try_completed_ref(&self) -> Option<&[T]>
    {
	if self.populated.load(atomic::Ordering::SeqCst) == self.len() {
	    Some(unsafe { iter::assume_init_ref(&*self.values.get()) })
	} else {
	    None
	}
    }

    /// Returns a slice of all the elements
    ///
    /// # Panics
    /// If the collection was not full.
    #[inline] 
    pub fn completed_ref(&self) -> &[T]
    {
	self.try_completed_ref().expect("Collection was not fully populated")
    }
    
    /// Returns a mutable slice of all the elements
    ///
    /// # Panics
    /// If the collection was not full.
    #[inline] 
    pub fn completed_mut(&mut self) -> &mut [T]
    {
	self.try_completed_mut().expect("Collection was not fully populated")
    }
    
    /// If all values are populated, returns a mutable slice of all the elements.
    ///
    /// # Exclusivity
    /// No atomic operations are required since this exclusive reference to `self` ensures there are no other references to this instance.
    pub fn try_completed_mut(&mut self) -> Option<&mut [T]>
    {
	if *self.populated.get_mut() == self.len() {
	    Some(unsafe { iter::assume_init_mut(self.values.get_mut()) })
	} else {
	    None
	}
    }
    /// Returns the completed population.
    ///
    /// # Panics
    /// If the collection is not fully populated.
    #[inline] // Maybe?
    pub fn complete(self) -> Box<[T]>
    {
	#[inline(never)]
	#[cold]
	fn panic_uncomplete() -> !
	{
	    panic!("Not all values had been populated")
	}
	
	match self.try_complete() {
	    Ok(v) => v,
	    Err(_) => panic_uncomplete(),
	}
    }

    /// Returns the completed population from an `Arc` that has no more than 1 reference.
    ///
    /// # Panics
    /// * If the `Arc` has more than one strong reference
    /// * If the collection is not fully populated
    pub fn complete_owned(self: Arc<Self>) -> Box<[T]>
    {
	#[inline(never)]
	#[cold]
	fn panic_uncomplete_or_shared<T>(values: &Arc<T>) -> !
	{
	    let sc = Arc::strong_count(values);
	    if sc > 1 {
		panic!("More than one ({}) reference to the `Arc` holding this instance", sc)
	    } else  {
		panic!("Not all values had been populated")
	    }
	}
	match self.try_complete_owned() {
	    Ok(v) => v,
	    Err(ref e) => panic_uncomplete_or_shared(e),
	}
    }

    /// Create an iterator over references to a completed population if it is completed.
    #[inline] 
    pub fn try_completed_iter(&self) -> Option<iter::FullIterRef<'a, '_, T>>
    {
	self.try_completed_ref().map(iter::FullIterRef::new)
    }
    
    /// Create an iterator over references to a completed population.
    ///
    /// # Panics
    /// If the collection is not fully populated
    #[inline] 
    pub fn completed_iter(&self) -> iter::FullIterRef<'a, '_, T>
    {
	iter::FullIterRef::new(self.completed_ref())
    }

    
    /// Create a mutable iterator over references to a completed population if it is completed.
    #[inline] 
    pub fn try_completed_iter_mut(&mut self) -> Option<iter::FullIterMut<'a, '_, T>>
    {
	self.try_completed_mut().map(iter::FullIterMut::new)
    }
    
    /// Create a mutable iterator over references to a completed population.
    ///
    /// # Panics
    /// If the collection is not fully populated
    #[inline] 
    pub fn completed_iter_mut(&mut self) -> iter::FullIterMut<'a, '_, T>
    {
	iter::FullIterMut::new(self.completed_mut())
    }

    /// Create an iterator of references to this `Populator<'a, T>`
    #[inline]
    pub fn iter(&self) -> iter::Iter<'a, '_, T>
    {
	iter::Iter::new(self)
    }

    /// Create an iterator of references to a slice of this `Populator<'a, T>`
    ///
    /// # Panics
    /// if `range.end` is larger than `len()`.
    #[inline]
    pub fn iter_slice(&self, range: impl Into<std::ops::Range<usize>>) -> iter::Iter<'a, '_, T>
    {
	let range = range.into();
	if range.end > self.len() {
	    #[inline(never)] 
	    #[cold]
	    fn _panic_oob(end: usize, idx: usize) -> !
	    {
		panic!("Range ends out of the bounds of the populator (length is {end}, range end is {idx})")
	    }
	    _panic_oob(self.len(), range.end);
	}
	
	iter::Iter::new_range(self, range)
    }
    //TODO: `self: Arc<Self>` version of iter_slice()/iter()
}

impl<'a, T: 'a> FromIterator<Option<T>> for Populator<'a, T>
{
    fn from_iter<I: IntoIterator<Item = Option<T>>>(iter: I) -> Self {
	let mut v =0usize;
	let (items, bools) : (Vec<_>, Vec<AtomicBool>) = iter.into_iter()
	    .map(|x| x.map(|item| { v +=1; (MaybeUninit::new(item), true.into()) })
		 .unwrap_or((MaybeUninit::uninit(),false.into())))
	    .unzip();
	
	debug_assert_eq!(items.len(), bools.len(), "invalid ");

	Self {
	    populated: v.into(),
	    values: UnsafeCell::new(items.into_boxed_slice()),
	    populates: bools.into_boxed_slice(),
	    
	    _lt: PhantomLifetime::new(),
	}
    }
}

impl<'a, T: 'a> IntoIterator for Populator<'a, T>
{
    type Item = T;
    type IntoIter = iter::IntoIter<'a, T>;
    
    #[inline]
    fn into_iter(self) -> Self::IntoIter {
	iter::IntoIter::create_from(self)
    }
}

#[cfg(test)]
mod tests;