grafos_dsp/

pipeline.rs

//! DSP pipeline builder and execution engine.

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

use grafos_std::error::FabricError;

use crate::placement::StagePlacement;
use crate::stages::DspStage;
use crate::types::{Block, DspConfig, DspStats};

/// Source of audio blocks for a DSP pipeline.
pub trait DspSource {
    fn next_block(&mut self) -> Result<Option<Block>, FabricError>;
}

/// Sink for audio blocks at the end of a DSP pipeline.
pub trait DspSink {
    fn accept_block(&mut self, block: Block) -> Result<(), FabricError>;
}

/// A collecting sink that stores all output blocks.
pub struct CollectSink {
    blocks: Vec<Block>,
}

impl CollectSink {
    pub fn new() -> Self {
        Self { blocks: Vec::new() }
    }

    pub fn into_blocks(self) -> Vec<Block> {
        self.blocks
    }
}

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

impl DspSink for CollectSink {
    fn accept_block(&mut self, block: Block) -> Result<(), FabricError> {
        self.blocks.push(block);
        Ok(())
    }
}

/// A source that yields blocks from a Vec.
pub struct VecSource {
    blocks: Vec<Block>,
    index: usize,
}

impl VecSource {
    pub fn new(blocks: Vec<Block>) -> Self {
        Self { blocks, index: 0 }
    }
}

impl DspSource for VecSource {
    fn next_block(&mut self) -> Result<Option<Block>, FabricError> {
        if self.index < self.blocks.len() {
            let block = self.blocks[self.index].clone();
            self.index += 1;
            Ok(Some(block))
        } else {
            Ok(None)
        }
    }
}

struct StageEntry {
    stage: Box<dyn DspStage>,
    _placement: StagePlacement,
}

/// Builder for constructing a DSP processing pipeline.
pub struct DspPipeline {
    config: DspConfig,
    source: Option<Box<dyn DspSource>>,
    stages: Vec<StageEntry>,
    sink: Option<Box<dyn DspSink>>,
}

impl DspPipeline {
    /// Create a new DSP pipeline with the given configuration.
    pub fn new(config: DspConfig) -> Self {
        Self {
            config,
            source: None,
            stages: Vec::new(),
            sink: None,
        }
    }

    /// Set the audio input source.
    pub fn source(mut self, source: impl DspSource + 'static) -> Self {
        self.source = Some(Box::new(source));
        self
    }

    /// Add a processing stage with placement hint.
    pub fn stage(mut self, stage: impl DspStage + 'static, placement: StagePlacement) -> Self {
        self.stages.push(StageEntry {
            stage: Box::new(stage),
            _placement: placement,
        });
        self
    }

    /// Set the audio output sink.
    pub fn sink(mut self, sink: impl DspSink + 'static) -> Self {
        self.sink = Some(Box::new(sink));
        self
    }

    /// Build the pipeline into a runnable handle.
    pub fn build(self) -> Result<DspPipelineHandle, FabricError> {
        let source = self.source.ok_or(FabricError::CapacityExceeded)?;
        let sink = self.sink.ok_or(FabricError::CapacityExceeded)?;
        Ok(DspPipelineHandle {
            config: self.config,
            source,
            stages: self.stages,
            sink,
        })
    }
}

/// A built DSP pipeline ready for execution.
pub struct DspPipelineHandle {
    config: DspConfig,
    source: Box<dyn DspSource>,
    stages: Vec<StageEntry>,
    sink: Box<dyn DspSink>,
}

