//! Benchmarking `takeaway` using Criterion.

use std::{
    num::NonZeroUsize,
    ops::ControlFlow,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
        mpsc,
    },
    task::{Poll, Waker},
};

use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use takeaway::util::CachePadded;

//----------- Implementations --------------------------------------------------

/// An implementation of a task queue.
pub trait Implementation<T: Task>: Send + Sync + Sized + 'static {
    /// Thread-local state for this implementation.
    type Local: Send + 'static;

    /// The name of this implementation.
    const NAME: &'static str;

    /// Initialize the implementation.
    fn init(
        num_workers: NonZeroUsize,
        batch_size: NonZeroUsize,
    ) -> (Self, Vec<Self::Local>);

    /// Drive a worker thread.
    ///
    /// Return the number of locally executed tasks.
    fn drive(&self, local: Self::Local, context: &T::Context, limit: usize);

    /// Shut down a running instance.
    fn stop(&self);
}

pub struct Takeaway<T: takeaway::Task>(takeaway::Queue<T>);

impl<T: Task + takeaway::Task> Implementation<T> for Takeaway<T> {
    // The ID of each worker.
    type Local = usize;

    const NAME: &'static str = "takeaway";

    fn init(
        num_workers: NonZeroUsize,
        batch_size: NonZeroUsize,
    ) -> (Self, Vec<Self::Local>) {
        let queue = takeaway::Config::new(num_workers)
            .with_batch_size(batch_size)
            .build();
        let locals = (0..num_workers.get()).collect();
        (Self(queue), locals)
    }

    fn drive(&self, local: Self::Local, context: &T::Context, limit: usize) {
        let mut executed = 0usize;
        let mut worker = takeaway::Worker::new(&self.0, local);
        while executed < limit {
            match worker.poll(Waker::noop()) {
                Poll::Ready(Some(task)) => {
                    task.execute(context, |t| worker.enqueue_one(t));
                    executed += 1;
                }
                Poll::Ready(None) => break,
                Poll::Pending => {
                    // The worker ran out of tasks, add some more.
                    worker.enqueue((0..16).map(|_| T::random()));
                }
            }
        }
    }

    fn stop(&self) {
        self.0.shutdown();
    }
}

pub struct CrossbeamDeque<T> {
    /// The configured batch size.
    batch_size: NonZeroUsize,

    /// The stealer of every worker.
    stealers: Vec<crossbeam_deque::Stealer<T>>,

    /// Whether a shutdown has occurred.
    shutdown: Box<CachePadded<AtomicBool>>,
}

impl<T: Task> Implementation<T> for CrossbeamDeque<T> {
    type Local = crossbeam_deque::Worker<T>;

    const NAME: &'static str = "crossbeam-deque";

    fn init(
        num_workers: NonZeroUsize,
        batch_size: NonZeroUsize,
    ) -> (Self, Vec<Self::Local>) {
        let workers = (0..num_workers.get())
            .map(|_| crossbeam_deque::Worker::new_fifo())
            .collect::<Vec<_>>();
        let stealers = workers.iter().map(|w| w.stealer()).collect();
        let shutdown = Box::new(CachePadded::new(AtomicBool::new(false)));
        (
            Self {
                batch_size,
                stealers,
                shutdown,
            },
            workers,
        )
    }

    fn drive(&self, local: Self::Local, context: &T::Context, limit: usize) {
        let mut executed = 0usize;
        'main: while executed < limit {
            let task = 'fetch: {
                // Stop if all tasks are complete.
                if self.shutdown.load(Ordering::Relaxed) {
                    break 'main;
                }

                // Pop a task from the local queue, if not empty.
                if let Some(task) = local.pop() {
                    break 'fetch task;
                }

                let split = rand::random_range(0..self.stealers.len());
                let mut stealers = self.stealers[split..]
                    .iter()
                    .chain(&self.stealers[..split]);

                match stealers.try_for_each(|stealer| {
                    match stealer.steal_batch_with_limit_and_pop(
                        &local,
                        self.batch_size.get(),
                    ) {
                        crossbeam_deque::Steal::Success(task) => {
                            ControlFlow::Break(task)
                        }
                        _ => ControlFlow::Continue(()),
                    }
                }) {
                    ControlFlow::Break(task) => break 'fetch task,
                    ControlFlow::Continue(()) => {}
                }

                // We ran out of tasks; add some more.
                for _ in 0..16 {
                    local.push(T::random());
                }
                continue 'main;
            };

            task.execute(context, |t| local.push(t));
            executed += 1;
        }
    }

    fn stop(&self) {
        self.shutdown.store(true, Ordering::Relaxed);
    }
}

//----------- Tasks ------------------------------------------------------------

/// A benchmarking task.
pub trait Task: Send + Sized + 'static {
    /// Context for executing tasks.
    type Context;

    /// The name of this task type.
    const NAME: &'static str;

    /// Generate a random task.
    fn random() -> Self;

    /// Execute the task.
    ///
    /// The returned tasks should be enqueued for future execution.
    fn execute(self, context: &Self::Context, enqueue: impl Fn(Self));
}

