#![allow(dead_code)] use std::sync::atomic::{ self, AtomicBool, AtomicUsize, }; use std::mem::{self, MaybeUninit}; use std::cell::UnsafeCell; use std::ops::Drop; /// A parallel, atomic populator of items pub struct Populator { values: Box<[UnsafeCell>]>, populates: Box<[AtomicBool]>, // populated: AtomicUsize, // number of populated items } impl Drop for Populator { fn drop(&mut self) { if mem::needs_drop::() { if self.populated() == self.values.len() { // Fully populated, drop whole slice in place unsafe { std::ptr::drop_in_place(self.values.as_mut() as *mut [UnsafeCell>] as *mut [T]) } } else if self.values.len() > 0 { // If values is 0, then that means `[try_]complete()` has been called. // Partially populated, drop individual parts for value in self.values.iter_mut().zip(self.populates.iter().map(|x| x.load(atomic::Ordering::Acquire))).filter_map(|(v, prod)| prod.then(move || v.get_mut().as_mut_ptr())) { unsafe { std::ptr::drop_in_place(value) } } } } // Both boxes will be dealloced after this, the values are dropped. } } unsafe impl Send for Populator where Box: Send {} unsafe impl Sync for Populator{} // Populator is always sync impl Populator { /// How many items are populated pub fn populated(&self) -> usize { self.populated.load(atomic::Ordering::Acquire) } /// Is the populator full? pub fn is_full(&self) -> bool { self.populated() == self.values.len() } /// Number of items held by the populator pub fn len(&self) -> usize { self.values.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: 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(), } } /// 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 { //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.values[idx].get(); unsafe { *ptr = MaybeUninit::new(value); } })) { Err(p) => 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.values[idx].get() = MaybeUninit::new(value); } } // Value is inserted, increment `populated` Ok(self.populated.fetch_add(1, atomic::Ordering::SeqCst) + 1) } else { Err(value) } } /// Insert `value` into `idx`. /// /// # Panics /// If `idx` already has a value inserted. 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), } } /// If all values are populated, then convert it into a boxed slice and return it. pub fn try_complete(mut self) -> Result, Self> { if *self.populated.get_mut() == self.values.len() { let ptr = Box::into_raw(std::mem::replace(&mut self.values, vec![].into_boxed_slice())); Ok(unsafe { Box::from_raw(ptr as *mut [T]) }) } else { Err(self) } } /// Returns the completed population. /// /// # Panics /// If the collection is not fully populated. 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(), } } } #[cfg(test)] mod tests { #[test] fn it_works() { let result = 2 + 2; assert_eq!(result, 4); } }