takeaway 0.1.2

//! Public queues.
//!
//! A "pubqueue" (public queue) holds a small number of tasks owned by a worker
//! thread.  Public queues can be stolen with just a few atomic operations, via
//! a [`Stealer`].

use alloc::{boxed::Box, vec::Vec};
use atomig::Atomic;
use core::{
    cell::UnsafeCell,
    mem::{MaybeUninit, transmute},
    num::NonZeroUsize,
    sync::atomic::{AtomicU32, Ordering},
};

use crate::{
    Config, Task, TaskPriority,
    util::{Backoff, extend_vec_from_slice},
};

//----------- Priority ---------------------------------------------------------

/// An atomic priority value.
#[repr(transparent)]
pub struct Priority<T: Task> {
    /// The raw value.
    raw: Atomic<<T::Priority as TaskPriority>::Repr>,
}

impl<T: Task> Default for Priority<T> {
    #[inline]
    fn default() -> Self {
        Self {
            raw: Atomic::new(<T::Priority as TaskPriority>::pack(None)),
        }
    }
}

impl<T: Task> Priority<T> {
    /// Set this [`Priority`].
    ///
    /// A worker thread can call this on its own priority to publish a public
    /// queue it has prepared.
    ///
    /// ## Safety
    ///
    /// Worker thread `i` can call `priority[i].set(p)` when `priority[i] ==
    /// None`.  `stealer[i]` must have been initialized already.  After this,
    /// `i` will no longer own its public queue.
    #[inline]
    pub unsafe fn set(&self, value: T::Priority) {
        let pack = <T::Priority as TaskPriority>::pack;
        let unpack = <T::Priority as TaskPriority>::unpack;

        debug_assert_eq!(unpack(self.raw.load(Ordering::Relaxed)), None);

        self.raw.store(pack(Some(value)), Ordering::Release);
    }

    /// Steal this thread's own [`Priority`].
    ///
    /// The worker thread can call this to take ownership of a previously
    /// published public queue, if it is still available.
    ///
    /// ## Safety
    ///
    /// Worker thread `i` can call `priority[i].steal_self(expected)`.  If the
    /// result is `Ok(())`, `priority[i]` was `Some(expected)`, and `i` now owns
    /// the public queue at `stealer[i]`.  If the result is `Err(())`, some
    /// thread `j != i` has begun stealing from `i` (and may have finished).
    #[inline]
    pub unsafe fn steal_self(&self, expected: T::Priority) -> Result<(), ()> {
        let pack = <T::Priority as TaskPriority>::pack;
        let unpack = <T::Priority as TaskPriority>::unpack;

        // Make writes to the steal-batch visible.
        let prev = unpack(self.raw.swap(pack(None), Ordering::Acquire));

        debug_assert!(
            prev.is_none_or(|prev| prev == expected),
            "'prev' ({prev:?}) can only be empty or 'expected' ({expected:?})"
        );

        if prev.is_some() { Ok(()) } else { Err(()) }
    }

    /// Steal this [`Priority`] if it is high enough.
    ///
    /// If the priority of the worker thread is strictly greater than the
    /// specified minimum priority, it will be stolen and returned as `Some`.
    /// Otherwise, `None` is returned.
    ///
    /// ## Safety
    ///
    /// Worker thread `i` can call `priority[j].steal_if_above(min)` while
    /// `priority[i] == None`.  If the result is `Some(p)`, `priority[j]` has
    /// been set to `None`, `i` is now stealing from `j`, `i != j`, and
    /// `p > min`.
    #[inline]
    pub unsafe fn steal_if_above(
        &self,
        min: Option<T::Priority>,
    ) -> Option<T::Priority> {
        let pack = <T::Priority as TaskPriority>::pack;
        let unpack = <T::Priority as TaskPriority>::unpack;

        let result = self.raw.fetch_update(
            Ordering::Acquire,
            Ordering::Relaxed,
            |cur| {
                if unpack(cur) > min {
                    Some(pack(None))
                } else {
                    None
                }
            },
        );

        match result {
            // SAFETY: 'prev > min >= 0' thus 'prev > 0'.
            Ok(prev) => Some(unsafe { unpack(prev).unwrap_unchecked() }),
            Err(_) => None,
        }
    }
}

//----------- Stealer ----------------------------------------------------------

