takeaway 0.1.4

//! Enqueuing tasks.

use alloc::{boxed::Box, vec::Vec};
use core::{
    alloc::Layout,
    cell::UnsafeCell,
    mem::MaybeUninit,
    ptr::{self, NonNull},
    slice,
    task::Waker,
};

use crate::Task;

//----------- Enqueuer ---------------------------------------------------------

/// A handle to enqueue tasks for a [`Worker`].
///
/// [`Worker::enqueue()`] / [`enqueue_one()`] is the most convenient way to
/// enqueue new tasks, but can be troublesome if the [`Worker`] is not directly
/// available or borrowed elsewhere.  An [`Enqueuer`] is a dedicated handle to
/// the [`Worker`] for enqueueing new tasks, and can help work around such
/// issues.  Most notably, it can be used while the worker is asleep.  It can be
/// accessed via [`Worker::enqueuer()`] (which returns `&Rc<Enqueuer<T>>`).
///
/// [`Worker`]: crate::Worker
/// [`Worker::enqueue()`]: crate::Worker::enqueue()
/// [`enqueue_one()`]: crate::Worker::enqueue_one()
/// [`Worker::enqueuer()`]: crate::Worker::enqueuer()
pub struct Enqueuer<T: Task> {
    inner: UnsafeCell<Inner<T>>,
}

struct Inner<T: Task> {
    /// The task counter.
    ///
    /// This is a very compact state variable for enqueueing new tasks.  It is
    /// optimized for the hot path of adding new tasks while not borrowed.
    ///
    /// ## States
    ///
    /// - Normal (`v <= -2 * esz`): The buffer has space for `-v - esz` more
    ///   bytes (`-v / esz - 1` more tasks), and is not borrowed.  The next
    ///   task should be written to `end + v + esz`.
    ///
    /// - Full (`v = -esz`, no waker): The buffer is out of space, but is not
    ///   borrowed.  It must be reallocated to add new tasks.
    ///
    /// - Asleep (`v = -esz`, waker set): The buffer was recently cleared, and a
    ///   waker has been set.  It is not borrowed, and a normal state can be set
    ///   once the waker is used.
    ///
    /// - Locked (`v = 0`): The buffer has been locked by another actor.  This
    ///   is usually grounds for panicking.
    ///
    /// The value is always a multiple of the element size, and is always zero
    /// or negative.  If the element size is 0, `esz` is treated as 1.
    ///
    /// There was a choice of positive or negative states, and of counting up
    /// or counting down.  The counter uses negative states and counts up so
    /// that tasks can be appended to the buffer instead of prepended, and the
    /// counter (with a single addition) determines the address to write to.
    ///
    /// This is especially well-suited to architectures with flag registers,
    /// where the state transition for enqueueing new tasks also immediately
    /// indicates whether it is full or locked.
    counter: isize,

    /// The buffer of tasks.
    end: NonNull<T>,

    /// The capacity of the buffer, in bytes.
    capacity: isize,

    /// A waker for the worker, if it is asleep.
    waker: Option<Waker>,
}

impl<T: Task> Enqueuer<T> {
    /// Construct a new [`Enqueuer`].
    pub(crate) const fn new() -> Self {
        let capacity = if Self::ZST { isize::MAX } else { 0 };
        let counter = -capacity - Self::ESZ;
        Self {
            inner: UnsafeCell::new(Inner {
                counter,
                end: NonNull::dangling(),
                capacity,
                waker: None,
            }),
        }
    }

    /// Set a waker in this [`Enqueuer`].
    ///
    /// # Panics
    ///
    /// Panics if:
    /// - The enqueuer is already borrowed.
    /// - The enqueuer has any tasks.
    pub(crate) fn set_waker(&self, waker: Waker) {
        match self.waker() {
            Some(slot) => {
                // State: Asleep.
                *slot = waker;
            }

            slot @ None => {
                // State: Normal / Empty / Full / Locked.
                assert!(!self.is_locked(), "the enqueuer is locked");

                // State: Normal / Empty / Full.
                let counter = self.counter();
                let empty = -self.byte_capacity() - Self::ESZ;
                assert!(*counter == empty, "the enqueuer is not empty");

                // State: Empty.
                *slot = Some(waker);
                *counter = -Self::ESZ;
            }
        }
    }

    /// Take the waker from this [`Enqueuer`].
    pub(crate) fn take_waker(&self) -> Option<Waker> {
        let waker = self.waker().take()?;
        *self.counter() = -self.byte_capacity() - Self::ESZ;
        Some(waker)
    }

    /// Enqueue a task.
    ///
    /// If the worker is asleep, it will be woken up.
    ///
    /// # Panics
    ///
    /// Panics if this is called while the enqueuer is borrowed, e.g. it is
    /// within a call to [`Enqueuer::extend()`].
    pub fn add(&self, task: T) {
        let inner = self.inner.get();
        let counter = unsafe { (*inner).counter };
        let next = counter + Self::ESZ;

        if next < 0 {
            // Current state: Normal.
            let end = unsafe { (*inner).end };
            if !Self::ZST {
                let ptr = unsafe { end.byte_offset(next) };
                unsafe { ptr.write(task) };
            }
            unsafe { (*inner).counter = next };
            return;
        }

        self.add_slow(task, next)
    }

