//! [![crates.io version](https://img.shields.io/crates/v/safina-executor.svg)](https://crates.io/crates/safina-executor)
//! [![license: Apache 2.0](https://gitlab.com/leonhard-llc/safina-rs/-/raw/main/license-apache-2.0.svg)](http://www.apache.org/licenses/LICENSE-2.0)
//! [![unsafe forbidden](https://gitlab.com/leonhard-llc/safina-rs/-/raw/main/unsafe-forbidden-success.svg)](https://github.com/rust-secure-code/safety-dance/)
//! [![pipeline status](https://gitlab.com/leonhard-llc/safina-rs/badges/main/pipeline.svg)](https://gitlab.com/leonhard-llc/safina-rs/-/pipelines)
//!
//! This is a safe Rust async executor.
//!
//! It is part of [`safina`](https://crates.io/crates/safina), a safe async runtime.
//!
//! # Features
//! - `forbid(unsafe_code)`
//! - Depends only on `std`
//! - Good test coverage (100%)
//!
//! # Limitations
//! - Requires Rust `nightly`, for
//!   [Wake trait](https://doc.rust-lang.org/std/task/trait.Wake.html)
//! - Allocates memory
//! - Not optimized
//!
//! # Documentation
//! <https://docs.rs/safina-executor>
//!
//! # Examples
//! ```rust
//! # fn f() {
//! let executor = safina_executor::Executor::default();
//! let (sender, receiver) = std::sync::mpsc::channel();
//! executor.spawn(async move {
//!     sender.send(()).unwrap();
//! });
//! receiver.recv().unwrap();
//! # }
//! ```
//!
//! ```rust
//! # async fn prepare_request() -> Result<(), std::io::Error> { Ok(()) }
//! # async fn execute_request() -> Result<(), std::io::Error> { Ok(()) }
//! # fn f() -> Result<(), std::io::Error> {
//! let result = safina_executor::block_on(async {
//!     prepare_request().await?;
//!     execute_request().await
//! })?;
//! # Ok(())
//! # }
//! ```
//!
//! # Alternatives
//! - [`async-executor`](https://crates.io/crates/async-executor)
//!   - Popular
//!   - Dependencies have some `unsafe` code
//! - [`futures-executor`](https://crates.io/crates/futures-executor)
//!   - Very popular
//!   - Full of `unsafe`
//! - [`tokio-executor`](https://crates.io/crates/tokio-executor)
//!   - Very popular
//!   - Fast
//!   - Internally complicated
//!   - Full of `unsafe`
//! - [`executors`](https://crates.io/crates/executors)
//!   - Dependencies have lots of `unsafe` code
//! - [`bastion-executor`](https://crates.io/crates/bastion-executor)
//!   - Lots of `unsafe` code
//! - [`rayon_core`](https://crates.io/crates/rayon-core)
//!   - Generous amounts of `unsafe` code
//! - [`pollster`](https://crates.io/crates/pollster)
//!   - Minimal
//!   - Has a little `unsafe` code
//! - [`lelet`](https://crates.io/crates/lelet)
//!   - Automatically scales worker thread pool
//!   - Lots of `unsafe` code
//! - [`fibers`](https://crates.io/crates/fibers)
//!   - Dependencies are full of unsafe
//! - [`nostd_async`](https://crates.io/crates/nostd_async)
//!   - Has some `unsafe` code
//! - [`embedded-executor`](https://crates.io/crates/embedded-executor)
//!   - Generous amounts of `unsafe` code
//! - [`spin_on`](https://crates.io/crates/spin_on)
//!   - Minimal
//! - [`pasts`](https://crates.io/crates/pasts)
//! - [`switchyard`](https://crates.io/crates/switchyard)
//!   - Lots of `unsafe` code
//! - [`sealrs`](https://crates.io/crates/sealrs)
//! - [`rusty_pool`](https://crates.io/crates/rusty_pool)
//!   - Automatically adjusts thread pool
//!   - Dependencies are full of unsafe
//!
//! # Changelog
//! - v0.1.3
//!   - Removed global executor.  Users must explicitly create executor.
//!   - Removed dependency on unstable
//!     [`OnceCell`](https://doc.rust-lang.org/std/lazy/struct.OnceCell.html).
//!   - Uses [`safina_threadpool`](https://crates.io/crates/safina-threadpool)
//!     internally.
//! - v0.1.2 - Let callers pass futures to `spawn` and `block_on` without
//!   using `Box::pin`.
//!   Add `spawn_unpin` and `block_on_unpin` for folks who need to avoid allocating.
//!   so callers don't have to.
//! - v0.1.1 - Fix badge and update readme
//! - v0.1.0 - Renamed from `safina`
//!
//! # TO DO
//! - DONE - Implement `spawn`
//! - DONE - Implement `block_on`
//! - DONE - Implement `increase_threads_to`
//! - DONE - Drop finished futures
//! - DONE - Handle task panic
//! - DONE - Add `stop_threads`, `allow_threads`, and `increase_threads_to`.
//! - DONE - Add tests
//! - DONE - Add docs
//! - DONE - Publish on crates.io
//! - DONE - Add an `#[async_test]` macro
//! - Add a stress test
//! - Add a benchmark.  See benchmarks in <https://crates.io/crates/executors>
//! - Make a version of the crate that uses
//!   unsafe [`once_cell`](https://crates.io/crates/once_cell)
//!   and unsafe [`RawWaker`](https://doc.rust-lang.org/std/task/struct.RawWaker.html)
//!   and builds with Rust `stable`.
//! - Add an `#[async_main]` macro
//!
//! # Release Process
//! 1. Edit `Cargo.toml` and bump version number.
//! 1. Run `./release.sh`
#![forbid(unsafe_code)]
#![feature(wake_trait)]

