grafos_stream/

pipeline.rs

//! Pipeline builder and execution engine.
//!
//! Pipelines are constructed by starting with a [`Source`] and chaining
//! transforms via `.map()` and `.filter()`. The pipeline is terminated
//! with `.sink()`, `.fold()`, or `.collect()` and executed with `.run()`.
//!
//! # Inter-stage buffering
//!
//! The [`BufferedPipeline`] variant materializes a [`FabricQueue`] at each
//! stage boundary, proving the integration with `grafos-collections`.
//!
//! # Placement
//!
//! `.on_node()` records a [`NodeConstraint`] for the next stage. In the
//! current implementation all stages run locally; the constraint is stored
//! for future remote dispatch.

extern crate alloc;
use alloc::boxed::Box;
use alloc::vec::Vec;

use grafos_std::error::FabricError;
use serde::{de::DeserializeOwned, Serialize};

use grafos_collections::queue::FabricQueue;

use crate::placement::NodeConstraint;
use crate::stage::{Sink, Source};

/// Type alias for a boxed map function used in pipeline stages.
type MapFn = Box<dyn FnMut(&[u8]) -> Result<Option<Vec<u8>>, FabricError>>;
/// Type alias for a boxed filter predicate used in pipeline stages.
type FilterFn = Box<dyn FnMut(&[u8]) -> Result<bool, FabricError>>;

/// Entry point for building a pipeline.
pub struct Pipeline;

impl Pipeline {
    /// Start a pipeline from a source.
    pub fn from_source<T>(source: impl Source<T> + 'static) -> PipelineBuilder<T>
    where
        T: Serialize + DeserializeOwned + 'static,
    {
        // Type-erase the source into a bytes-producing closure so that
        // subsequent `.map()` calls can change the item type freely.
        let mut source = source;
        let erased_source: Box<dyn FnMut() -> Result<Option<Vec<u8>>, FabricError>> =
            Box::new(move || match source.next()? {
                Some(item) => {
                    let bytes =
                        postcard::to_allocvec(&item).map_err(|_| FabricError::IoError(-1))?;
                    Ok(Some(bytes))
                }
                None => Ok(None),
            });

        PipelineBuilder {
            source: erased_source,
            stages: Vec::new(),
            buffer_size: 1024,
            next_constraint: NodeConstraint::Any,
            _marker: core::marker::PhantomData,
        }
    }
}

/// Describes one transform step (type-erased over byte buffers).
enum StageKind {
    Map(MapFn),
    Filter(FilterFn),
}

/// One stage entry with its metadata.
pub struct StageEntry {
    kind: StageKind,
    #[allow(dead_code)] // scaffolded for placement-aware scheduling
    constraint: NodeConstraint,
    #[allow(dead_code)] // scaffolded for backpressure tuning
    buffer_size: usize,
}

/// Builder for constructing a stream processing pipeline.
///
/// Start with [`Pipeline::from_source`], chain transforms with `.map()`
/// and `.filter()`, then terminate with `.sink()`, `.fold()`, or
/// `.collect()`.
pub struct PipelineBuilder<T: 'static> {
    /// Type-erased source that yields serialized items as byte vectors.
    source: Box<dyn FnMut() -> Result<Option<Vec<u8>>, FabricError>>,
    stages: Vec<StageEntry>,
    buffer_size: usize,
    next_constraint: NodeConstraint,
    _marker: core::marker::PhantomData<T>,
}