    #[cold]
    fn add_slow(&self, task: T, next: isize) {
        let inner = self.inner.get();

        // Current state: Full / Asleep / Borrowed.
        assert!(next == 0, "the enqueuer is already borrowed");

        // Current state: Full / Asleep.
        let capacity = unsafe { (*inner).capacity };
        if capacity == 0 {
            // Current state: Empty / Full / Asleep.
            cold_path();
            debug_assert!(!Self::ZST);

            let mut buffer = Box::new_uninit_slice(16);
            buffer[0].write(task);
            let ptr = Box::into_raw(buffer);
            let end = unsafe { ptr.cast::<T>().add(16) };
            let end = unsafe { NonNull::new_unchecked(end) };
            unsafe { (*inner).end = end };
            unsafe { (*inner).capacity = 16 * Self::ESZ };
            unsafe { (*inner).counter = -16 * Self::ESZ };

            if let Some(waker) = unsafe { (*inner).waker.take() } {
                waker.wake();
            }

            return;
        }

        // Current state: Full / Asleep.
        let end = unsafe { (*inner).end };
        let ptr = unsafe { end.byte_offset(-capacity) };
        if unsafe { (*inner).waker.is_some() } {
            // Current state: Asleep.
            let waker = unsafe { (*inner).waker.take() }.unwrap();
            unsafe { ptr.write(task) };
            unsafe { (*inner).counter = -capacity };
            waker.wake();
        } else {
            // Current state: Full.
            let mut vec = unsafe {
                Vec::from_raw_parts(
                    ptr.as_ptr(),
                    (capacity / Self::ESZ) as usize,
                    (capacity / Self::ESZ) as usize,
                )
            };
            vec.push(task);
            let old_cap = capacity;
            let capacity = vec.capacity() as isize * Self::ESZ;
            let ptr = unsafe { NonNull::new_unchecked(vec.as_mut_ptr()) };
            let end = unsafe { ptr.byte_offset(capacity) };
            core::mem::forget(vec);
            unsafe { (*inner).end = end };
            unsafe { (*inner).capacity = capacity };
            unsafe { (*inner).counter = -capacity + old_cap };
        }
    }

    /// Enqueue a set of tasks.
    ///
    /// If the worker is asleep, it will be woken up.
    ///
    /// # Panics
    ///
    /// Panics if this is called while the enqueuer is borrowed, e.g. it is
    /// within a(nother) call to [`Enqueuer::extend()`].
    pub fn extend(&self, tasks: impl IntoIterator<Item = T>) {
        // TODO
        for task in tasks {
            self.add(task);
        }
    }

    /// Drain the enqueued tasks.
    ///
    /// Returns the number of freshly enqueued tasks.
    ///
    /// # Panics
    ///
    /// Panics if this is called while the enqueuer is borrowed, e.g. it is
    /// within a call to [`Enqueuer::extend()`].
    pub(crate) fn drain(&self, tasks: &mut Vec<T>) -> usize {
        // Determine how many tasks are in the enqueuer.
        let inner = self.inner.get();
        let counter = unsafe { (*inner).counter };
        let capacity = unsafe { (*inner).capacity };
        assert!(counter < 0, "the enqueuer is already borrowed");

        let end = unsafe { (*inner).end };
        let ptr = unsafe { end.byte_offset(-capacity) };
        let ptr = ptr.cast::<MaybeUninit<T>>().as_ptr();
        let len = ((capacity + counter) / Self::ESZ + 1) as usize;

        tasks.reserve(len);
        let space = tasks.spare_capacity_mut().as_mut_ptr();
        unsafe { ptr::copy_nonoverlapping(ptr, space, len) };
        unsafe { tasks.set_len(tasks.len() + len) };

        unsafe { (*inner).counter = -capacity - Self::ESZ };

        len
    }
}

impl<T: Task> Extend<T> for &Enqueuer<T> {
    /// Enqueue a set of tasks.
    ///
    /// This is equivalent to (and is defined as) [`Enqueuer::extend()`].
    ///
    /// If the worker is asleep, it will be woken up.
    ///
    /// # Panics
    ///
    /// Panics if this is called while the enqueuer is borrowed, e.g. it is
    /// within a(nother) call to [`Enqueuer::extend()`].
    fn extend<I: IntoIterator<Item = T>>(&mut self, tasks: I) {
        (*self).extend(tasks)
    }
}