use core::cell::Cell;
use safina_threadpool::ThreadPool;
use std::future::Future;
use std::pin::Pin;
use std::sync::mpsc::SyncSender;
use std::sync::{Arc, Mutex, Weak};
use std::task::Poll;

thread_local! {
    pub static EXECUTOR_THREAD_POOL: Cell<Weak<ThreadPool>> = Cell::new(Weak::new());
}

struct OnDrop<R, F: FnOnce() -> R>(Option<F>);
impl<R, F: FnOnce() -> R> OnDrop<R, F> {
    pub fn new(f: F) -> Self {
        Self(Some(f))
    }
}
impl<R, F: FnOnce() -> R> Drop for OnDrop<R, F> {
    fn drop(&mut self) {
        if let Some(f) = self.0.take() {
            f();
        }
    }
}

pub struct Executor {
    pool: Arc<ThreadPool>,
}

impl Executor {
    /// Creates a new executor with the name `async` and 4 threads.
    #[must_use]
    pub fn default() -> Self {
        Self::new(4)
    }

    /// Creates a new executor with the name `async`.
    ///
    /// You probably want to use [`default`](#method.default) instead of this.
    #[must_use]
    pub fn new(num_threads: usize) -> Self {
        Self::with_name("async", num_threads)
    }

    // Creates a new executor.
    #[must_use]
    pub fn with_name(name: &'static str, num_threads: usize) -> Self {
        Self {
            pool: Arc::new(ThreadPool::new(name, num_threads)),
        }
    }

    /// Adds a task that will execute `fut`.
    /// The task runs on any available worker thread.
    /// The task runs until `fut` completes or the Executor is dropped.
    ///
    /// Returns immediately.
    ///
    /// Uses
    /// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
    /// to make the future
    /// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
    /// You can use [`spawn_unpin`](#method.spawn_unpin) to avoid this allocation.
    ///
    /// Example:
    /// ```rust
    /// # async fn an_async_fn() -> Result<(), std::io::Error> { Ok(()) }
    /// # fn f() {
    /// let executor = safina_executor::Executor::default();
    /// executor.spawn(async move {
    ///     an_async_fn().await.unwrap();
    /// });
    /// # }
    /// ```
    pub fn spawn(&self, fut: impl (Future<Output = ()>) + Send + 'static) {
        self.spawn_unpin(Box::pin(fut))
    }

    /// Adds a task that will execute `fut`.
    /// The task runs on any available worker thread.
    /// The task runs until `fut` completes or the Executor is dropped.
    ///
    /// Returns immediately.
    ///
    /// Note that `fut` must be
    /// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
    /// You can use
    /// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
    /// to make it Unpin.  The [`spawn`](#method.spawn) function does this for you.
    /// Or use [`pin_utils::pin_mut`](https://docs.rs/pin-utils/latest/pin_utils/macro.pin_mut.html)
    /// to do it with unsafe code that does not allocate memory.
    pub fn spawn_unpin(&self, fut: impl (Future<Output = ()>) + Send + Unpin + 'static) {
        let task = Arc::new(Mutex::new(fut));
        let pool = Arc::downgrade(&self.pool);
        self.pool.schedule(move || poll_task(task, pool));
    }