impl<T: Serialize + DeserializeOwned + 'static> PipelineBuilder<T> {
    /// Apply a map transformation to each item.
    ///
    /// The closure receives each item and returns a transformed item.
    /// The output type must be serializable for inter-stage buffering.
    pub fn map<U, F>(mut self, mut f: F) -> PipelineBuilder<U>
    where
        U: Serialize + DeserializeOwned + 'static,
        F: FnMut(T) -> U + 'static,
    {
        let constraint = core::mem::replace(&mut self.next_constraint, NodeConstraint::Any);
        let buf_size = self.buffer_size;

        let erased: MapFn = Box::new(move |bytes: &[u8]| {
            let item: T = postcard::from_bytes(bytes).map_err(|_| FabricError::IoError(-1))?;
            let out = f(item);
            let out_bytes = postcard::to_allocvec(&out).map_err(|_| FabricError::IoError(-1))?;
            Ok(Some(out_bytes))
        });

        let mut stages = self.stages;
        stages.push(StageEntry {
            kind: StageKind::Map(erased),
            constraint,
            buffer_size: buf_size,
        });

        PipelineBuilder {
            source: self.source,
            stages,
            buffer_size: 1024,
            next_constraint: NodeConstraint::Any,
            _marker: core::marker::PhantomData,
        }
    }

    /// Filter items by a predicate. Items for which the predicate returns
    /// `false` are dropped.
    pub fn filter<F>(mut self, mut pred: F) -> PipelineBuilder<T>
    where
        F: FnMut(&T) -> bool + 'static,
    {
        let constraint = core::mem::replace(&mut self.next_constraint, NodeConstraint::Any);
        let buf_size = self.buffer_size;

        let erased: FilterFn = Box::new(move |bytes: &[u8]| {
            let item: T = postcard::from_bytes(bytes).map_err(|_| FabricError::IoError(-1))?;
            Ok(pred(&item))
        });

        let mut stages = self.stages;
        stages.push(StageEntry {
            kind: StageKind::Filter(erased),
            constraint,
            buffer_size: buf_size,
        });

        PipelineBuilder {
            source: self.source,
            stages,
            buffer_size: 1024,
            next_constraint: NodeConstraint::Any,
            _marker: core::marker::PhantomData,
        }
    }

    /// Set a placement constraint for the next stage.
    pub fn on_node(mut self, constraint: NodeConstraint) -> Self {
        self.next_constraint = constraint;
        self
    }

    /// Set the inter-stage buffer capacity for the next stage boundary.
    pub fn buffer_size(mut self, n: usize) -> Self {
        self.buffer_size = n;
        self
    }

    /// Terminate the pipeline with a sink, then call `.run()` to execute.
    pub fn sink<K: Sink<T> + 'static>(self, sink: K) -> SinkPipeline<T, K> {
        SinkPipeline {
            source: self.source,
            stages: self.stages,
            sink,
            _marker: core::marker::PhantomData,
        }
    }

    /// Terminate the pipeline with a fold, then call `.run()` to execute.
    pub fn fold<Acc, F>(self, init: Acc, f: F) -> FoldPipeline<T, Acc, F>
    where
        F: FnMut(Acc, T) -> Acc,
    {
        FoldPipeline {
            source: self.source,
            stages: self.stages,
            init: Some(init),
            f,
            _marker: core::marker::PhantomData,
        }
    }

    /// Convenience: terminate with a collect and call `.run()` to get items.
    pub fn collect(self) -> CollectPipeline<T> {
        CollectPipeline {
            source: self.source,
            stages: self.stages,
            _marker: core::marker::PhantomData,
        }
    }

    /// Build the pipeline into a handle that can be run.
    pub fn build(self) -> Result<PipelineHandle<T>, FabricError> {
        Ok(PipelineHandle {
            source: self.source,
            stages: self.stages,
            _marker: core::marker::PhantomData,
        })
    }
}

/// Pull all items from the erased source, push through type-erased
/// stages, and return the serialized output items.
fn run_stages(
    source: &mut dyn FnMut() -> Result<Option<Vec<u8>>, FabricError>,
    stages: &mut [StageEntry],
) -> Result<Vec<Vec<u8>>, FabricError> {
    let mut items: Vec<Vec<u8>> = Vec::new();
    while let Some(bytes) = source()? {
        items.push(bytes);
    }

    for stage in stages.iter_mut() {
        let mut output = Vec::new();
        match &mut stage.kind {
            StageKind::Map(f) => {
                for item_bytes in &items {
                    if let Some(out_bytes) = f(item_bytes)? {
                        output.push(out_bytes);
                    }
                }
            }
            StageKind::Filter(pred) => {
                for item_bytes in &items {
                    if pred(item_bytes)? {
                        output.push(item_bytes.clone());
                    }
                }
            }
        }
        items = output;
    }

    Ok(items)
}

fn deserialize_items<T: DeserializeOwned>(items: Vec<Vec<u8>>) -> Result<Vec<T>, FabricError> {
    let mut result = Vec::with_capacity(items.len());
    for item_bytes in items {
        let item: T = postcard::from_bytes(&item_bytes).map_err(|_| FabricError::IoError(-1))?;
        result.push(item);
    }
    Ok(result)
}

/// A built pipeline ready for execution.
pub struct PipelineHandle<T: 'static> {
    source: Box<dyn FnMut() -> Result<Option<Vec<u8>>, FabricError>>,
    stages: Vec<StageEntry>,
    _marker: core::marker::PhantomData<T>,
}