impl DspPipelineHandle {
    /// Run the pipeline to completion, processing all blocks from source to sink.
    pub fn run(&mut self) -> Result<DspStats, FabricError> {
        let mut stats = DspStats::default();
        let mut min_latency_us: u64 = u64::MAX;
        let mut total_latency_us: u64 = 0;

        // Block duration in microseconds: (block_size / sample_rate) * 1_000_000
        let block_duration_us = if self.config.sample_rate > 0 {
            (self.config.block_size as u64 * 1_000_000) / self.config.sample_rate as u64
        } else {
            0
        };

        while let Some(block) = self.source.next_block()? {
            let start = Self::timestamp_us();

            let mut current = block;
            for entry in self.stages.iter_mut() {
                current = entry.stage.process(&current)?;
            }

            self.sink.accept_block(current)?;

            let end = Self::timestamp_us();
            let elapsed = end.saturating_sub(start);

            stats.blocks_processed += 1;
            total_latency_us += elapsed;

            if elapsed > stats.max_latency_us {
                stats.max_latency_us = elapsed;
            }
            if elapsed < min_latency_us {
                min_latency_us = elapsed;
            }

            // Overrun detection: processing took longer than the block duration
            if block_duration_us > 0 && elapsed > block_duration_us {
                stats.overruns += 1;
            }
        }

        if stats.blocks_processed > 0 {
            stats.avg_latency_us = total_latency_us / stats.blocks_processed;
            stats.jitter_us = stats.max_latency_us.saturating_sub(min_latency_us);
        }

        Ok(stats)
    }

    #[cfg(feature = "std")]
    fn timestamp_us() -> u64 {
        // Use std::time for host-side timing
        use std::time::{SystemTime, UNIX_EPOCH};
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_micros() as u64)
            .unwrap_or(0)
    }

    #[cfg(not(feature = "std"))]
    fn timestamp_us() -> u64 {
        // On bare metal, return 0 — timing comes from hardware counters
        0
    }
}

#[cfg(all(test, feature = "std"))]
mod tests {
    use super::*;
    use crate::stages::DspStage;
    use crate::types::{Block, Sample};
    use alloc::vec;

    /// A stage that sleeps for a configurable duration per block,
    /// used to intentionally trigger overrun detection.
    struct SlowStage {
        delay: std::time::Duration,
    }

    impl DspStage for SlowStage {
        fn process(&mut self, block: &Block) -> Result<Block, FabricError> {
            std::thread::sleep(self.delay);
            Ok(block.clone())
        }
    }

    fn make_block(block_size: usize, sample_rate: u32) -> Block {
        Block {
            data: vec![0.0 as Sample; block_size],
            sample_rate,
            channels: 1,
        }
    }

    #[test]
    fn overrun_detected_for_slow_stage() {
        // 48kHz, 256 samples per block => block duration = 256/48000 ~= 5333 us
        let config = DspConfig {
            block_size: 256,
            sample_rate: 48_000,
            channels: 1,
        };

        let blocks = vec![
            make_block(256, 48_000),
            make_block(256, 48_000),
            make_block(256, 48_000),
        ];

        // Sleep 20ms per block — well over the ~5.3ms block duration
        let slow = SlowStage {
            delay: std::time::Duration::from_millis(20),
        };

        let mut handle = DspPipeline::new(config)
            .source(VecSource::new(blocks))
            .stage(slow, StagePlacement::default())
            .sink(CollectSink::new())
            .build()
            .unwrap();

        let stats = handle.run().unwrap();

        assert_eq!(stats.blocks_processed, 3);
        assert_eq!(stats.overruns, 3, "all 3 blocks should be overruns");
    }

    #[test]
    fn no_overrun_for_fast_stage() {
        // 48kHz, 256 samples => ~5333 us block duration
        // GainStage is trivially fast, no overruns expected
        let config = DspConfig {
            block_size: 256,
            sample_rate: 48_000,
            channels: 1,
        };

        let blocks = vec![make_block(256, 48_000), make_block(256, 48_000)];

        let gain = crate::stages::GainStage::new(1.0);

        let mut handle = DspPipeline::new(config)
            .source(VecSource::new(blocks))
            .stage(gain, StagePlacement::default())
            .sink(CollectSink::new())
            .build()
            .unwrap();

        let stats = handle.run().unwrap();

        assert_eq!(stats.blocks_processed, 2);
        assert_eq!(stats.overruns, 0, "fast stage should not trigger overruns");
    }
}