/// A channel for stealing a public queue.
///
/// A stealer is associated with a particular worker thread, and allows it to
/// expose its public queue globally.  Other threads can then try to steal the
/// queue.  The owning thread can extract tasks from the exposed queue, so that
/// progress is made even if no thefts are attempted.
///
/// A stealer can be in one of the following states, depending on its intrinsic
/// value and external conditions:
///
/// - Available: the stealer holds a public queue on behalf of its owning
///   thread, and other threads can steal it.  The stealer holds the ID of the
///   public queue and the (non-zero) number of tasks within it.  The owning
///   thread may take some tasks from the queue by atomically reducing that
///   number.
///
/// - Stolen: some thread (not the owning thread) stole the public queue, and
///   left a different public queue to the stealer.  The new public queue is
///   empty, and the owning thread can take ownership of it to fill it with new
///   tasks.
///
/// - Unavailable: the owning thread has taken ownership of the public queue.
///   Other threads cannot observe the stealer at this time.
///
/// The stealer acts as a bidirectional channel between the owning thread and at
/// most one thief thread.
#[repr(transparent)]
pub struct Stealer {
    /// The raw value.
    raw: AtomicU32,
}

impl Stealer {
    /// Initialize a [`Stealer`].
    ///
    /// ## Safety
    ///
    /// - `id` is valid, i.e. `id < num_workers`.
    #[inline]
    pub const unsafe fn new(id: usize, config: &Config) -> Self {
        debug_assert!(id < config.num_workers.get());

        let value = unsafe { PubQueue::new(id, config.batch_size, config) };
        Self {
            raw: AtomicU32::new(value.raw),
        }
    }

    /// Set this [`Stealer`] to a particular value.
    ///
    /// ## Safety
    ///
    /// Worker thread `i` can call `stealer[i].set(expected, new)` while
    /// `priority[i] == None`.  `stealer[i]` should have the value `expected`.
    /// `i` must own `new.id()`.
    #[inline]
    pub unsafe fn set(&self, expected: PubQueue, new: PubQueue) {
        debug_assert_eq!(
            self.raw.load(Ordering::Relaxed) % (1u32 << 31),
            expected.raw
        );

        self.raw.store(new.raw, Ordering::Relaxed);
    }

    /// Try to decrement the length of this [`Stealer`].
    ///
    /// ## Safety
    ///
    /// Worker thread `i` can call `stealer[i].try_dec(expected)`.
    /// `stealer[i]` should have the value `expected`.
    ///
    /// `expected.len()` must be strictly greater than 1.
    #[inline]
    pub unsafe fn try_dec(
        &self,
        expected: PubQueue,
        config: &Config,
    ) -> Result<PubQueue, PubQueue> {
        debug_assert!(expected.len(config).get() > 1);

        // Decrement the value in place, and find out what happened.
        let res = self.raw.fetch_sub(1, Ordering::Relaxed);
        if res == expected.raw {
            let len = expected.len(config).get() - 1;
            let len = unsafe { NonZeroUsize::new_unchecked(len) };
            Ok(unsafe { expected.with_len(len, config) })
        } else {
            Err(unsafe { PubQueue::from_raw(res - 1) })
        }
    }

    /// Steal this [`Stealer`].
    ///
    /// A worker thread owning an empty public queue can use this to swap with a
    /// victim thread and steal their public queue contents.  If the victim was
    /// asleep waiting for the theft, it will be woken up.
    ///
    /// ## Safety
    ///
    /// Worker thread `i` can call `stealer[j].steal(replacement)` while `i`
    /// owns `replacement` and `i` is stealing from `j`, for `j != i`.
    ///
    /// `replacement` is a valid queue ID, i.e. `replacement < num_workers`.
    #[inline]
    pub unsafe fn steal(
        &self,
        replacement: usize,
        config: &Config,
    ) -> PubQueue {
        debug_assert!(replacement < config.num_workers.get());

        // SAFETY: 'replacement' is a valid queue ID as per the caller.
        let replacement =
            unsafe { PubQueue::new(replacement, config.batch_size, config) };

        // NOTE:
        // - The priority value was loaded using 'Acquire', which happened after
        //   the writes of the priority queue, so the new queue's contents are
        //   visible to us already.
        // - Our writes to the replacement queue need to be visible to the
        //   victim thread so that their writes to it happen afterwards.
        let mut value = self.raw.swap(replacement.raw, Ordering::Release);

        // Wake up the owning thread if it went to sleep.
        if value >= (1u32 << 31) {
            value -= 1u32 << 31;
            atomic_wait::wake_one(&self.raw);
        }

        // SAFETY: 'value % 2^31' was initialized by 'PubQueue.raw'.
        unsafe { PubQueue::from_raw(value) }
    }

