//! Lockless Work Stealing Job System.
//!
//! # Features
//! - Lockless work stealing queue
//! - 0 runtime allocations
//! - Jobs are executed from closures
//! - Chaining and grouping of jobs
//! - Parallel for_each abstraction
//!
//! # General Notes
//! This crate provides job system with 0 allocation overhead at runtime. What this means in
//! practice is, that once the instance is created, there are no more memory allocations required
//! to execute jobs.
//!
//! Jobs are allocated from a Job Pool with a fixed capacity. Each thread gets their own pool and
//! queue. The system uses a fixed storage space where it stores the closure for each job. The
//! default storage is set to 64 bytes. This can be extended to 128 by enabling the feature
//! _"job_storage_128"_.
//!
//! When a Job finishes execution there is a possibility to start other jobs. See
//! [JobSystem::chain()] for more details.
//!
//! Additionally, it's recommended that, in order to avoid running out of jobs during execution, you
//! regularly ensure that you all your previous jobs have finished executing.
//!
//! Finally, this crate runs on rust stable and has been tested with rust 1.50.0 with the 2018 edition.
//! # Panics
//!
//! Due to the queues having a fixed size, if we start filling all the queues up to full capacity there
//! is a small chance that a worker thread may not be able to queue chained jobs. When this happens
//! the system will panic as there is no way to recover from this.
//! The queues come with debug asserts which will detect this situation.
//!
//! # Examples
//! ## Start and Wait
//! ```
//! let mut js = jobsys::JobSystem::new(4, 512).unwrap();
//! let mut handle = js.create(|| { println!("Hello World!");}).unwrap();
//! js.run(&mut handle).expect("Failed to run job");
//! js.wait(&mut handle);
//! ```
//! ## Grouping
//! Ensure one job does not complete until other jobs have finished as well.
//! ```
//! let mut js = jobsys::JobSystem::new(4, 512).unwrap();
//! let mut parent = js.create(|| {}).unwrap();
//! let mut child = js.create_with_parent(&mut parent, || { println!("Hello World!");}).unwrap();
//! js.run(&child).expect("Failed to start child");
//! js.run(&parent).expect("Failed to start parent");
//! js.wait(&parent); // Parent will only finish when both it and its child have finished
//! ```
//! ## Chaining
//! Launch new jobs as soon as one job completes.
//! ```
//! let mut js = jobsys::JobSystem::new(4, 512).unwrap();
//! let mut first = js.create(|| { println!("Hello World Chained!");}).unwrap();
//! let mut second = js.create(|| { println!("Hello World Chained!");}).unwrap();
//! js.chain(&mut first, &second).expect("Failed to chain job, maximum chain count exceeded");
//! js.run(&first).expect("Failed to start job");
//! js.wait(&second); // Second will only be executed after first completes
//! ```
//! ## Parallel
//! See [for_each()][JobSystem::for_each()] and [for_each_with_result()][JobSystem::for_each_with_result]
//! for more details.
//! ```
//! let mut js = jobsys::JobSystem::new(4, 512).unwrap();
//! let mut array = [0_u32; 100];
//! js.for_each(|slice: &mut [u32], start, _end| {
//!         for i in 0..slice.len() {
//!             slice[i] = (start + i) as u32;
//!         }
//!     },
//!     &mut array).expect("Failed to start jobs");
//! ```
mod job;
mod queues;
mod random;
mod thread;

use crate::job::JobHandlePrivate;
use crate::thread::JobThread;
use crate::thread::ThreadDataWrapper;
use std::marker::PhantomData;
use std::sync::Arc;

/// Handle which represents an allocated job.
pub struct JobHandle<T>
where
    T: Sized + FnMut() + Send,
{
    h: JobHandlePrivate,
    p: PhantomData<T>,
}

type ThreadDataList = Vec<ThreadDataWrapper>;
/// Work Stealing JobSystem.
pub struct JobSystem {
    thread_data: std::sync::Arc<ThreadDataList>,
    threads: Vec<JobThread>,
}

#[derive(Debug)]
pub enum Error {
    /// Failed to create/start a job Thread
    ThreadCreate,
    /// Requested job pool/queue capacity is not a power of 2
    CapacityNotPowerOf2,
    /// The submitted closure exceeds the storage capacity
    StorageSizeExceeded,
    /// The thread queue is full and can't be submitted to
    QueueFull,
    /// The number of chained jobs has exceed the limit
    ChainCountExceeded,
}

