photon-ring 3.0.0

// Copyright 2026 Photon Ring Contributors
// SPDX-License-Identifier: MIT OR Apache-2.0

use crate::channel::{self, Publisher, Subscribable, Subscriber};
use crate::pod::Pod;
use crate::wait::WaitStrategy;

use super::fan_out::FanOutBuilder;
use super::pipeline::Pipeline;
use super::{SharedState, DEFAULT_CAPACITY};

// ---------------------------------------------------------------------------
// PipelineBuilder
// ---------------------------------------------------------------------------

/// Builder for constructing a [`Pipeline`].
///
/// # Example
///
/// ```
/// use photon_ring::topology::Pipeline;
///
/// let (mut pub_, stages) = Pipeline::builder()
///     .capacity(128)
///     .input::<u32>();
///
/// let (mut output, pipeline) = stages
///     .then(|x| x + 1)
///     .build();
///
/// pub_.publish(41);
/// assert_eq!(output.recv(), 42);
/// pipeline.shutdown();
/// pipeline.join();
/// ```
pub struct PipelineBuilder {
    capacity: usize,
}

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

impl PipelineBuilder {
    /// Create a new builder with default capacity (1024).
    pub fn new() -> Self {
        PipelineBuilder {
            capacity: DEFAULT_CAPACITY,
        }
    }

    /// Set the ring capacity for all inter-stage channels.
    ///
    /// Any capacity of at least 2. Powers of two index with a mask; other
    /// sizes use a reciprocal multiply.
    pub fn capacity(mut self, cap: usize) -> Self {
        self.capacity = cap;
        self
    }

    /// Declare the input type and create the first channel.
    ///
    /// Returns the publisher for writing into the pipeline and a
    /// [`StageBuilder`] for attaching processing stages.
    pub fn input<T: Pod>(self) -> (Publisher<T>, StageBuilder<T>) {
        let (pub_, subs) = channel::channel::<T>(self.capacity);
        let subscriber = subs.subscribe();
        (
            pub_,
            StageBuilder {
                subscriber,
                subscribable: subs,
                capacity: self.capacity,
                state: SharedState::new(),
            },
        )
    }
}

// ---------------------------------------------------------------------------
// StageBuilder<T> — typed builder for chaining stages
// ---------------------------------------------------------------------------

/// Intermediate builder for adding processing stages to a pipeline.
///
/// The type parameter `T` is the output type of the most recently added
/// stage (or the input type, if no stages have been added yet).
pub struct StageBuilder<T: Pod> {
    pub(super) subscriber: Subscriber<T>,
    pub(super) subscribable: Subscribable<T>,
    pub(super) capacity: usize,
    pub(super) state: SharedState,
}

impl<T: Pod> StageBuilder<T> {
    /// Add a processing stage that transforms `T -> U`.
    ///
    /// Spawns a dedicated thread that reads from the current stage's
    /// output channel, applies `f`, and publishes to a new channel.
    /// Uses [`WaitStrategy::default()`] (adaptive) for the stage's
    /// wait behavior. Use [`then_with`](Self::then_with) to specify a
    /// custom wait strategy.
    ///
    /// # Example
    ///
    /// ```
    /// use photon_ring::topology::Pipeline;
    ///
    /// let (mut pub_, stages) = Pipeline::builder()
    ///     .capacity(64)
    ///     .input::<i32>();
    ///
    /// let (mut out, pipe) = stages
    ///     .then(|x| x * 2)
    ///     .then(|x| x + 1)
    ///     .build();
    ///
    /// pub_.publish(5);
    /// assert_eq!(out.recv(), 11);
    /// pipe.shutdown();
    /// pipe.join();
    /// ```
    pub fn then<U: Pod>(self, f: impl Fn(T) -> U + Send + 'static) -> StageBuilder<U> {
        self.then_with(f, WaitStrategy::default())
    }

    /// Add a processing stage with a custom wait strategy.
    ///
    /// Identical to [`then`](Self::then), but allows specifying a
    /// [`WaitStrategy`] that controls how the stage waits when no
    /// message is available.
    ///
    /// # Example
    ///
    /// ```
    /// use photon_ring::topology::Pipeline;
    /// use photon_ring::WaitStrategy;
    ///
    /// let (mut pub_, stages) = Pipeline::builder()
    ///     .capacity(64)
    ///     .input::<u64>();
    ///
    /// let (mut out, pipe) = stages
    ///     .then_with(|x| x * 2, WaitStrategy::YieldSpin)
    ///     .then_with(|x| x + 1, WaitStrategy::BackoffSpin)
    ///     .build();
    ///
    /// pub_.publish(10);
    /// assert_eq!(out.recv(), 21);
    /// pipe.shutdown();
    /// pipe.join();
    /// ```
    pub fn then_with<U: Pod>(
        mut self,
        f: impl Fn(T) -> U + Send + 'static,
        strategy: WaitStrategy,
    ) -> StageBuilder<U> {
        let (next_sub, next_subs) =
            self.state
                .add_stage(self.subscriber, self.capacity, f, strategy);
        StageBuilder {
            subscriber: next_sub,
            subscribable: next_subs,
            capacity: self.capacity,
            state: self.state,
        }
    }

    /// Fan out to two parallel processing stages.
    ///
    /// Both stages receive the same input (via two subscribers on the
    /// same ring). Each applies its own transformation and publishes
    /// to its own output channel.
    ///
    /// Returns a [`FanOutBuilder`] that can be finalized with `.build()`.
    ///
    /// # Example
    ///
    /// ```
    /// use photon_ring::topology::Pipeline;
    ///
    /// let (mut pub_, stages) = Pipeline::builder()
    ///     .capacity(64)
    ///     .input::<u64>();
    ///
    /// let ((mut out_a, mut out_b), pipe) = stages
    ///     .fan_out(|x: u64| x * 2, |x: u64| x + 100)
    ///     .build();
    ///
    /// pub_.publish(5);
    /// let a = out_a.recv();
    /// let b = out_b.recv();
    /// assert_eq!(a, 10);
    /// assert_eq!(b, 105);
    /// pipe.shutdown();
    /// pipe.join();
    /// ```
    pub fn fan_out<A, B>(
        mut self,
        fa: impl Fn(T) -> A + Send + 'static,
        fb: impl Fn(T) -> B + Send + 'static,
    ) -> FanOutBuilder<A, B>
    where
        A: Pod,
        B: Pod,
    {
        // Branch B gets a fresh subscriber from the same source ring; branch A
        // reuses the existing one.
        let input_b = self.subscribable.subscribe();
        let strategy = WaitStrategy::default();
        let (sub_a_out, subs_a) =
            self.state
                .add_stage(self.subscriber, self.capacity, fa, strategy);
        let (sub_b_out, subs_b) = self.state.add_stage(input_b, self.capacity, fb, strategy);

        FanOutBuilder {
            sub_a: sub_a_out,
            subs_a,
            sub_b: sub_b_out,
            subs_b,
            capacity: self.capacity,
            state: self.state,
        }
    }

    /// Finalize the pipeline, returning the output subscriber and the
    /// [`Pipeline`] handle.
    ///
    /// The subscriber reads from the last stage's output channel (or
    /// directly from the input channel if no stages were added).
    /// The pipeline handle is used for shutdown and health monitoring.
    pub fn build(self) -> (Subscriber<T>, Pipeline) {
        (self.subscriber, self.state.into())
    }
}