    /// Wait for a thief to finish stealing this [`Stealer`].
    ///
    /// When a worker thread discovers it is being stolen from, it should use
    /// this to wait until the theft finishes, so it has a public queue to work
    /// with.
    ///
    /// ## Safety
    ///
    /// Worker thread `i` can call `stealer[i].wait_for_theft(original)` while
    /// `original` is the last value `i` observed/wrote in `stealer[i]` and
    /// `priority[i]` was set to `None` by a thread other than `i`.
    pub unsafe fn wait_for_theft(&self, original: PubQueue) -> PubQueue {
        let backoff = Backoff::new();
        while !backoff.is_completed() {
            let raw = self.raw.load(Ordering::Acquire);
            if raw != original.raw {
                // The new value has been found, return it.
                debug_assert!(raw < (1u32 << 31));
                return unsafe { PubQueue::from_raw(raw) };
            }

            // The thief has not yet stolen the value.  Wait.
            backoff.snooze();
        }

        // After waiting a substantial amount, the thief has not yet stolen the
        // value.  Put the thread to sleep, and let the waker restore it.

        // Set 'contended' to mark the owner as going to sleep.
        let raw = self.raw.fetch_or(1u32 << 31, Ordering::Relaxed);
        if raw != original.raw {
            // The new value has been found, return it.
            debug_assert!(raw < (1u32 << 31));

            // Now that the public queue is stolen, nobody else will observe the
            // stealer.  We can freely write to it to remove the contention bit.
            self.raw.store(raw, Ordering::Relaxed);

            return unsafe { PubQueue::from_raw(raw) };
        }

        loop {
            // Go to sleep.
            let old = original.raw | (1u32 << 31);
            atomic_wait::wait(&self.raw, old);

            // Check whether the value has changed.
            let raw = self.raw.load(Ordering::Relaxed);
            if raw != old {
                // The new value has been found, return it.
                debug_assert!(raw < (1u32 << 31));
                return unsafe { PubQueue::from_raw(raw) };
            }
        }
    }
}

//----------- PubQueue ---------------------------------------------------------

/// A reference to a public queue.
///
/// This combines the ID of a public queue with its length.  This is the value
/// stored in a [`Stealer`], depending on its state.
#[derive(Copy, Clone)]
#[repr(transparent)]
pub struct PubQueue {
    /// The raw queue value.
    ///
    /// This is always strictly less than `2^31`.
    raw: u32,
}

impl PubQueue {
    /// Construct a new [`PubQueue`].
    ///
    /// ## Safety
    ///
    /// - `id` is a valid public queue ID, i.e. `id < num_workers`.
    /// - `len` is a valid public queue length, i.e. `len <= pq_size`.
    #[inline]
    pub const unsafe fn new(
        id: usize,
        len: NonZeroUsize,
        config: &Config,
    ) -> Self {
        debug_assert!(id < config.num_workers.get());
        debug_assert!(len.get() <= config.batch_size.get());

        Self {
            // SAFETY:
            // As per 'params', 'num_workers * pq_size <= 2^31'.
            //
            // - 'id < num_workers' as per the caller.
            // - Thus 'id + 1 <= num_workers'.
            // - Thus '(id + 1) * pq_size <= num_workers * pq_size <= 2^31'.
            // - Thus 'id * pq_size + pq_size <= 2^31'.
            //
            // - 'len <= pq_size' as per the caller.
            // - Thus 'len - 1 < pq_size'.
            // - Thus 'id * pq_size + (len - 1) < 2^31'.
            //
            // - 'raw = id * pq_size + (len - 1)'.
            // - Thus 'raw' fits in a 'u32' without overflow.
            // - Thus 'self.id() == raw // pq_size == id'.
            // - Thus 'self.len() == raw % pq_size == len'.
            raw: (id * config.batch_size.get() + (len.get() - 1)) as u32,
        }
    }