impl JobSystem {
    /// Create a new instance of a JobSystem.
    ///
    /// * `thread_count` - Number of worker threads.
    /// * `job_capacity` - Maximum number of jobs the system can allocate. Must be a power of 2.
    pub fn new(thread_count: usize, job_capacity: usize) -> Result<Self, Error> {
        if !job_capacity.is_power_of_two() {
            return Err(Error::CapacityNotPowerOf2);
        }
        let actual_thread_count = thread_count.max(1) + 1;
        let mut job_sys = Self {
            thread_data: std::sync::Arc::new(ThreadDataList::new()),
            threads: vec![],
        };
        let mut data_vec: ThreadDataList = Vec::with_capacity(actual_thread_count);
        data_vec.resize_with(actual_thread_count, || ThreadDataWrapper::new(job_capacity));
        job_sys.thread_data = std::sync::Arc::new(data_vec);
        job_sys
            .threads
            .resize_with(actual_thread_count - 1, JobThread::new);
        for index in 1..actual_thread_count {
            let thread = &mut job_sys.threads[index - 1];
            thread.set_data(job_sys.thread_data.clone(), index);
            if thread.start().is_err() {
                return Err(Error::ThreadCreate);
            }
        }
        Ok(job_sys)
    }

    fn shutdown(&mut self) {
        for thread in &mut self.threads {
            thread.finish().unwrap()
        }
        self.threads.clear();
        self.thread_data = Arc::new(vec![]);
    }

    /// Allocate a new Job
    ///
    /// This function will check whether the closure fits in the job storage space and return an
    /// error if that's the case.
    pub fn create<T>(&mut self, job: T) -> Result<JobHandle<T>, Error>
    where
        T: Sized + FnMut() + Send + Sync,
    {
        debug_assert!(std::mem::size_of::<T>() <= job::JOB_STORAGE_SIZE);
        if std::mem::size_of::<T>() > job::JOB_STORAGE_SIZE {
            return Err(Error::StorageSizeExceeded);
        }
        let mut handle = thread::alloc_job(&self.thread_data);
        handle.store(job);
        Ok(JobHandle {
            h: handle,
            p: PhantomData,
        })
    }

    /// Allocate a new Job as a child of another job.
    ///
    /// The newly allocated will be created as a child of a parent job. What this means in practice
    /// is that the new job (child) can run parallel with the parent job and the parent job will
    /// not reach completion status until all of it's children have finished.
    pub fn create_with_parent<T, Y>(
        &mut self,
        parent: &mut JobHandle<Y>,
        job: T,
    ) -> Result<JobHandle<T>, Error>
    where
        T: Sized + FnMut() + Send,
        Y: Sized + FnMut() + Send,
    {
        debug_assert!(std::mem::size_of::<T>() <= job::JOB_STORAGE_SIZE);
        if std::mem::size_of::<T>() > job::JOB_STORAGE_SIZE {
            return Err(Error::StorageSizeExceeded);
        }
        let mut handle = thread::alloc_job(&self.thread_data);
        handle.store(job);
        handle.set_parent_job(parent.h);
        Ok(JobHandle {
            h: handle,
            p: PhantomData,
        })
    }

    /// Register the child job as a follow up to the parent job.
    ///
    /// Every Job has the ability to chain follow up jobs on completion. This ensures the child
    /// job only runs when the parent is finished.
    ///
    /// There is a limited capacity for the number of jobs you can chain with one job. This function
    /// will return error if we are unable to register the child job.
    pub fn chain<T, Y>(&self, parent: &mut JobHandle<T>, child: &JobHandle<Y>) -> Result<(), Error>
    where
        T: Sized + FnMut() + Send,
        Y: Sized + FnMut() + Send,
    {
        match parent.h.chain_job(child.h) {
            Err(_) => Err(Error::ChainCountExceeded),
            _ => Ok(()),
        }
    }

    /// Execute a job.
    pub fn run<T>(&mut self, handle: &JobHandle<T>) -> Result<(), Error>
    where
        T: Sized + FnMut() + Send,
    {
        match thread::start_job(&self.thread_data, handle.h) {
            true => Ok(()),
            false => Err(Error::QueueFull),
        }
    }

    /// Block and wait until the job has finished.
    ///
    /// While we wait, the system will try to execute other jobs.
    pub fn wait<T>(&mut self, handle: &JobHandle<T>)
    where
        T: Sized + FnMut() + Send,
    {
        while !handle.h.is_finished() {
            if let Some(job) = thread::get_job(&self.thread_data) {
                debug_assert!(job.is_valid());
                thread::run_job(&self.thread_data, job);
            }
        }
    }

