corophage 0.4.0

use std::future::Future;
use std::marker::PhantomData;
use std::ops::Add;

use frunk_core::coproduct::{CNil, CoproductSubsetter};
use frunk_core::hlist::{HCons, HNil};

use crate::control::Cancelled;
use crate::coproduct::{AsyncHandleMut, AsyncHandleWith, HandleMut, HandleWith, HandlersToEffects};
use crate::coroutine::{Co, CoSend, GenericCo, Yielder};
use crate::effect::{CanStart, Effects, Resumes};
use crate::locality::{Local, Locality, Sendable};

/// An effectful computation with no handlers attached.
///
/// This is a convenience alias for [`Program`] where all effects in `Effs`
/// still need to be handled. Use this as the return type of functions that
/// create programs:
///
/// ```ignore
/// use corophage::prelude::*;
///
/// type MyEffs = Effects![Counter, Ask];
///
/// fn my_computation<'a>() -> Effectful<'a, MyEffs, String> {
///     Program::new(|y: Yielder<'_, MyEffs>| async move {
///         // ...
///         "result".to_string()
///     })
/// }
/// ```
pub type Effectful<'a, Effs, R, L = Local> = Program<'a, Effs, R, L, Effs, HNil>;

/// Deprecated alias for [`Effectful`].
#[deprecated(since = "0.3.0", note = "renamed to `Effectful`")]
pub type Eff<'a, Effs, R, L = Local> = Effectful<'a, Effs, R, L>;

/// A computation with incrementally attached effect handlers.
///
/// Handlers are added one at a time via [`handle`](Program::handle),
/// or in bulk via [`handle_all`](Program::handle_all).
/// Once all effects are handled (`Remaining = CNil`), the computation
/// can be executed via [`run`](Program::run) or [`run_sync`](Program::run_sync).
pub struct Program<'a, Effs: Effects<'a>, Result, L: Locality, Remaining, Handlers> {
    pub(crate) co: GenericCo<'a, Effs, Result, L>,
    handlers: Handlers,
    _remaining: PhantomData<Remaining>,
}