    /// Reconstruct a [`PubQueue`] from a raw integer.
    ///
    /// ## Safety
    ///
    /// - `raw` must be the result of `PubQueue.raw`.
    #[inline]
    pub const unsafe fn from_raw(raw: u32) -> Self {
        debug_assert!(raw < (1u32 << 31));

        Self { raw }
    }

    /// The queue ID.
    ///
    /// ## Safety
    ///
    /// - `0 <= id < 2^31 / pq_size`.
    #[inline]
    pub const fn id(&self, config: &Config) -> usize {
        debug_assert!(self.raw < (1u32 << 31));

        self.raw as usize / config.batch_size.get()
    }

    /// The number of tasks in the batch.
    ///
    /// ## Safety
    ///
    /// - `1 <= len <= pq_size`.
    /// - `len` fits in [`u32`] and [`usize`].
    #[inline]
    pub const fn len(&self, config: &Config) -> NonZeroUsize {
        let batch_size = config.batch_size.get();

        // SAFETY:
        // - 'raw % pq_size < pq_size <= 2^31'.
        // - Thus '(raw % pq_size) + 1 <= pq_size <= 2^31'.
        unsafe {
            NonZeroUsize::new_unchecked((self.raw as usize % batch_size) + 1)
        }
    }

    /// Use this [`PubQueue`], with a new length.
    ///
    /// ## Invariants
    ///
    /// - `len` is a valid queue length, i.e. `len <= pq_size`.
    #[inline]
    pub const unsafe fn with_len(
        self,
        len: NonZeroUsize,
        config: &Config,
    ) -> Self {
        debug_assert!(len.get() <= config.batch_size.get());

        // SAFETY: 'len' is a valid queue length, as per the caller.
        unsafe { Self::new(self.id(config), len, config) }
    }
}

//----------- PQContents -------------------------------------------------------

/// The contents of a public queue.
#[repr(transparent)]
pub struct PQContents<T> {
    /// The tasks in the queue.
    tasks: [UnsafeCell<MaybeUninit<T>>],
}

impl<T> PQContents<T> {
    /// Construct a new [`PQContents`] on the heap.
    pub fn new_boxed(config: &Config) -> Box<Self> {
        // SAFETY: 'pq_size' fits in a 'usize'.
        let size: usize = config.batch_size.get();

        let ptr = Box::<[MaybeUninit<T>]>::new_uninit_slice(size);
        // SAFETY:
        // - 'UnsafeCell<T>' has the same layout as 'T'.
        // - 'PQContents<T>' has the same layout as '[UnsafeCell<...>]'.
        unsafe { transmute(ptr) }
    }

    /// Read an element out of the queue.
    ///
    /// The element might not be initialized, or might even be corrupted; it is
    /// up to the caller to unwrap it.
    ///
    /// ## Safety
    ///
    /// - `0 <= index < pq_size`.
    #[inline]
    pub unsafe fn read(&self, index: usize) -> MaybeUninit<T> {
        debug_assert!(index < self.tasks.len());

        // SAFETY:
        // - Another thread may be reading or writing the same data, leaving it
        //   in a corrupted state; it is up to the caller to determine this.
        unsafe { self.tasks.get_unchecked(index).get().read() }
    }

    /// Move all tasks out of this queue and into the given [`Vec`].
    ///
    /// ## Safety
    ///
    /// - `0 <= len <= SIZE`.
    /// - The first `len` elements of `self` must be initialized.
    /// - `self` must not be written to over this time.
    pub unsafe fn move_to(&self, len: usize, vec: &mut Vec<T>) {
        // SAFETY: 'tasks.len()' fits in a 'u32'.
        debug_assert!(len <= self.tasks.len());

        // SAFETY: 'len <= pq_size', thus 'len' fits in a 'usize'.
        let tasks = unsafe { self.tasks.get_unchecked(..len) };

        // SAFETY: 'data' is not being written to right now.
        let data = unsafe {
            transmute::<&[UnsafeCell<MaybeUninit<T>>], &[MaybeUninit<T>]>(tasks)
        };

        unsafe { extend_vec_from_slice(vec, data) };
    }

