pipe_it/

ext.rs

use crate::{Context, handler::Handler, Pipeline};
use std::future::Future;
use std::marker::PhantomData;

/// Extension trait for any type that implements the Handler trait.
/// Provides a way to convert a Handler into a Pipeline and combinators.
pub trait HandlerExt<I, O, Args>: Handler<I, O, Args> {
    fn pipe(self) -> Pipe<Self, Args>
    where
        Self: Sized,
    {
        Pipe {
            p: self,
            _marker: PhantomData,
        }
    }

    /// Connects two pipelines together. Output of the first becomes the input of the second.
    ///
    /// # Example
    ///
    /// ```rust
    /// use pipe_it::{Context, Pipeline, Input, ext::HandlerExt};
    ///
    /// async fn add_one(n: Input<i32>) -> i32 { *n + 1 }
    /// async fn to_string(n: Input<i32>) -> String { n.to_string() }
    /// 
    /// # #[tokio::main]
    /// # async fn main() {
    /// let pipe = add_one.pipe().connect(to_string);
    /// let result = pipe.apply(Context::empty(1)).await;
    /// assert_eq!(result, "2");
    /// # }
    /// ```
    fn connect<O2, G, Args2>(self, g: G) -> Connect<Pipe<Self, Args>, Pipe<G, Args2>, I, O, O2>
    where
        G: Handler<O, O2, Args2>,
        Self: Sized,
    {
        Connect {
            f: self.pipe(),
            g: g.pipe(),
            _marker: PhantomData,
        }
    }

    /// Pulls back the domain of the pipeline.
    /// This allows a pipeline defined on `I` to be used for input `I2` given a mapping `I2 -> I`.
    /// The mapping function is now a full Handler (async, supports DI).
    ///
    /// # Example
    ///
    /// ```rust
    /// use pipe_it::{Context, Pipeline, Input, ext::HandlerExt};
    ///
    /// #[derive(Clone)]
    /// struct User { age: i32 }
    /// 
    /// // Pre-processing step: extract age from User.
    /// // Supports DI and ownership extraction.
    /// async fn get_age(user: Input<User>) -> i32 { 
    ///     // We can try_unwrap to get ownership if needed (optimization)
    ///     user.try_unwrap().map(|u| u.age).unwrap_or_else(|u| u.age)
    /// }
    ///
    /// async fn check_adult(age: Input<i32>) -> bool { *age >= 18 }
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// // Input: User -> (get_age) -> i32 -> (check_adult) -> bool
    /// let pipe = check_adult.pipe().pullback(get_age);
    /// let result = pipe.apply(Context::empty(User { age: 20 })).await;
    /// assert_eq!(result, true);
    /// # }
    /// ```
    fn pullback<I2, H, Args2>(self, handler: H) -> Pullback<Pipe<Self, Args>, Pipe<H, Args2>, I2, I, O>
    where
        H: Handler<I2, I, Args2>,
        Self: Sized,
    {
        Pullback {
            p: self.pipe(),
            map: handler.pipe(),
            _marker: PhantomData,
        }
    }

    /// Lifts the domain requirement to anything that can be converted into the original input.
    /// Uses `From` / `Into` trait.
    ///
    /// # Example
    ///
    /// ```rust
    /// use pipe_it::{Context, Pipeline, Input, ext::HandlerExt};
    ///
    /// async fn process_string(s: Input<String>) -> usize { s.len() }
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// // Accepts &str because String implements From<&str>
    /// let pipe = process_string.pipe().lift::<&str>();
    /// let result = pipe.apply(Context::empty("hello")).await;
    /// assert_eq!(result, 5);
    /// # }
    /// ```
    fn lift<I2>(self) -> Pullback<Pipe<Self, Args>, Pipe<LiftHandler<I, I2>, (crate::Input<I2>,)>, I2, I, O>
    where
        I: From<I2> + Send + Sync + 'static,
        I2: Clone + Send + Sync + 'static,
        Self: Sized,
    {
        self.pullback(LiftHandler(PhantomData))
    }

    /// Extends the output of the pipeline by applying a transformation.
    ///
    /// # Example
    ///
    /// ```rust
    /// use pipe_it::{Context, Pipeline, Input, ext::HandlerExt};
    ///
    /// async fn compute(n: Input<i32>) -> i32 { *n * 2 }
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// // Changes output from i32 to String
    /// let pipe = compute.pipe().extend(|n| format!("Result: {}", n));
    /// let result = pipe.apply(Context::empty(10)).await;
    /// assert_eq!(result, "Result: 20");
    /// # }
    /// ```
    fn extend<O2, F>(self, map: F) -> Extend<Pipe<Self, Args>, F, I, O, O2>
    where
        F: Fn(O) -> O2 + Send + Sync + 'static,
        Self: Sized,
    {
        Extend {
            p: self.pipe(),
            map,
            _marker: PhantomData,
        }
    }

