pipe_it/

lib.rs

use std::{
    any::{self, TypeId},
    collections::HashMap,
    future::Future,
    ops::{Deref, DerefMut},
    sync::Arc,
};

use tokio::sync;

pub use pipeit_derive::node;

pub extern crate pipeit_derive;

pub mod cocurrency;
pub mod ext;
pub mod handler;
pub mod tag;
// store dynamic data
/// A map storing shared dependencies, keyed by their TypeId.
/// Wrapped in an Arc for thread-safe sharing and using Arc<RwLock> to support owned guards.
#[derive(Clone)]
pub(crate) struct DendencyMap(
    Arc<HashMap<TypeId, Arc<sync::RwLock<Box<dyn any::Any + Send + Sync>>>>>,
);

/// A collection of shared resources that can be injected into pipeline functions.
///
/// This acts as a builder to conveniently register resources before creating a Context.
#[derive(Default, Clone)]
pub struct Shared {
    inner: HashMap<TypeId, Arc<sync::RwLock<Box<dyn any::Any + Send + Sync>>>>,
}

impl Shared {
    /// Creates a new, empty shared resource collection.
    pub fn new() -> Self {
        Self::default()
    }

    /// Adds a new resource to the collection.
    /// The resource will be wrapped in an Arc<RwLock> to allow concurrent access.
    pub fn insert<T: Send + Sync + 'static>(mut self, resource: T) -> Self {
        self.inner.insert(
            TypeId::of::<T>(),
            Arc::new(sync::RwLock::new(
                Box::new(resource) as Box<dyn any::Any + Send + Sync>
            )),
        );
        self
    }
}

/// A read-only resource container that mirrors Bevy's Res.
/// It holds an owned read lock on a shared dependency and provides strong typing via Deref.
pub struct Res<T>(
    sync::OwnedRwLockReadGuard<Box<dyn any::Any + Send + Sync>>,
    std::marker::PhantomData<T>,
);

/// A mutable resource container that mirrors Bevy's ResMut.
/// It holds an owned write lock on a shared dependency and provides strong typing via Deref/DerefMut.
pub struct ResMut<T>(
    sync::OwnedRwLockWriteGuard<Box<dyn any::Any + Send + Sync>>,
    std::marker::PhantomData<T>,
);

impl<T: 'static> Deref for Res<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        // Here we perform the downcast from dyn Any to the strong type T.
        // This is exactly how Bevy provides strong typing from a generic container.
        (**self.0)
            .downcast_ref::<T>()
            .expect("Resource type mismatch during Res deref")
    }
}

impl<T: 'static> Deref for ResMut<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        (**self.0)
            .downcast_ref::<T>()
            .expect("Resource type mismatch during ResMut deref")
    }
}

impl<T: 'static> DerefMut for ResMut<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        // Performs mutable downcast to provide &mut T.
        (**self.0)
            .downcast_mut::<T>()
            .expect("Resource type mismatch during ResMut deref_mut")
    }
}

/// The execution context for a pipeline, carrying both shared dependencies and the specific input.
pub struct Context<Input>
where
    Input: ?Sized,
{
    shared: DendencyMap,
    input: Arc<Input>,
}

impl<I> Context<I> {
    /// Creates a new context with the provided input and shared dependency collection.
    pub fn new(input: I, shared: Shared) -> Self {
        Self {
            shared: DendencyMap(Arc::new(shared.inner)),
            input: Arc::new(input),
        }
    }
    pub fn empty(input: I) -> Self {
        Self {
            shared: DendencyMap(Arc::new(Shared::new().inner)),
            input: Arc::new(input),
        }
    }
    /// Replaces the input of the current context while keeping the same shared dependencies.
    /// Returns a new Context with the updated input type.
    pub(crate) fn replace<NewInput>(self, input: NewInput) -> Context<NewInput> {
        Context {
            shared: self.shared,
            input: Arc::new(input),
        }
    }

    /// Returns a reference to the current input wrapped in an Arc.
    pub(crate) fn input(&self) -> Arc<I> {
        self.input.clone()
    }

    /// Consumes the Context and returns the input Arc and the shared dependencies.
    /// This allows avoiding the clone of the input if the Context is no longer needed.
    pub(crate) fn into_parts(self) -> (Arc<I>, DendencyMap) {
        (self.input, self.shared)
    }