    /// Executes the future on the current thread and returns its result.
    /// Panics if the future panics.
    ///
    /// `fut` can call [`spawn`] to create tasks.
    /// Those tasks run on the executor and will continue even after
    /// `fut` completes and this call returns.
    ///
    /// Uses
    /// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
    /// to make the future
    /// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
    /// You can use [`block_on_unpin`](#method.block_on_unpin) to avoid this allocation.
    ///
    /// ```rust
    /// # async fn prepare_request() -> Result<(), std::io::Error> { Ok(()) }
    /// # async fn execute_request() -> Result<(), std::io::Error> { Ok(()) }
    /// # fn f() -> Result<(), std::io::Error> {
    /// let executor = safina_executor::Executor::default();
    /// let result = executor.block_on(async {
    ///     prepare_request().await?;
    ///     execute_request().await
    /// })?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn block_on<R>(&self, fut: impl (Future<Output = R>) + Send + 'static) -> R {
        self.block_on_unpin(Box::pin(fut))
    }

    /// Executes the future on the current thread and returns its result.
    /// Panics if the future panics.
    ///
    /// `fut` can call [`spawn`] to create tasks.
    /// Those tasks run on the executor and will continue even after
    /// `fut` completes and this call returns.
    ///
    /// Note that `fut` must be
    /// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
    /// You can use
    /// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
    /// to make it Unpin.  The [`block_on`](#method.block_on) function does this for you.
    /// Or use [`pin_utils::pin_mut`](https://docs.rs/pin-utils/latest/pin_utils/macro.pin_mut.html)
    /// to do it with unsafe code that does not allocate memory.
    pub fn block_on_unpin<R>(&self, fut: impl (Future<Output = R>) + Send + Unpin + 'static) -> R {
        EXECUTOR_THREAD_POOL.with(|cell| cell.set(Arc::downgrade(&self.pool)));
        let _guard = OnDrop::new(|| EXECUTOR_THREAD_POOL.with(std::cell::Cell::take));
        block_on_unpin(fut)
    }
}

impl Default for Executor {
    fn default() -> Self {
        Executor::default()
    }
}

fn poll_task(task: Arc<Mutex<dyn Future<Output = ()> + Send + Unpin>>, pool: Weak<ThreadPool>) {
    if pool.strong_count() > 0 {
        let pool_clone = pool.clone();
        let task_clone = task.clone();
        let waker = std::task::Waker::from(Arc::new(TaskWaker::new(task, pool)));
        let mut cx = std::task::Context::from_waker(&waker);
        let mut fut_guard = task_clone.lock().unwrap();
        EXECUTOR_THREAD_POOL.with(|cell| cell.set(pool_clone));
        let _guard = OnDrop::new(|| EXECUTOR_THREAD_POOL.with(std::cell::Cell::take));
        let _ = Pin::new(&mut *fut_guard).poll(&mut cx);
    }
}

// TODO(mleonhard) Make TaskWaker and its clones single-use.
struct TaskWaker {
    task: Arc<Mutex<dyn Future<Output = ()> + Send + Unpin>>,
    pool: Weak<ThreadPool>,
}
impl TaskWaker {
    pub fn new(
        task: Arc<Mutex<dyn Future<Output = ()> + Send + Unpin>>,
        pool: Weak<ThreadPool>,
    ) -> Self {
        Self { task, pool }
    }
}
impl std::task::Wake for TaskWaker {
    fn wake(self: Arc<Self>) {
        if let Some(ref pool) = self.pool.upgrade() {
            let task_clone = self.task.clone();
            let pool_weak = Arc::downgrade(pool);
            pool.schedule(move || poll_task(task_clone, pool_weak));
        }
    }
}

