use std::{cell, marker, mem, num::NonZero, ptr, vec};
use dope_sys::thread;
use o3::alloc::{self, VecExt};
use crate::Errno;
mod raw;
union Value<T, O> {
task: mem::ManuallyDrop<T>,
output: mem::ManuallyDrop<O>,
}
pub trait Task: Send {
type Output: Send;
fn run(self) -> Self::Output;
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
Attributes(Errno),
Create(Errno),
}
pub struct Builder<T, O, P> {
group: raw::Group<T, O, P>,
}
impl<T, O, P> Builder<T, O, P> {
pub fn try_new(capacity: usize) -> Result<Self, alloc::Error> {
Ok(Self {
group: raw::Group {
slots: Vec::try_new(capacity)?,
started: 0,
},
})
}
pub fn try_push(&mut self, task: T, parent: P) -> Result<(), alloc::Error> {
alloc::Vacant::try_new(&mut self.group.slots)?.insert(Slot {
value: cell::UnsafeCell::new(Value {
task: mem::ManuallyDrop::new(task),
}),
thread: cell::Cell::new(mem::MaybeUninit::uninit()),
parent,
});
Ok(())
}
pub fn scope<R>(self, run: impl for<'scope> FnOnce(Workers<'scope, T, O, P>) -> R) -> R {
let mut group = self.group;
run(Workers {
group: &mut group,
owner: marker::PhantomData,
})
}
}
pub struct Slot<T, O, P> {
value: cell::UnsafeCell<Value<T, O>>,
thread: cell::Cell<mem::MaybeUninit<thread::Handle>>,
parent: P,
}
impl<T, O, P> Slot<T, O, P> {
pub fn parent(&self) -> &P {
&self.parent
}
}
pub struct Workers<'scope, T, O, P> {
group: &'scope mut raw::Group<T, O, P>,
owner: marker::PhantomData<fn(&'scope ()) -> &'scope ()>,
}
impl<T, O, P> Workers<'_, T, O, P> {
pub fn entries(&self) -> &[Slot<T, O, P>] {
&self.group.slots
}
pub fn parents_mut(&mut self) -> impl ExactSizeIterator<Item = &mut P> {
self.group.parents_mut()
}
pub fn finish(self) -> Finished<T, O, P> {
self.group.take()
}
}
pub struct Finished<T, O, P> {
slots: vec::IntoIter<Slot<T, O, P>>,
}
impl<T, O, P> Workers<'_, T, O, P> {
pub fn start(&mut self, stack: NonZero<usize>) -> Result<(), (&P, Error)>
where
T: Task<Output = O>,
{
let pending = unsafe { self.group.slots.get_unchecked(self.group.started..) };
let Some(first) = pending.first() else {
return Ok(());
};
let attr = thread::Attrs::with_stack(stack)
.map_err(|error| (&first.parent, Error::Attributes(error)))?;
for slot in pending {
let entry = unsafe { thread::raw::Entry::new(ptr::NonNull::from(slot)) };
let handle = entry
.spawn(&attr)
.map_err(|error| (&slot.parent, Error::Create(error)))?;
slot.thread.set(mem::MaybeUninit::new(handle));
self.group.started += 1;
}
Ok(())
}
}
impl<T, O, P> Iterator for Finished<T, O, P> {
type Item = (P, O);
fn next(&mut self) -> Option<Self::Item> {
let slot = self.slots.next()?;
let Slot { value, parent, .. } = slot;
Some((
parent,
mem::ManuallyDrop::into_inner(unsafe { value.into_inner().output }),
))
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.slots.size_hint()
}
}
impl<T, O, P> ExactSizeIterator for Finished<T, O, P> {}
impl<T, O, P> Drop for Finished<T, O, P> {
fn drop(&mut self) {
for slot in self.slots.as_mut_slice() {
unsafe { mem::ManuallyDrop::drop(&mut slot.value.get_mut().output) };
}
}
}