    /// Reconstructs a Context from its parts.
    pub(crate) fn from_parts(input: Arc<I>, shared: DendencyMap) -> Self {
        Self { shared, input }
    }
    /// # Example
    ///
    /// ```rust
    /// use pipe_it::{Context, Input, Res, ResMut, Shared};
    ///
    /// #[derive(Debug, Clone)]
    /// struct Counter {
    ///     c: i32,
    /// }
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let ctx = Context::new(3, Shared::new().insert(Counter { c: 1 }));
    ///     ctx.invoke(async |x: Input<i32>, mut counter: ResMut<Counter>| {
    ///         counter.c += 1;
    ///         *x + 1
    ///     })
    ///     .await
    ///     .invoke(async |x: Input<i32>, counter: Res<Counter>| *x + counter.c)
    ///     .await
    ///     .invoke(async |x: Input<i32>| assert_eq!(*x, 6))
    ///     .await;
    /// }
    /// ```
    pub async fn invoke<H, Args, O>(self, handler: H) -> Context<O>
    where
        H: handler::Handler<I, O, Args>,
        I: Clone + Send + Sync + 'static,
        O: Send + Sync + 'static,
        Args: Send + Sync + 'static,
    {
        let (input, shared) = self.into_parts();
        let shared_cloned = shared.clone();
        // move input to avoid increasing ref counter
        let new_ctx = Context {
            input,
            shared: shared_cloned,
        };
        let output = handler.call(new_ctx).await;
        Context {
            shared: shared,
            input: Arc::new(output),
        }
    }
}
impl<Input> Clone for Context<Input>
where
    Input: ?Sized,
{
    fn clone(&self) -> Self {
        Self {
            shared: self.shared.clone(),
            input: self.input.clone(),
        }
    }
}

/// A trait for types that can be extracted from a Context.
/// Similar to Axum's FromRequest or Bevy's SystemParam.
pub trait FromContext<Input>: Send {
    fn from(ctx: Context<Input>) -> impl Future<Output = Self> + Send;
}

/// A wrapper for the primary input data of the pipeline.
#[derive(Clone)]
pub struct Input<T>(pub Arc<T>);

impl<T> Input<T> {
    /// Returns the inner value, if the Arc has exactly one strong reference.
    /// Otherwise, an Err is returned with the same Arc that was passed in.
    /// // Optimized: consume context to steal Arc without cloning if possible
    pub fn try_unwrap(self) -> Result<T, Arc<T>> {
        Arc::try_unwrap(self.0)
    }
}
impl<T> Deref for Input<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<I: Send + Sync + 'static> FromContext<I> for Input<I> {
    fn from(ctx: Context<I>) -> impl Future<Output = Self> + Send {
        let input = ctx.input.clone();
        async move { Input(input) }
    }
}

impl<I, T> FromContext<I> for Res<T>
where
    I: Send + Sync + 'static,
    T: Send + Sync + 'static,
{
    fn from(ctx: Context<I>) -> impl Future<Output = Self> + Send {
        async move {
            let shared = ctx.shared.0.clone();
            let dep = shared
                .get(&TypeId::of::<T>())
                .expect("Dependency not found")
                .clone();
            let guard = dep.read_owned().await;
            Res(guard, std::marker::PhantomData)
        }
    }
}

impl<I, T> FromContext<I> for ResMut<T>
where
    I: Send + Sync + 'static,
    T: Send + Sync + 'static,
{
    fn from(ctx: Context<I>) -> impl Future<Output = Self> + Send {
        async move {
            let shared = ctx.shared.0.clone();
            let dep = shared
                .get(&TypeId::of::<T>())
                .expect("Dependency not found")
                .clone();
            let guard = dep.write_owned().await;
            ResMut(guard, std::marker::PhantomData)
        }
    }
}

/// Represents a pipeline unit that can be applied to a Context.
/// Implementations are automatically provided for functions that match the signature.
pub trait Pipeline<I, O>: Send + Sync + 'static {
    fn apply(&self, ctx: Context<I>) -> impl Future<Output = O> + Send;
}

/// Errors occurring during pipeline execution.
#[derive(Debug, Clone)]
pub enum PipelineError {
    Failure { msg: String, expected: String },
    Fatal { msg: String },
}

/// Standard Result type for pipeline operations.
pub type PResult<O, E = PipelineError> = Result<O, E>;

/// A pipeline that executes only if a predicate on the input is met.
pub struct Cond<F, P, I, O> {
    predicate: F,
    next: P,
    _marker: std::marker::PhantomData<fn(I, O)>,
}

impl<F, P, I, O> Pipeline<I, PResult<O>> for Cond<F, P, I, O>
where
    F: Pipeline<I, bool>,
    P: Pipeline<I, O>,
    I: Clone + Send + Sync + 'static,
    O: Send + 'static,
    F: Send + Sync + 'static,
    P: Send + Sync + 'static,
{
    fn apply(&self, ctx: Context<I>) -> impl Future<Output = PResult<O>> + Send {
        async move {
            let matched = self.predicate.apply(ctx.clone()).await;
            if matched {
                Ok(self.next.apply(ctx).await)
            } else {
                Err(PipelineError::Failure {
                    msg: "Condition not met".to_string(),
                    expected: "true".to_string(),
                })
            }
        }
    }
}