impl<T: Serialize + DeserializeOwned + 'static> PipelineHandle<T> {
    /// Run the pipeline to completion, discarding output.
    pub fn run(&mut self) -> Result<(), FabricError> {
        let _ = run_stages(self.source.as_mut(), &mut self.stages)?;
        Ok(())
    }
}

/// Pipeline terminated with a [`Sink`].
pub struct SinkPipeline<T: 'static, K: Sink<T>> {
    source: Box<dyn FnMut() -> Result<Option<Vec<u8>>, FabricError>>,
    stages: Vec<StageEntry>,
    sink: K,
    _marker: core::marker::PhantomData<T>,
}

impl<T, K> SinkPipeline<T, K>
where
    T: Serialize + DeserializeOwned + 'static,
    K: Sink<T>,
{
    /// Run the pipeline to completion, delivering items to the sink.
    pub fn run(&mut self) -> Result<(), FabricError> {
        let items = run_stages(self.source.as_mut(), &mut self.stages)?;
        for item_bytes in items {
            let item: T =
                postcard::from_bytes(&item_bytes).map_err(|_| FabricError::IoError(-1))?;
            self.sink.accept(item)?;
        }
        Ok(())
    }
}

/// Pipeline terminated with a fold operation.
pub struct FoldPipeline<T: 'static, Acc, F: FnMut(Acc, T) -> Acc> {
    source: Box<dyn FnMut() -> Result<Option<Vec<u8>>, FabricError>>,
    stages: Vec<StageEntry>,
    init: Option<Acc>,
    f: F,
    _marker: core::marker::PhantomData<T>,
}

impl<T, Acc, F> FoldPipeline<T, Acc, F>
where
    T: Serialize + DeserializeOwned + 'static,
    F: FnMut(Acc, T) -> Acc,
{
    /// Run the pipeline to completion and return the fold result.
    pub fn run(&mut self) -> Result<Acc, FabricError> {
        let items = run_stages(self.source.as_mut(), &mut self.stages)?;
        let mut acc = self.init.take().ok_or(FabricError::IoError(-10))?;
        for item_bytes in items {
            let item: T =
                postcard::from_bytes(&item_bytes).map_err(|_| FabricError::IoError(-1))?;
            acc = (self.f)(acc, item);
        }
        Ok(acc)
    }
}

/// Pipeline terminated with collect — returns items as `Vec<T>`.
pub struct CollectPipeline<T: 'static> {
    source: Box<dyn FnMut() -> Result<Option<Vec<u8>>, FabricError>>,
    stages: Vec<StageEntry>,
    _marker: core::marker::PhantomData<T>,
}

impl<T: Serialize + DeserializeOwned + 'static> CollectPipeline<T> {
    /// Run the pipeline and return collected items.
    pub fn run(&mut self) -> Result<Vec<T>, FabricError> {
        let items = run_stages(self.source.as_mut(), &mut self.stages)?;
        deserialize_items(items)
    }
}

/// Pipeline variant using FabricQueue inter-stage buffering.
///
/// Materializes a [`FabricQueue`] at each stage boundary, proving the
/// integration with `grafos-collections`. Items are serialized into the
/// queue by the producer and deserialized by the consumer.
pub struct BufferedPipeline<T: 'static> {
    source: Box<dyn FnMut() -> Result<Option<Vec<u8>>, FabricError>>,
    stages: Vec<StageEntry>,
    queue_capacity: usize,
    queue_stride: usize,
    _marker: core::marker::PhantomData<T>,
}