/// A trivial no-op task.
///
/// Each iteration performs a few arithmetic operations.
pub struct Noop {
    data: usize,
}

impl Task for Noop {
    type Context = ();
    const NAME: &'static str = "noop";

    fn random() -> Self {
        Self {
            data: rand::random::<u32>() as usize,
        }
    }

    fn execute(self, _context: &Self::Context, enqueue: impl Fn(Self)) {
        for i in 0..rand::random_range(0..4) {
            enqueue(Self {
                data: self.data.wrapping_mul(4).wrapping_add(i),
            })
        }
    }
}

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

/// A lightweight task.
///
/// Each iteration performs string formatting, on an ever-growing string.
pub struct Light {
    data: String,
}

impl Task for Light {
    type Context = ();
    const NAME: &'static str = "light";

    fn random() -> Self {
        Self {
            data: format!("{}", rand::random::<u32>() as usize),
        }
    }

    fn execute(self, _context: &Self::Context, enqueue: impl Fn(Self)) {
        for i in 0..rand::random_range(0..4) {
            let data = format!("{}{i}", self.data);
            let data = data.replace(['8', '9'], "");
            let data = usize::from_str_radix(&data, 8).unwrap();
            let data = format!("{data}");
            enqueue(Self { data })
        }
    }
}

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

//----------- Running ----------------------------------------------------------

/// A benchmarker for a task queue implementation.
pub struct Runner<I, T>
where
    I: Implementation<T>,
    T: Task<Context = ()>,
{
    /// Control of every benchmarking thread.
    ctrls: Vec<ChannelControl<I, I::Local>>,
}

struct ChannelControl<I, L> {
    start: mpsc::SyncSender<(Arc<I>, L, usize)>,
    finish: mpsc::Receiver<()>,
}

impl<I, T> Runner<I, T>
where
    I: Implementation<T>,
    T: Task<Context = ()>,
{
    /// Initialize a [`Runner`].
    pub fn new() -> Self {
        let num = std::thread::available_parallelism().unwrap();
        let ctrls = (0..num.get())
            .map(|_| {
                let (start_tx, start_rx) = mpsc::sync_channel(0);
                let (finish_tx, finish_rx) = mpsc::sync_channel(0);
                std::thread::spawn(|| Self::drive(start_rx, finish_tx));
                ChannelControl {
                    start: start_tx,
                    finish: finish_rx,
                }
            })
            .collect();
        Self { ctrls }
    }

    /// Drive a thread for this runner.
    fn drive(
        start: mpsc::Receiver<(Arc<I>, I::Local, usize)>,
        finish: mpsc::SyncSender<()>,
    ) {
        while let Ok((global, local, limit)) = start.recv() {
            global.drive(local, &(), limit);
            let _ = finish.send(());
        }
    }

    /// Perform a single iteration of a benchmark.
    fn iterate(
        &mut self,
        num_workers: NonZeroUsize,
        batch_size: NonZeroUsize,
        limit: usize,
    ) {
        let (global, locals) = I::init(num_workers, batch_size);
        let global = Arc::new(global);
        let ctrls = &self.ctrls[..num_workers.get()];

        for (ctrl, local) in ctrls.iter().zip(locals) {
            ctrl.start.send((global.clone(), local, limit)).unwrap();
        }

        for ctrl in ctrls {
            ctrl.finish.recv().unwrap();
        }
    }
}

impl<I, T> Default for Runner<I, T>
where
    I: Implementation<T>,
    T: Task<Context = ()>,
{
    fn default() -> Self {
        Self::new()
    }
}

//----------- main() -----------------------------------------------------------

pub fn evaluate<T>(criterion: &mut Criterion, limit: usize)
where
    T: Task<Context = ()> + takeaway::Task,
{
    let mut group = criterion.benchmark_group(T::NAME);
    let mut takeaway = Runner::<Takeaway<_>, T>::new();
    let mut crossbeam_deque = Runner::<CrossbeamDeque<_>, T>::new();

    let parallelism = std::thread::available_parallelism().unwrap();
    for num_workers in 1..=parallelism.get() {
        let num_workers = NonZeroUsize::new(num_workers).unwrap();

        for batch_size in [4, 16, 64, 256, 4096] {
            let batch_size = NonZeroUsize::new(batch_size).unwrap();

            let id = format!("{num_workers}t-{batch_size}b");

            group.bench_function(BenchmarkId::new("takeaway", &id), |b| {
                b.iter(|| takeaway.iterate(num_workers, batch_size, limit))
            });
            group.bench_function(
                BenchmarkId::new("crossbeam-deque", &id),
                |b| {
                    b.iter(|| {
                        crossbeam_deque.iterate(num_workers, batch_size, limit)
                    })
                },
            );
        }
    }

    group.finish();
}

fn benchmark(criterion: &mut Criterion) {
    let limit = 16384;
    evaluate::<Noop>(criterion, limit);
    evaluate::<Light>(criterion, limit);
}

criterion_group!(benches, benchmark);
criterion_main!(benches);