    /// Fill this queue from the given iterator.
    ///
    /// ## Safety
    ///
    /// - The iterator must fit in the batch.
    /// - `self` must not be written to (externally) over this time.
    pub unsafe fn fill(&self, iter: impl Iterator<Item = T>) {
        iter.zip(&self.tasks).for_each(|(elem, slot)| {
            // SAFETY: 'self' is not being written to over this time.
            unsafe { slot.get().write(MaybeUninit::new(elem)) };
        });
    }
}

unsafe impl<T> Send for PQContents<T> {}
unsafe impl<T> Sync for PQContents<T> {}

//----------- Tests ------------------------------------------------------------

#[cfg(test)]
mod tests {
    use std::{
        num::{NonZeroU32, NonZeroUsize},
        sync::{Barrier, atomic::Ordering},
        thread,
        vec::Vec,
    };

    use crate::{Config, Task};

    use super::{PQContents, Priority, PubQueue, Stealer};

    #[test]
    fn priority_theft() {
        enum Foo {}

        impl Task for Foo {
            type Priority = NonZeroU32;

            fn priority(&self) -> Self::Priority {
                match *self {}
            }
        }

        let priority = Priority::<Foo>::default();
        unsafe { priority.set(NonZeroU32::new(42).unwrap()) };
        let barrier = Barrier::new(2);
        let mut stealer_success = None;
        let mut owner_success = None;

        thread::scope(|s| {
            // Stealer thread.
            s.spawn(|| {
                barrier.wait();
                let min = Some(NonZeroU32::new(30).unwrap());
                let res = unsafe { priority.steal_if_above(min) };
                stealer_success = Some(res.is_some());
                if let Some(res) = res {
                    assert_eq!(res.get(), 42);
                }
            });

            // Owning thread.
            s.spawn(|| {
                barrier.wait();
                let orig = NonZeroU32::new(42).unwrap();
                let res = unsafe { priority.steal_self(orig) };
                owner_success = Some(res.is_ok());
            });
        });

        assert!(stealer_success.unwrap() != owner_success.unwrap());
    }

    #[test]
    fn id_len() {
        let num_workers = NonZeroUsize::new(4).unwrap();
        let batch_size = NonZeroUsize::new(4).unwrap();
        let config = Config::new(num_workers).with_batch_size(batch_size);

        let cases = [(0, 1), (3, 4)];
        for (id, len) in cases {
            let len = NonZeroUsize::new(len).unwrap();
            // SAFETY: 'id < 4' and 'len <= 4'.
            let pq = unsafe { PubQueue::new(id, len, &config) };

            assert_eq!(pq.id(&config), id);
            assert_eq!(pq.len(&config), len);
        }
    }

    #[test]
    fn pq_theft() {
        let num_workers = NonZeroUsize::new(2).unwrap();
        let batch_size = NonZeroUsize::new(4).unwrap();
        let config = Config::new(num_workers).with_batch_size(batch_size);

        let stealer = unsafe { Stealer::new(0, &config) };
        let barrier = Barrier::new(2);

        assert_eq!(stealer.raw.load(Ordering::Relaxed), 3);

        thread::scope(|s| {
            // Stealer thread.
            s.spawn(|| {
                barrier.wait();
                let pq = unsafe { stealer.steal(1, &config) };
                assert_eq!(pq.id(&config), 0);
                assert_eq!(pq.len(&config).get(), 4);
            });

            // Owning thread.
            s.spawn(|| {
                barrier.wait();
                let orig = unsafe { PubQueue::new(0, batch_size, &config) };
                let pq = unsafe { stealer.wait_for_theft(orig) };
                assert_eq!(pq.id(&config), 1);
                assert_eq!(pq.len(&config).get(), 4);
            });
        })
    }

    #[test]
    fn contents() {
        let num_workers = NonZeroUsize::new(2).unwrap();
        let batch_size = NonZeroUsize::new(4).unwrap();
        let config = Config::new(num_workers).with_batch_size(batch_size);

        // Initialize the public queue contents.
        let contents = PQContents::<u32>::new_boxed(&config);
        unsafe { contents.fill(4..8) };

        // Try reading from the contents.
        for index in 0..4 {
            assert_eq!(
                unsafe { contents.read(index).assume_init() },
                index as u32 + 4
            );
        }

        // Try moving the contents out.
        let mut elems = Vec::new();
        unsafe { contents.move_to(4, &mut elems) };
        assert_eq!(elems, [4, 5, 6, 7]);
    }
}