impl<T: Serialize + DeserializeOwned + 'static> BufferedPipeline<T> {
    /// Create a buffered pipeline that uses FabricQueues between stages.
    pub fn new(
        source: impl Source<T> + 'static,
        stages: Vec<StageEntry>,
        capacity: usize,
        stride: usize,
    ) -> Self {
        let mut source = source;
        let erased: Box<dyn FnMut() -> Result<Option<Vec<u8>>, FabricError>> =
            Box::new(move || match source.next()? {
                Some(item) => {
                    let bytes =
                        postcard::to_allocvec(&item).map_err(|_| FabricError::IoError(-1))?;
                    Ok(Some(bytes))
                }
                None => Ok(None),
            });

        BufferedPipeline {
            source: erased,
            stages,
            queue_capacity: capacity,
            queue_stride: stride,
            _marker: core::marker::PhantomData,
        }
    }

    /// Run the pipeline using FabricQueues for inter-stage buffering.
    ///
    /// Each stage boundary creates a FabricQueue, pushes items in, applies
    /// the transform, and collects the output.
    pub fn run(&mut self) -> Result<Vec<T>, FabricError> {
        // Pull all source items into the first queue.
        let mut queue: FabricQueue<Vec<u8>> =
            FabricQueue::with_capacity(self.queue_capacity, self.queue_stride)?;

        while let Some(bytes) = (self.source)()? {
            if !queue.push(&bytes)? {
                return Err(FabricError::CapacityExceeded);
            }
        }

        // Process through each stage: drain queue -> transform -> new queue.
        for stage in self.stages.iter_mut() {
            let mut next_queue: FabricQueue<Vec<u8>> =
                FabricQueue::with_capacity(self.queue_capacity, self.queue_stride)?;

            while let Some(item_bytes) = queue.pop()? {
                match &mut stage.kind {
                    StageKind::Map(f) => {
                        if let Some(out_bytes) = f(&item_bytes)? {
                            if !next_queue.push(&out_bytes)? {
                                return Err(FabricError::CapacityExceeded);
                            }
                        }
                    }
                    StageKind::Filter(pred) => {
                        if pred(&item_bytes)? && !next_queue.push(&item_bytes)? {
                            return Err(FabricError::CapacityExceeded);
                        }
                    }
                }
            }

            queue = next_queue;
        }

        // Drain final queue into output.
        let mut result = Vec::new();
        while let Some(item_bytes) = queue.pop()? {
            let item: T =
                postcard::from_bytes(&item_bytes).map_err(|_| FabricError::IoError(-1))?;
            result.push(item);
        }
        Ok(result)
    }
}

/// Helper to construct a map stage entry for [`BufferedPipeline`].
pub fn make_map_stage(f: MapFn) -> StageEntry {
    StageEntry {
        kind: StageKind::Map(f),
        constraint: NodeConstraint::Any,
        buffer_size: 1024,
    }
}