impl<T: Task> Drop for Enqueuer<T> {
    fn drop(&mut self) {
        if self.is_locked() {
            // State: Locked.
            return;
        }

        // Drop any live tasks.
        if self.waker().is_none() {
            // State: Normal or Full.

            // SAFETY: State is neither Asleep nor Locked.
            for task in unsafe { self.buffer_init() } {
                // SAFETY: All tasks in 'buffer_init()' are initialized.
                unsafe { task.assume_init_drop() };
            }
        }

        // De-allocate the buffer itself.
        if !Self::ZST && self.byte_capacity() != 0 {
            // SAFETY: State is not Locked.
            let buffer = unsafe { self.buffer() };

            // SAFETY: Such an allocation exists, so the layout must be valid.
            let layout =
                unsafe { Layout::array::<T>(buffer.len()).unwrap_unchecked() };

            // SAFETY: 'buffer' was allocated with this layout.
            unsafe {
                alloc::alloc::dealloc(buffer.as_mut_ptr().cast(), layout)
            };
        }
    }
}

impl<T: Task> Enqueuer<T> {
    /// Whether the task is zero-sized.
    const ZST: bool = size_of::<T>() == 0;

    /// The element size.
    const ESZ: isize = if !Self::ZST {
        size_of::<T>() as isize
    } else {
        1
    };

    /// Whether the state is Locked.
    fn is_locked(&self) -> bool {
        unsafe { (*self.inner.get()).counter == 0 }
    }

    /// The capacity, in bytes.
    fn byte_capacity(&self) -> isize {
        unsafe { (*self.inner.get()).capacity }
    }

    /// The enqueued task length, in bytes.
    ///
    /// ## Safety
    ///
    /// - State must not be Asleep or Locked.
    fn byte_len(&self) -> isize {
        let inner = self.inner.get();
        unsafe { (*inner).capacity + (*inner).counter + Self::ESZ }
    }

    /// The length.
    ///
    /// ## Safety
    ///
    /// - State must not be Asleep or Locked.
    fn len(&self) -> usize {
        (self.byte_len() / Self::ESZ) as usize
    }

    /// The capacity.
    fn capacity(&self) -> usize {
        (self.byte_capacity() / Self::ESZ) as usize
    }

    /// The start of the buffer.
    fn ptr(&self) -> *mut T {
        let inner = self.inner.get();
        let off = if Self::ZST {
            0
        } else {
            unsafe { (*inner).capacity }
        };
        unsafe { (*inner).end.as_ptr().byte_offset(-off) }
    }

    /// The whole buffer.
    ///
    /// ## Safety
    ///
    /// - State must not be Locked.
    #[allow(clippy::mut_from_ref)] // Internal API
    unsafe fn buffer(&self) -> &mut [MaybeUninit<T>] {
        unsafe { slice::from_raw_parts_mut(self.ptr().cast(), self.capacity()) }
    }

    /// The initialized part of the buffer.
    ///
    /// ## Safety
    ///
    /// - State must not be Asleep or Locked.
    #[allow(clippy::mut_from_ref)] // Internal API
    unsafe fn buffer_init(&self) -> &mut [MaybeUninit<T>] {
        unsafe { slice::from_raw_parts_mut(self.ptr().cast(), self.len()) }
    }

    /// The waker.
    #[allow(clippy::mut_from_ref)] // Internal API
    fn waker(&self) -> &mut Option<Waker> {
        unsafe { &mut (*self.inner.get()).waker }
    }

    /// The counter.
    #[allow(clippy::mut_from_ref)] // Internal API
    fn counter(&self) -> &mut isize {
        unsafe { &mut (*self.inner.get()).counter }
    }
}

#[cold]
fn cold_path() {}

#[cfg(test)]
mod test {
    use std::vec::Vec;

    use crate::Task;

    use super::Enqueuer;

    #[derive(Copy, Clone, PartialEq, Eq, Debug)]
    struct Foo(u32);

    impl Task for Foo {
        type Priority = ();
        fn priority(&self) -> Self::Priority {}
    }

    #[test]
    fn new() {
        let _ = Enqueuer::<()>::new();
        let _ = Enqueuer::<Foo>::new();
    }

    #[test]
    fn add() {
        let enqueuer = Enqueuer::<()>::new();
        enqueuer.add(());
        enqueuer.add(());
        enqueuer.add(());
        enqueuer.add(());
        assert_eq!(enqueuer.len(), 4);

        let enqueuer = Enqueuer::<Foo>::new();
        enqueuer.add(Foo(0));
        enqueuer.add(Foo(1));
        enqueuer.add(Foo(2));
        enqueuer.add(Foo(3));
        assert_eq!(enqueuer.len(), 4);
        assert_eq!(enqueuer.byte_len(), 16);
        let mut tasks = Vec::new();
        enqueuer.drain(&mut tasks);
        assert_eq!(tasks, [Foo(0), Foo(1), Foo(2), Foo(3)]);
    }

    #[test]
    fn realloc() {
        let enqueuer = Enqueuer::<Foo>::new();
        for n in 1..4 {
            let mut expected = Vec::new();
            for i in 0..32 * n {
                let task = Foo(n * 32 + i);
                enqueuer.add(task);
                expected.push(task);
            }
            assert_eq!(enqueuer.len(), 32 * n as usize);
            assert_eq!(enqueuer.byte_len(), 32 * n as isize * 4);
            let mut tasks = Vec::new();
            enqueuer.drain(&mut tasks);
            assert_eq!(tasks, expected);
        }
    }
}