/// Creates a new task to execute `fut` and schedules it for immediate execution.
///
/// Panics if the caller is not running on an [`Executor`](struct.Executor.html).
///
/// Returns immediately.
///
/// Uses
/// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
/// to make the future
/// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
/// You can use [`spawn_unpin`](#method.spawn_unpin) to avoid this allocation.
///
/// Example:
/// ```rust
/// # async fn an_async_fn() -> Result<(), std::io::Error> { Ok(()) }
/// # fn f() {
/// let executor = safina_executor::Executor::default();
/// executor.spawn(async move {
///     safina_executor::spawn(async move {
///         an_async_fn().await.unwrap();
///     });
/// });
/// # }
/// ```
pub fn spawn(fut: impl (Future<Output = ()>) + Send + 'static) {
    spawn_unpin(Box::pin(fut))
}

/// Creates a new task to execute `fut` and schedules it for immediate execution.
///
/// Panics if the caller is not running on an [`Executor`](struct.Executor.html).
///
/// Returns immediately.
///
/// Note that `fut` must be
/// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
/// You can use
/// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
/// to make it Unpin.  The [`spawn`](#method.spawn) function does this for you.
/// Or use [`pin_utils::pin_mut`](https://docs.rs/pin-utils/latest/pin_utils/macro.pin_mut.html)
/// to do it with unsafe code that does not allocate memory.
pub fn spawn_unpin(fut: impl (Future<Output = ()>) + Send + Unpin + 'static) {
    let option_arc_pool = EXECUTOR_THREAD_POOL.with(|cell| {
        let weak = cell.take();
        let result = weak.upgrade();
        cell.set(weak);
        result
    });
    if let Some(pool) = option_arc_pool {
        let task = Arc::new(Mutex::new(fut));
        let pool_weak = Arc::downgrade(&pool);
        pool.schedule(move || poll_task(task, pool_weak));
    } else {
        panic!("safina_executor::spawn_unpin called from outside a task")
    }
}

/// Executes the future on the current thread and returns its result.
/// Panics if the future panics.
///
/// Panics if `fut` calls [`spawn`](fn.spawn.html) to create a new task.
/// Does not create an executor.
/// Use [`Executor::block_on`](struct.Executor.html#method.block_on)
/// if you need `fut` to spawn new tasks.
///
/// Uses
/// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
/// to make the future
/// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
/// You can use [`block_on_unpin`](#method.block_on_unpin) to avoid this allocation.
///
/// ```rust
/// # async fn prepare_request() -> Result<(), std::io::Error> { Ok(()) }
/// # async fn execute_request() -> Result<(), std::io::Error> { Ok(()) }
/// # fn f() -> Result<(), std::io::Error> {
/// let result = safina_executor::block_on(async {
///     prepare_request().await?;
///     execute_request().await
/// })?;
/// # Ok(())
/// # }
/// ```
pub fn block_on<R>(fut: impl (Future<Output = R>) + Send + 'static) -> R {
    block_on_unpin(Box::pin(fut))
}

/// Executes the future on the current thread and returns its result.
/// Panics if the future panics.
///
/// Panics if `fut` calls [`spawn`](fn.spawn.html) to create a new task.
/// Does not create an executor.
/// Use [`Executor::block_on`](struct.Executor.html#method.block_on)
/// if you need `fut` to spawn new tasks.
///
/// Note that `fut` must be
/// [`Unpin`](https://doc.rust-lang.org/stable/core/marker/trait.Unpin.html).
/// You can use
/// [`std::boxed::Box::pin`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.pin)
/// to make it Unpin.  The [`block_on`](#method.block_on) function does this for you.
/// Or use [`pin_utils::pin_mut`](https://docs.rs/pin-utils/latest/pin_utils/macro.pin_mut.html)
/// to do it with unsafe code that does not allocate memory.
pub fn block_on_unpin<R>(mut fut: impl (Future<Output = R>) + Send + Unpin + 'static) -> R {
    struct BlockOnTaskWaker(SyncSender<()>);
    impl std::task::Wake for BlockOnTaskWaker {
        fn wake(self: Arc<Self>) {
            let _ = self.0.send(());
        }
    }
    let (sender, receiver) = std::sync::mpsc::sync_channel(1);
    let waker = std::task::Waker::from(Arc::new(BlockOnTaskWaker(sender)));
    let mut cx = std::task::Context::from_waker(&waker);
    loop {
        if let Poll::Ready(result) = Pin::new(&mut fut).poll(&mut cx) {
            return result;
        }
        receiver.recv().unwrap();
    }
}