/// Creates a conditional pipeline.
/// If the predicate returns true, the next pipeline is executed.
/// If false, it returns a PipelineError::Failure.
///
/// # Example
///
/// ```rust
/// use pipe_it::{cond, Context, Pipeline, Input, ext::HandlerExt};
///
/// async fn is_even(n: Input<i32>) -> bool { *n % 2 == 0 }
/// async fn process(n: Input<i32>) -> String { "Even".to_string() }
///
/// # #[tokio::main]
/// # async fn main() {
/// let pipe = cond(is_even, process);
///
/// // Success case
/// let result = pipe.apply(Context::empty(2)).await;
/// assert_eq!(result.unwrap(), "Even");
///
/// // Failure case
/// let result = pipe.apply(Context::empty(1)).await;
/// assert!(result.is_err());
/// # }
/// ```
pub fn cond<I, O, F, P, ArgsF, ArgsP>(
    predicate: F,
    next: P,
) -> Cond<crate::ext::Pipe<F, ArgsF>, crate::ext::Pipe<P, ArgsP>, I, O>
where
    F: crate::handler::Handler<I, bool, ArgsF>,
    P: crate::handler::Handler<I, O, ArgsP>,
    I: Clone + Send + Sync + 'static,
    O: Send + 'static,
    ArgsF: Send + Sync + 'static,
    ArgsP: Send + Sync + 'static,
{
    use crate::ext::HandlerExt;
    Cond {
        predicate: predicate.pipe(),
        next: next.pipe(),
        _marker: std::marker::PhantomData,
    }
}

/// A wrapper used as a marker for choice-based pipeline execution.
pub struct Choice<T>(pub T);

macro_rules! impl_pipeline_for_tuple {
    ($($P:ident),+) => {
        impl<I, O, $($P),+ > Pipeline<I, PResult<O>> for ($($P,)+)
        where
            I: Clone + Send + Sync + 'static,
            O: Send + 'static,
            $( $P: Pipeline<I, PResult<O>> ),+
        {
            fn apply(&self, ctx: Context<I>) -> impl Future<Output = PResult<O>> + Send {
                #[allow(non_snake_case)]
                let ($($P,)+) = self;
                async move {
                    $(
                        match $P.apply(ctx.clone()).await {
                            Ok(res) => return Ok(res),
                            Err(PipelineError::Fatal { msg }) => return Err(PipelineError::Fatal { msg }),
                            Err(PipelineError::Failure { .. }) => {}
                        }
                    )*
                    Err(PipelineError::Fatal { msg: "All pipeline branches failed".to_string() })
                }
            }
        }
    };
}

impl_pipeline_for_tuple!(P1);
impl_pipeline_for_tuple!(P1, P2);
impl_pipeline_for_tuple!(P1, P2, P3);
impl_pipeline_for_tuple!(P1, P2, P3, P4);
impl_pipeline_for_tuple!(P1, P2, P3, P4, P5);
impl_pipeline_for_tuple!(P1, P2, P3, P4, P5, P6);
impl_pipeline_for_tuple!(P1, P2, P3, P4, P5, P6, P7);
impl_pipeline_for_tuple!(P1, P2, P3, P4, P5, P6, P7, P8);

/// An identity pipeline that simply returns the current input.
/// This acts as a neutral element in pipeline composition.
pub async fn identity<I: Clone>(input: Input<I>) -> I {
    (*input).clone()
}

/// An identity pipeline that returns the current input wrapped in a successful Result.
/// Useful as a terminal fallback in a choice-based pipeline (tuple).
pub async fn identity_res<I: Clone>(input: Input<I>) -> PResult<I> {
    Ok((*input).clone())
}

#[cfg(test)]
mod tests {
    use crate::{Context, Input, Res, ResMut, Shared};
    #[derive(Debug, Clone)]
    struct Counter {
        c: i32,
    }
    #[tokio::test]
    // test context invoke api
    async fn test_chain_invoke() {
        let ctx = Context::new(3, Shared::new().insert(Counter { c: 1 }));
        ctx.invoke(async |x: Input<i32>, mut counter: ResMut<Counter>| {
            counter.c += 1;
            *x + 1
        })
        .await
        .invoke(async |x: Input<i32>, counter: Res<Counter>| *x + counter.c)
        .await
        .invoke(async |x: Input<i32>| assert_eq!(*x, 6))
        .await;
    }
}