    /// Check if a job has finished execution. Does not block.
    pub fn is_finished<T>(&mut self, handle: &JobHandle<T>) -> bool
    where
        T: Sized + FnMut() + Send,
    {
        handle.h.is_finished()
    }

    /// Given a slice of data and read-only closure, divide the slice into unique sub-slices which
    /// are distributed to the worker threads.
    ///
    /// The closure will receive the following parameters in order:
    /// * unique sub-slice over which the the current thread is operating on
    /// * start index of the full slice
    /// * end index of the the full slice
    /// Note use [for_each_with_result()][JobSystem::for_each_with_result()] if you wish to
    /// produce output.
    pub fn for_each<'env, T, Y>(&mut self, cb: T, slice: &'env mut [Y]) -> Result<(), Error>
    where
        T: Fn(&mut [Y], usize, usize) + 'env + Send + Sync,
        Y: Send,
    {
        let mut parent_job = self.create(move || {})?;

        //const DEFAULT_GROUP_SIZE: usize = 64;
        let divisor = self.threads.len();
        let group_size = (slice.len() / divisor).max(divisor);
        let mut offset = 0_usize;
        let mut remaining = slice.len();

        while remaining != 0 {
            let range = group_size.min(remaining);
            let slice_ptr = unsafe { slice.as_mut_ptr().add(offset) };
            let work_slice = unsafe { std::slice::from_raw_parts_mut(slice_ptr, range) };
            let callback = &cb;
            let child_job = self.create_with_parent(&mut parent_job, move || {
                callback(work_slice, offset, offset + work_slice.len());
            })?;
            self.run(&child_job)?;
            remaining -= range;
            offset += range;
        }
        self.run(&parent_job)?;
        self.wait(&parent_job);
        Ok(())
    }

    /// Same as [for_each()][JobSystem::for_each()] but allows a result type to be returned for every individual group.
    pub fn for_each_with_result<'env, T, Y, Z>(
        &mut self,
        cb: T,
        slice: &'env mut [Y],
    ) -> Result<Vec<Z>, Error>
    where
        T: Fn(&mut [Y], usize, usize) -> Z + 'env + Send + Sync,
        Z: Sized + Default + Send,
        Y: Send,
    {
        let mut parent_job = self.create(|| {})?;

        //const DEFAULT_GROUP_SIZE: usize = 64;
        let divisor = self.threads.len();
        let group_size = (slice.len() / divisor).max(divisor);
        let group_count = slice.len() % divisor;
        let vec_size = if group_count == 0 {
            divisor
        } else {
            divisor + 1
        };
        let mut offset = 0_usize;
        let mut remaining = slice.len();
        let mut group_index = 0_usize;
        let mut result_vec = Vec::<Z>::with_capacity(vec_size);
        for _ in 0..vec_size {
            result_vec.push(Z::default());
        }
        while remaining != 0 {
            let range = group_size.min(remaining);
            let slice_ptr = unsafe { slice.as_mut_ptr().add(offset) };
            let result_ref = unsafe { &mut *result_vec.as_mut_ptr().add(group_index) };
            let work_slice = unsafe { std::slice::from_raw_parts_mut(slice_ptr, range) };
            let callback = &cb;
            let child_job = self.create_with_parent(&mut parent_job, move || {
                *result_ref = callback(work_slice, offset, offset + work_slice.len());
            })?;
            self.run(&child_job)?;
            remaining -= range;
            offset += range;
            group_index += 1;
        }
        self.run(&parent_job)?;
        self.wait(&parent_job);
        Ok(result_vec)
    }
}

impl Drop for JobSystem {
    fn drop(&mut self) {
        self.shutdown();
    }
}

#[cfg(test)]
mod tests {
    use crate::JobSystem;
    use std::cell::RefCell;

    const THREAD_COUNT: usize = 4;
    const JOB_CAPACITY: usize = 1024;

    #[test]
    fn start_stop() {
        let r = JobSystem::new(THREAD_COUNT, JOB_CAPACITY);
        assert!(r.is_ok());
    }