/// Helper to construct a filter stage entry for [`BufferedPipeline`].
pub fn make_filter_stage(f: FilterFn) -> StageEntry {
    StageEntry {
        kind: StageKind::Filter(f),
        constraint: NodeConstraint::Any,
        buffer_size: 1024,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sink::CountSink;
    use crate::source::VecSource;
    use grafos_std::host;

    fn setup_mock() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);
    }

    #[test]
    fn simple_map_pipeline() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 2, 3, 4, 5]);
        let mut pipeline = Pipeline::from_source(source).map(|x: u32| x * 2).collect();
        let result = pipeline.run().expect("run");
        assert_eq!(result, vec![2, 4, 6, 8, 10]);
    }

    #[test]
    fn filter_pipeline() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 5, 10, 15, 20, 25]);
        let mut pipeline = Pipeline::from_source(source)
            .filter(|x: &u32| *x > 10)
            .collect();
        let result = pipeline.run().expect("run");
        assert_eq!(result, vec![15, 20, 25]);
    }

    #[test]
    fn fold_pipeline() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 2, 3, 4, 5]);
        let mut pipeline = Pipeline::from_source(source).fold(0u64, |acc, x: u32| acc + x as u64);
        let result = pipeline.run().expect("run");
        assert_eq!(result, 15);
    }

    #[test]
    fn sink_pipeline() {
        setup_mock();
        let source = VecSource::new(vec![10u32, 20, 30]);
        let sink: CountSink<u32> = CountSink::new();
        let mut pipeline = Pipeline::from_source(source).sink(sink);
        pipeline.run().expect("run");
    }

    #[test]
    fn map_then_filter() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 2, 3, 4, 5]);
        let mut pipeline = Pipeline::from_source(source)
            .map(|x: u32| x * 2)
            .filter(|x: &u32| *x > 4)
            .collect();
        let result = pipeline.run().expect("run");
        assert_eq!(result, vec![6, 8, 10]);
    }

    #[test]
    fn multi_stage_map_filter_fold() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        let mut pipeline = Pipeline::from_source(source)
            .map(|x: u32| x * 3)
            .filter(|x: &u32| *x > 10)
            .fold(0u64, |acc, x: u32| acc + x as u64);
        // Items > 10 after *3: 12, 15, 18, 21, 24, 27, 30
        // Sum: 12+15+18+21+24+27+30 = 147
        let result = pipeline.run().expect("run");
        assert_eq!(result, 147);
    }

    #[test]
    fn empty_source() {
        setup_mock();
        let source = VecSource::<u32>::new(vec![]);
        let mut pipeline = Pipeline::from_source(source).collect();
        let result = pipeline.run().expect("run");
        assert!(result.is_empty());
    }

    #[test]
    fn empty_source_fold() {
        setup_mock();
        let source = VecSource::<u32>::new(vec![]);
        let mut pipeline = Pipeline::from_source(source).fold(42u64, |acc, _x: u32| acc + 1);
        let result = pipeline.run().expect("run");
        assert_eq!(result, 42);
    }

    #[test]
    fn filter_removes_everything() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 2, 3]);
        let mut pipeline = Pipeline::from_source(source)
            .filter(|_x: &u32| false)
            .collect();
        let result = pipeline.run().expect("run");
        assert!(result.is_empty());
    }

    #[test]
    fn on_node_is_accepted() {
        setup_mock();
        let source = VecSource::new(vec![1u32]);
        let mut pipeline = Pipeline::from_source(source)
            .on_node(NodeConstraint::HasMemory(1024))
            .map(|x: u32| x + 1)
            .collect();
        let result = pipeline.run().expect("run");
        assert_eq!(result, vec![2]);
    }

    #[test]
    fn buffer_size_is_accepted() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 2, 3]);
        let mut pipeline = Pipeline::from_source(source)
            .buffer_size(64)
            .map(|x: u32| x + 10)
            .collect();
        let result = pipeline.run().expect("run");
        assert_eq!(result, vec![11, 12, 13]);
    }

    #[test]
    fn buffered_pipeline_with_fabric_queue() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 2, 3, 4, 5]);

        let map_stage = make_map_stage(Box::new(|bytes: &[u8]| {
            let item: u32 = postcard::from_bytes(bytes).map_err(|_| FabricError::IoError(-1))?;
            let out = item * 2;
            let out_bytes = postcard::to_allocvec(&out).map_err(|_| FabricError::IoError(-1))?;
            Ok(Some(out_bytes))
        }));

        let mut buffered = BufferedPipeline::new(source, vec![map_stage], 32, 64);
        let result = buffered.run().expect("run");
        assert_eq!(result, vec![2u32, 4, 6, 8, 10]);
    }

    #[test]
    fn backpressure_small_buffer() {
        setup_mock();
        // capacity=4 means 3 usable slots — just enough for 3 items.
        let source = VecSource::new(vec![1u32, 2, 3]);
        let mut buffered = BufferedPipeline::new(source, vec![], 4, 64);
        let result = buffered.run().expect("run");
        assert_eq!(result, vec![1u32, 2, 3]);
    }

    #[test]
    fn pipeline_handle_run() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 2, 3]);
        let mut handle = Pipeline::from_source(source)
            .map(|x: u32| x + 100)
            .build()
            .expect("build");
        handle.run().expect("run");
    }

    #[test]
    fn count_sink_standalone() {
        use crate::stage::Sink as _;
        let mut sink: CountSink<u32> = CountSink::new();
        sink.accept(10).unwrap();
        sink.accept(20).unwrap();
        sink.accept(30).unwrap();
        assert_eq!(sink.count(), 3);
    }

    #[test]
    fn buffered_pipeline_filter() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 2, 3, 4, 5]);

        let filter_stage = make_filter_stage(Box::new(|bytes: &[u8]| {
            let item: u32 = postcard::from_bytes(bytes).map_err(|_| FabricError::IoError(-1))?;
            Ok(item > 3)
        }));

        let mut buffered = BufferedPipeline::new(source, vec![filter_stage], 32, 64);
        let result = buffered.run().expect("run");
        assert_eq!(result, vec![4u32, 5]);
    }

    #[test]
    fn buffered_pipeline_map_then_filter() {
        setup_mock();
        let source = VecSource::new(vec![1u32, 2, 3, 4, 5]);

        let map_stage = make_map_stage(Box::new(|bytes: &[u8]| {
            let item: u32 = postcard::from_bytes(bytes).map_err(|_| FabricError::IoError(-1))?;
            let out = item * 10;
            let out_bytes = postcard::to_allocvec(&out).map_err(|_| FabricError::IoError(-1))?;
            Ok(Some(out_bytes))
        }));

        let filter_stage = make_filter_stage(Box::new(|bytes: &[u8]| {
            let item: u32 = postcard::from_bytes(bytes).map_err(|_| FabricError::IoError(-1))?;
            Ok(item > 25)
        }));

        let mut buffered = BufferedPipeline::new(source, vec![map_stage, filter_stage], 32, 64);
        let result = buffered.run().expect("run");
        assert_eq!(result, vec![30u32, 40, 50]);
    }
}