impl<'a, Result> Program<'a, CNil, Result, Local, CNil, HNil> {
    /// Create a new program from a computation closure.
    ///
    /// The effect set can be inferred from the closure parameter annotation,
    /// or specified explicitly via turbofish:
    ///
    /// ```ignore
    /// // With closure annotation:
    /// Program::new(|y: Yielder<'_, Effs>| async move { ... })
    ///
    /// // With turbofish:
    /// Program::new::<Effects![Counter, Ask], _>(|y| async move { ... })
    /// ```
    pub fn new<Effs, F>(
        f: impl FnOnce(Yielder<'a, Effs>) -> F + 'a,
    ) -> Program<'a, Effs, Result, Local, Effs, HNil>
    where
        Effs: Effects<'a>,
        F: Future<Output = Result>,
    {
        Program::from_co(Co::new(f))
    }
}

impl<'a, Result> Program<'a, CNil, Result, Sendable, CNil, HNil> {
    /// Create a new `Send`-able program from a computation closure.
    ///
    /// The effect set can be inferred from the closure parameter annotation,
    /// or specified explicitly via turbofish:
    ///
    /// ```ignore
    /// // With closure annotation:
    /// Program::new_send(|y: Yielder<'_, Effs>| async move { ... })
    ///
    /// // With turbofish:
    /// Program::new_send::<Effects![Counter, Ask], _>(|y| async move { ... })
    /// ```
    pub fn new_send<Effs, F>(
        f: impl FnOnce(Yielder<'a, Effs>) -> F + Send + 'a,
    ) -> Program<'a, Effs, Result, Sendable, Effs, HNil>
    where
        Effs: Effects<'a>,
        for<'r> Resumes<'r, CanStart<Effs>>: Send + Sync,
        F: Future<Output = Result> + Send,
    {
        Program::from_co(CoSend::new(f))
    }
}

impl<'a, Effs, R, L> Program<'a, Effs, R, L, Effs, HNil>
where
    Effs: Effects<'a>,
    L: Locality,
{
    /// Create a program from an existing coroutine.
    #[inline]
    pub fn from_co(co: GenericCo<'a, Effs, R, L>) -> Self {
        Program {
            co,
            handlers: HNil,
            _remaining: PhantomData,
        }
    }
}

type HandleEffects<'a, Remaining, H, Effs, HandleIdx, SubsetIdx> =
    <Remaining as CoproductSubsetter<
        <H as HandlersToEffects<'a, Effs, HandleIdx>>::Effects,
        SubsetIdx,
    >>::Remainder;

impl<'a, Effs, R, L, Remaining, Handlers> Program<'a, Effs, R, L, Remaining, Handlers>
where
    Effs: Effects<'a>,
    L: Locality,
{
    /// Attach a handler for an unhandled effect.
    ///
    /// Handlers can be attached in any order — the effect type is
    /// inferred from the closure signature and removed from the
    /// remaining set via [`CoproductSubsetter`].
    #[allow(clippy::type_complexity)]
    #[inline]
    pub fn handle<F, HandleIdx, SubsetIdx>(
        self,
        handler: F,
    ) -> Program<
        'a,
        Effs,
        R,
        L,
        HandleEffects<'a, Remaining, HCons<F, HNil>, Effs, HandleIdx, SubsetIdx>,
        <Handlers as Add<HCons<F, HNil>>>::Output,
    >
    where
        HCons<F, HNil>: HandlersToEffects<'a, Effs, HandleIdx>,
        Remaining: CoproductSubsetter<
                <HCons<F, HNil> as HandlersToEffects<'a, Effs, HandleIdx>>::Effects,
                SubsetIdx,
            >,
        Handlers: Add<HCons<F, HNil>>,
    {
        self.handle_all(HCons {
            head: handler,
            tail: HNil,
        })
    }

    /// Attach multiple handlers at once from an HList.
    ///
    /// The handlers can be for any subset of the remaining effects,
    /// in any order. The effect types are inferred from the handler
    /// closure signatures, and `CoproductSubsetter` removes those
    /// effects from the `Remaining` set.
    #[allow(clippy::type_complexity)]
    #[inline]
    pub fn handle_all<H, HandleIdx, SubsetIdx>(
        self,
        handlers: H,
    ) -> Program<
        'a,
        Effs,
        R,
        L,
        HandleEffects<'a, Remaining, H, Effs, HandleIdx, SubsetIdx>,
        <Handlers as Add<H>>::Output,
    >
    where
        H: HandlersToEffects<'a, Effs, HandleIdx>,
        Remaining: CoproductSubsetter<H::Effects, SubsetIdx>,
        Handlers: Add<H>,
    {
        Program {
            co: self.co,
            handlers: self.handlers + handlers,
            _remaining: PhantomData,
        }
    }
}

impl<'a, Effs, R, L, Handlers> Program<'a, Effs, R, L, CNil, Handlers>
where
    Effs: Effects<'a>,
    L: Locality,
{
    /// Run the computation synchronously.
    #[inline]
    pub fn run_sync<Indices>(self) -> Result<R, Cancelled>
    where
        Effs: HandleMut<'a, Effs, Handlers, Indices>,
    {
        let mut handlers = self.handlers;
        crate::sync::run(self.co, &mut handlers)
    }

    /// Run the computation synchronously with shared state.
    #[inline]
    pub fn run_sync_stateful<S, Indices>(self, state: &mut S) -> Result<R, Cancelled>
    where
        Effs: HandleWith<'a, Effs, Handlers, S, Indices>,
    {
        let handlers = self.handlers;
        crate::sync::run_stateful(self.co, state, &handlers)
    }

    /// Run the computation asynchronously.
    #[inline]
    pub async fn run<Indices>(self) -> Result<R, Cancelled>
    where
        Effs: AsyncHandleMut<'a, Effs, Handlers, Indices>,
    {
        let mut handlers = self.handlers;
        crate::asynk::run(self.co, &mut handlers).await
    }

    /// Run the computation asynchronously with shared state.
    #[inline]
    pub async fn run_stateful<S, Indices>(self, state: &mut S) -> Result<R, Cancelled>
    where
        Effs: AsyncHandleWith<'a, Effs, Handlers, S, Indices>,
    {
        let handlers = self.handlers;
        crate::asynk::run_stateful(self.co, state, &handlers).await
    }
}