    #[test]
    fn launch_jobs() {
        let mut job_sys = JobSystem::new(THREAD_COUNT, 128).expect("Failed to init job system");
        let mut _counter = 0_usize;
        for _ in 0..8192_u32 {
            const JOB_COUNT: usize = 100;
            let mut jobs = Vec::<_>::with_capacity(JOB_COUNT);

            for _ in 0..JOB_COUNT {
                let handle = job_sys.create(|| {}).unwrap();
                assert!(job_sys.run(&handle).is_ok());
                jobs.push(handle);
            }
            for job in jobs {
                job_sys.wait(&job);
            }
            _counter += JOB_COUNT;
        }
    }

    #[test]
    fn launch_jobs_with_ref() {
        let mut job_sys =
            JobSystem::new(THREAD_COUNT, JOB_CAPACITY).expect("Failed to init job system");
        const JOB_COUNT: usize = 100;
        let mut jobs = Vec::<_>::with_capacity(JOB_COUNT);

        let val = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
        for _ in 0..JOB_COUNT {
            let val_copy = val.clone();
            let handle = job_sys
                .create(move || {
                    val_copy.fetch_add(10, std::sync::atomic::Ordering::Release);
                })
                .unwrap();
            assert!(job_sys.run(&handle).is_ok());
            jobs.push(handle);
        }
        for job in jobs {
            job_sys.wait(&job);
        }
        assert_eq!(
            val.load(std::sync::atomic::Ordering::Acquire),
            10 * JOB_COUNT as u32
        );
        job_sys.shutdown();
    }

    #[test]
    fn launch_jobs_chained() {
        let mut job_sys =
            JobSystem::new(THREAD_COUNT, JOB_CAPACITY).expect("Failed to init job system");
        const JOB_COUNT: usize = 20;
        let mut jobs = Vec::<_>::with_capacity(JOB_COUNT);
        for i in 0..JOB_COUNT {
            let handle = job_sys
                .create(move || {
                    println!("Chained {:?}: Job {:02}", std::thread::current().id(), i);
                })
                .unwrap();
            jobs.push(RefCell::new(handle));
            if i > 0 {
                let cur_handle = &jobs[i];
                let prev_handle = &jobs[i - 1];
                job_sys
                    .chain(&mut prev_handle.borrow_mut(), &cur_handle.borrow())
                    .expect("Failed to chain");
            }
        }
        assert!(job_sys.run(&jobs.first().unwrap().borrow_mut()).is_ok());
        job_sys.wait(&jobs.last().unwrap().borrow_mut());
        job_sys.shutdown();
    }

    #[test]
    fn parallel_for() {
        let mut job_sys =
            JobSystem::new(THREAD_COUNT, JOB_CAPACITY).expect("Failed to init job system");

        let mut array = [0_u32; 100];

        let r = job_sys.for_each(
            |slice: &mut [u32], start, _end| {
                for i in 0..slice.len() {
                    slice[i] = (start + i) as u32;
                }
            },
            &mut array,
        );
        assert!(r.is_ok());
        for i in 0..array.len() {
            assert_eq!(array[i] as usize, i);
        }
    }

    #[test]
    fn launch_with_parent() {
        let mut job_sys =
            JobSystem::new(THREAD_COUNT, JOB_CAPACITY).expect("Failed to init job system");
        const JOB_COUNT: usize = 20;
        let mut parent = job_sys.create(|| {}).unwrap();
        let mut jobs = Vec::<_>::with_capacity(JOB_COUNT);
        for _i in 1..JOB_COUNT {
            let handle = job_sys
                .create_with_parent(&mut parent, move || {
                    /*
                    println!(
                        "Hello from thread {:?}: Job {:02}",
                        std::thread::current().id(),
                        i
                    );*/
                })
                .unwrap();
            jobs.push(handle);
        }
        assert!(job_sys.run(&parent).is_ok());
        for job in &jobs {
            assert!(job_sys.run(job).is_ok());
        }
        job_sys.wait(&parent);
        job_sys.shutdown();
    }

    #[test]
    fn parallel_for_with_result() {
        let mut job_sys =
            JobSystem::new(THREAD_COUNT, JOB_CAPACITY).expect("Failed to init job system");

        let mut array = [0_u32; 100];

        let r = job_sys.for_each_with_result(
            |slice: &mut [u32], start, end| -> u32 {
                for i in 0..slice.len() {
                    slice[i] = (start + i) as u32;
                }
                (end - start) as u32
            },
            &mut array,
        );
        assert!(r.is_ok());
        let result: u32 = r.unwrap().iter().sum();
        assert_eq!(result, 100_u32);
    }
}