    /// Repeats the pipeline operation n times.
    /// Input and output types must be the same.
    ///
    /// # Example
    ///
    /// ```rust
    /// use pipe_it::{Context, Pipeline, Input, ext::HandlerExt};
    ///
    /// async fn add_one(n: Input<i32>) -> i32 { *n + 1 }
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// // Input: 0 -> 1 -> 2 -> 3
    /// let pipe = add_one.pipe().repeat(3);
    /// let result = pipe.apply(Context::empty(0)).await;
    /// assert_eq!(result, 3);
    /// # }
    /// ```
    fn repeat(self, times: usize) -> Repeat<Pipe<Self, Args>, I>
    where
        Self: Handler<I, I, Args> + Sized,
        I: Clone + Send + Sync + 'static,
    {
        Repeat {
            p: self.pipe(),
            times,
            _marker: PhantomData,
        }
    }

    /// Caches the results of the pipeline using an LRU strategy.
    /// Requires the input to implement `Hash + Eq + Clone` and output to implement `Clone`.
    ///
    /// # Example
    ///
    /// ```rust
    /// use pipe_it::{Context, Pipeline, Input, ext::HandlerExt};
    /// use std::sync::atomic::{AtomicUsize, Ordering};
    /// use std::sync::Arc;
    ///
    /// static CALL_COUNT: AtomicUsize = AtomicUsize::new(0);
    ///
    /// async fn heavy_calc(n: Input<i32>) -> i32 { 
    ///     CALL_COUNT.fetch_add(1, Ordering::SeqCst);
    ///     *n * *n 
    /// }
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let pipe = heavy_calc.pipe().cache(100);
    /// let result1 = pipe.apply(Context::empty(10)).await;
    /// assert_eq!(result1, 100);
    /// let result2 = pipe.apply(Context::empty(10)).await; // Should hit cache
    /// assert_eq!(result2, 100);
    /// assert_eq!(CALL_COUNT.load(Ordering::SeqCst), 1);
    /// # }
    /// ```
    fn cache(self, capacity: usize) -> Cache<Pipe<Self, Args>, I, O>
    where
        Self: Sized,
        I: std::hash::Hash + Eq + Clone + Send + Sync + 'static,
        O: Clone + Send + Sync + 'static,
    {
        Cache {
            p: self.pipe(),
            cache: quick_cache::sync::Cache::new(capacity),
        }
    }

    /// Maps the pipeline over a collection of inputs concurrently.
    /// Input for the new pipeline will be `Vec<I>` and output will be `Vec<O>`.
    ///
    /// # Example
    ///
    /// ```rust
    /// use pipe_it::{Context, Pipeline, Input, ext::HandlerExt};
    ///
    /// async fn process_item(n: Input<i32>) -> String { n.to_string() }
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// // Input: Vec<i32> -> Output: Vec<String>
    /// let pipe = process_item.pipe().map();
    /// let result = pipe.apply(Context::empty(vec![1, 2, 3])).await;
    /// assert_eq!(result, vec!["1", "2", "3"]);
    /// # }
    /// ```
    fn map(self) -> crate::cocurrency::Map<Pipe<Self, Args>, I, O>
    where
        Self: Sized,
    {
        crate::cocurrency::Map {
            p: self.pipe(),
            _marker: PhantomData,
        }
    }
}

impl<F, I, O, Args> HandlerExt<I, O, Args> for F where F: Handler<I, O, Args> {}

/// Wrapper struct to adapt a Handler into a Pipeline.
pub struct Pipe<P, Args> {
    p: P,
    _marker: PhantomData<Args>,
}

impl<P, Args, I, O> Pipeline<I, O> for Pipe<P, Args>
where
    P: Handler<I, O, Args>,
    I: Clone + Send + Sync + 'static,
    O: Send + 'static,
    Args: Send + Sync + 'static,
{
    fn apply(&self, ctx: Context<I>) -> impl Future<Output = O> + Send {
        self.p.call(ctx)
    }
}

// PipelineExt removed to avoid ambiguity with HandlerExt
// pub trait PipelineExt<I, O>: Pipeline<I, O> { ... }
// impl<P, I, O> PipelineExt<I, O> for P where P: Pipeline<I, O> {}

// --- Connect ---

pub struct Connect<F, G, I, O1, O2> {
    f: F,
    g: G,
    _marker: PhantomData<(I, O1, O2)>,
}

