takeaway 0.1.2

//! Representing tasks.

use core::{
    convert::Infallible,
    fmt,
    num::{NonZeroU8, NonZeroU16, NonZeroU32, NonZeroU64, NonZeroUsize},
};

use atomig::{Atom, impls::PrimitiveAtom};

//----------- Task -------------------------------------------------------------

/// A task.
///
/// To use a type for representing tasks in `takeaway`, it must implement this
/// trait.  It allows for additional functionality like prioritization.
///
/// # Usage
///
/// Due to language limitations, some of the trait's items cannot have defaults.
/// To start out, use the following template:
///
/// ```no_run
/// # struct MyTask;
/// impl takeaway::Task for MyTask {
///     type Priority = ();
///     fn priority(&self) -> Self::Priority {}
/// }
/// ```
///
/// If your tasks have a concept of priority, and you want higher-priority
/// tasks to be executed before others, you can override [`Task::Priority`] and
/// [`Task::priority()`].  Implementing your own priority type is involved, but
/// you can just use [`NonZeroU32`].
pub trait Task {
    /// The type of the priority of this task.
    type Priority: TaskPriority;

    /// The priority of this task.
    fn priority(&self) -> Self::Priority;
}

// TODO: Use '!' once it is stabilized.
impl Task for Infallible {
    type Priority = Infallible;

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

impl Task for () {
    type Priority = ();

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

//----------- TaskPriority -----------------------------------------------------

/// The priority of a task.
///
/// `takeaway` uses this type to locate high-priority tasks and execute them
/// before others.  [`NonZeroU32`] is a common choice, but you can use [`()`] to
/// disable prioritization entirely.
///
/// Internally, `takeaway` includes priority information when sharing tasks
/// between workers.  Tasks are only stolen when they have a higher priority
/// than the thief's local tasks, ensuring 1) high-priority tasks are propagated
/// through the system, and 2) theft only occurs when it will improve the
/// distribution of tasks.  Thus, task priorities are part of the delicate dance
/// `takeaway` performs to distribute tasks, and they have to be implemented
/// carefully.
///
/// If you don't want to implement a custom task priority, [`NonZeroU32`] is
/// often sufficient.  Note that higher priority values result in higher
/// priorities.
///
/// # Safety
///
/// A type `T` can soundly implement `TaskPriority` if and only if all of the
/// following conditions are satisfied:
///
/// - `T::TRIVIAL` is `true` if and only if `T` is `()`.
/// - `T::pack()` and `T::unpack()` are mutual inverses.
///   - For all `t`, `T::unpack(T::pack(t)) == t`.
/// - `T::pack()` and `T::unpack()` never panic.
pub unsafe trait TaskPriority: Copy + Ord + fmt::Debug {
    /// An atomic-capable type for representing these values.
    type Repr: Atom + PrimitiveAtom;

    /// Whether this is the "trivial" priority type `()`.
    ///
    /// If this is true, `takeaway` will perform some optimizations internally
    /// by relying on the fact that all tasks have equal priority.
    const TRIVIAL: bool = false;

    /// Pack this priority into the atomic-capable representation.
    fn pack(priority: Option<Self>) -> Self::Repr;

    /// Unpack a priority from the atomic-capable representation.
    fn unpack(repr: Self::Repr) -> Option<Self>;
}

// TODO: Use '!' once it is stabilized.
unsafe impl TaskPriority for Infallible {
    // TODO: Use '()' if that ever gets supported.
    type Repr = u8;

    fn pack(priority: Option<Self>) -> Self::Repr {
        let None = priority;
        0u8
    }

    fn unpack(repr: Self::Repr) -> Option<Self> {
        assert_eq!(repr, 0u8);
        None
    }
}

unsafe impl TaskPriority for () {
    type Repr = bool;
    const TRIVIAL: bool = true;

    fn pack(priority: Option<Self>) -> Self::Repr {
        match priority {
            Some(()) => true,
            None => false,
        }
    }

    fn unpack(repr: Self::Repr) -> Option<Self> {
        match repr {
            true => Some(()),
            false => None,
        }
    }
}

unsafe impl TaskPriority for NonZeroU8 {
    type Repr = u8;

    fn pack(priority: Option<Self>) -> Self::Repr {
        priority.map_or(0, NonZeroU8::get)
    }

    fn unpack(repr: Self::Repr) -> Option<Self> {
        NonZeroU8::new(repr)
    }
}

unsafe impl TaskPriority for NonZeroU16 {
    type Repr = u16;

    fn pack(priority: Option<Self>) -> Self::Repr {
        priority.map_or(0, NonZeroU16::get)
    }

    fn unpack(repr: Self::Repr) -> Option<Self> {
        NonZeroU16::new(repr)
    }
}

unsafe impl TaskPriority for NonZeroU32 {
    type Repr = u32;

    fn pack(priority: Option<Self>) -> Self::Repr {
        priority.map_or(0, NonZeroU32::get)
    }

    fn unpack(repr: Self::Repr) -> Option<Self> {
        NonZeroU32::new(repr)
    }
}

unsafe impl TaskPriority for NonZeroU64 {
    type Repr = u64;

    fn pack(priority: Option<Self>) -> Self::Repr {
        priority.map_or(0, NonZeroU64::get)
    }

    fn unpack(repr: Self::Repr) -> Option<Self> {
        NonZeroU64::new(repr)
    }
}

unsafe impl TaskPriority for NonZeroUsize {
    type Repr = usize;

    fn pack(priority: Option<Self>) -> Self::Repr {
        priority.map_or(0, NonZeroUsize::get)
    }

    fn unpack(repr: Self::Repr) -> Option<Self> {
        NonZeroUsize::new(repr)
    }
}