impl<F, G, I, O1, O2> Pipeline<I, O2> for Connect<F, G, I, O1, O2>
where
    F: Pipeline<I, O1>,
    G: Pipeline<O1, O2>,
    I: Clone + Send + Sync + 'static,
    O1: Send + Sync + 'static,
    O2: Send + Sync + 'static,
{
    fn apply(&self, ctx: Context<I>) -> impl Future<Output = O2> + Send {
        async move {
            // Optimized to release input ownership for the second stage
            let (input, shared) = ctx.into_parts();
            let ctx_f = Context::from_parts(input, shared.clone());
            let res = self.f.apply(ctx_f).await;
            self.g.apply(Context::from_parts(std::sync::Arc::new(res), shared)).await
        }
    }
}

// --- Pullback ---

pub struct Pullback<P, H, I2, I, O> {
    p: P,
    map: H,
    _marker: PhantomData<(I2, I, O)>,
}

impl<P, H, I2, I, O> Pipeline<I2, O> for Pullback<P, H, I2, I, O>
where
    P: Pipeline<I, O>,
    H: Pipeline<I2, I>,
    I2: Clone + Send + Sync + 'static,
    I: Send + Sync + 'static,
    O: Send + Sync + 'static,
{
    fn apply(&self, ctx: Context<I2>) -> impl Future<Output = O> + Send {
        async move {
            // Effectively H.connect(P)
            let (input, shared) = ctx.into_parts();
            // Run Mapper (H) on I2
            let ctx_h = Context::from_parts(input, shared.clone());
            let mapped_input = self.map.apply(ctx_h).await;
            // Run P on mapped input I
            self.p.apply(Context::from_parts(std::sync::Arc::new(mapped_input), shared)).await
        }
    }
}

/// Helper handler for Lift operation
pub struct LiftHandler<I, I2>(PhantomData<(I, I2)>);

impl<I, I2> Handler<I2, I, (crate::Input<I2>,)> for LiftHandler<I, I2>
where
    I: From<I2> + Send + Sync + 'static,
    I2: Clone + Send + Sync + 'static,
{
    fn call(&self, ctx: Context<I2>) -> impl Future<Output = I> + Send {
        async move {
            let (input, _) = ctx.into_parts();
            let val = match crate::Input::<I2>(input).try_unwrap() {
                Ok(v) => v,
                Err(arc) => (*arc).clone(),
            };
            val.into()
        }
    }
}

// --- Extend ---

pub struct Extend<P, F, I, O1, O2> {
    p: P,
    map: F,
    _marker: PhantomData<(I, O1, O2)>,
}

impl<P, F, I, O1, O2> Pipeline<I, O2> for Extend<P, F, I, O1, O2>
where
    P: Pipeline<I, O1>,
    F: Fn(O1) -> O2 + Send + Sync + 'static,
    I: Clone + Send + Sync + 'static,
    O1: Send + Sync + 'static,
    O2: Send + Sync + 'static,
{
    fn apply(&self, ctx: Context<I>) -> impl Future<Output = O2> + Send {
        async move {
            let res = self.p.apply(ctx).await;
            (self.map)(res)
        }
    }
}

// --- Repeat ---

pub struct Repeat<P, I> {
    p: P,
    times: usize,
    _marker: PhantomData<I>,
}

impl<P, I> Pipeline<I, I> for Repeat<P, I>
where
    P: Pipeline<I, I>,
    I: Clone + Send + Sync + 'static,
{
    fn apply(&self, ctx: Context<I>) -> impl Future<Output = I> + Send {
        async move {
            let (input_arc, shared) = ctx.into_parts();
            let mut val = match std::sync::Arc::try_unwrap(input_arc) {
                Ok(v) => v,
                Err(arc) => (*arc).clone(),
            };

            for _ in 0..self.times {
                let iter_ctx = Context::from_parts(std::sync::Arc::new(val), shared.clone());
                val = self.p.apply(iter_ctx).await;
            }
            val
        }
    }
}

// --- Cache ---

pub struct Cache<P, I, O> {
    p: P,
    cache: quick_cache::sync::Cache<I, O>,
}

impl<P, I, O> Pipeline<I, O> for Cache<P, I, O>
where
    P: Pipeline<I, O>,
    I: std::hash::Hash + Eq + Clone + Send + Sync + 'static,
    O: Clone + Send + Sync + 'static,
{
    fn apply(&self, ctx: Context<I>) -> impl Future<Output = O> + Send {
        async move {
            let input_arc = ctx.input();
            if let Some(val) = self.cache.get(&*input_arc) {
                return val;
            }

            let res = self.p.apply(ctx).await;
            self.cache.insert((*input_arc).clone(), res.clone());
            res
        }
    }
}

impl<P, I, O> Handler<I, O, ()> for P
where
    P: Pipeline<I, O>,
    I: Clone + Send + Sync + 'static,
    O: Send + 'static,
{
    fn call(&self, ctx: Context<I>) -> impl Future<Output = O> + Send {
        self.apply(ctx)
    }
}