tenshift_core/

transform.rs

//! Transform  -  composable operations on samples.
//!
//! Transforms are the second community extension point. Built-in transforms
//! cover the common cases (map, filter, shuffle, batch). Community transforms
//! can add augmentation, normalization, tokenization, etc.
//!
//! # Stateless vs Stateful
//!
//! - **[`Transform`]**: Stateless transforms process samples independently.
//!   They run in parallel across worker threads. Examples: `map`, `filter`, `flat_map`.
//!
//! - **[`StatefulTransform`]**: Stateful transforms maintain internal buffers
//!   and require ordered processing. They run serially in the collector thread.
//!   Examples: `shuffle`, `batch`.

use crate::error::Error;
use crate::sample::Sample;

/// Result of applying a transform to a sample.
#[non_exhaustive]
pub enum TransformResult {
    /// The transformed sample.
    Sample(Sample),
    /// Yield multiple samples from one input.
    Samples(Vec<Sample>),
    /// Skip this sample (e.g., filtered out).
    Skip,
    /// An error occurred during transformation.
    Error(Error),
}

/// A stateless transform that operates on individual samples.
///
/// For stateful operations (shuffle, batch), use [`StatefulTransform`].
///
/// # Implementing a Transform
///
/// ```rust
/// use tenshift_core::transform::{Transform, TransformResult};
/// use tenshift_core::sample::Sample;
///
/// struct Normalize {
///     mean: f32,
///     std: f32,
/// }
///
/// impl Transform for Normalize {
///     fn apply(&self, sample: Sample) -> TransformResult {
///         // Normalize a tensor field in-place
///         TransformResult::Sample(sample)
///     }
///
///     fn name(&self) -> &str {
///         "normalize"
///     }
/// }
/// ```
pub trait Transform: Send + Sync {
    /// Apply this transform to a sample.
    fn apply(&self, sample: Sample) -> TransformResult;

    /// Human-readable name for logging.
    fn name(&self) -> &str;
}

/// A stateful transform that can buffer, reorder, or group samples.
///
/// Unlike [`Transform`], stateful transforms receive samples one at a time
/// via [`push`](StatefulTransform::push) and yield zero or more samples via
/// [`finish`](StatefulTransform::finish) or internal buffering.
pub trait StatefulTransform: Send {
    /// Push a sample into this transform. Returns any immediately available outputs.
    fn push(&mut self, sample: Sample) -> Vec<Sample>;

    /// Signal that no more input is coming. Returns any buffered samples.
    fn finish(&mut self) -> Vec<Sample>;

    /// Human-readable name for logging.
    fn name(&self) -> &str;
}

/// Map transform  -  apply a function to each sample.
pub struct MapTransform<F> {
    func: F,
}

impl<F> MapTransform<F>
where
    F: Fn(Sample) -> crate::error::Result<Sample> + Send + Sync,
{
    /// Create a new map transform.
    pub fn new(func: F) -> Self {
        Self { func }
    }
}

impl<F> Transform for MapTransform<F>
where
    F: Fn(Sample) -> crate::error::Result<Sample> + Send + Sync,
{
    fn apply(&self, sample: Sample) -> TransformResult {
        match (self.func)(sample) {
            Ok(s) => TransformResult::Sample(s),
            Err(e) => TransformResult::Error(e),
        }
    }

    #[allow(clippy::needless_borrows_for_generic_args)]
    fn name(&self) -> &str {
        "map"
    }
}

/// `FlatMap` transform  -  apply a function that yields zero or more samples.
pub struct FlatMapTransform<F> {
    func: F,
}

impl<F> FlatMapTransform<F>
where
    F: Fn(Sample) -> crate::error::Result<Vec<Sample>> + Send + Sync,
{
    /// Create a new `flat_map` transform.
    pub fn new(func: F) -> Self {
        Self { func }
    }
}

impl<F> Transform for FlatMapTransform<F>
where
    F: Fn(Sample) -> crate::error::Result<Vec<Sample>> + Send + Sync,
{
    fn apply(&self, sample: Sample) -> TransformResult {
        match (self.func)(sample) {
            Ok(samples) => TransformResult::Samples(samples),
            Err(e) => TransformResult::Error(e),
        }
    }

    #[allow(clippy::needless_borrows_for_generic_args)]
    fn name(&self) -> &str {
        "flat_map"
    }
}

/// Filter transform  -  keep only samples that match a predicate.
pub struct FilterTransform<F> {
    predicate: F,
}

impl<F> FilterTransform<F>
where
    F: Fn(&Sample) -> bool + Send + Sync,
{
    /// Create a new filter transform.
    pub fn new(predicate: F) -> Self {
        Self { predicate }
    }
}

impl<F> Transform for FilterTransform<F>
where
    F: Fn(&Sample) -> bool + Send + Sync,
{
    fn apply(&self, sample: Sample) -> TransformResult {
        if (self.predicate)(&sample) {
            TransformResult::Sample(sample)
        } else {
            TransformResult::Skip
        }
    }

    fn name(&self) -> &str {
        "filter"
    }
}

/// Shuffle transform  -  randomize sample order using a reservoir buffer.
///
/// Maintains a buffer of `capacity` samples. When the buffer is full,
/// a random sample is evicted and yielded. This provides streaming
/// approximate shuffling with bounded memory.
pub struct ShuffleBuffer {
    buffer: Vec<Sample>,
    capacity: usize,
    rng_state: u64,
}

impl ShuffleBuffer {
    /// Create a new shuffle buffer with the given capacity.
    pub fn new(capacity: usize, seed: Option<u64>) -> Self {
        Self {
            buffer: Vec::with_capacity(capacity),
            capacity: capacity.max(1),
            rng_state: seed.unwrap_or(crate::pipeline::DEFAULT_SHUFFLE_SEED),
        }
    }

    /// Simple xorshift64 PRNG  -  fast, no dependencies.
    fn next_rand(&mut self) -> u64 {
        let mut x = self.rng_state;
        x ^= x << 13;
        x ^= x >> 7;
        x ^= x << 17;
        self.rng_state = x;
        x
    }
}

impl StatefulTransform for ShuffleBuffer {
    fn push(&mut self, sample: Sample) -> Vec<Sample> {
        if self.buffer.len() < self.capacity {
            self.buffer.push(sample);
            Vec::new()
        } else {
            // Buffer full  -  swap a random element out
            #[allow(clippy::cast_possible_truncation)]
            let idx = (self.next_rand() as usize) % self.buffer.len();
            let evicted = std::mem::replace(&mut self.buffer[idx], sample);
            vec![evicted]
        }
    }

    fn finish(&mut self) -> Vec<Sample> {
        // Drain remaining buffer in shuffled order
        let mut remaining = std::mem::take(&mut self.buffer);
        // Fisher-Yates shuffle on the remaining buffer
        for i in (1..remaining.len()).rev() {
            #[allow(clippy::cast_possible_truncation)]
            let j = (self.next_rand() as usize) % (i + 1);
            remaining.swap(i, j);
        }
        remaining
    }

    fn name(&self) -> &str {
        "shuffle"
    }
}

/// Batch transform  -  accumulate N samples into a group.
pub struct BatchAccumulator {
    batch_size: usize,
    drop_last: bool,
    buffer: Vec<Sample>,
}

impl BatchAccumulator {
    /// Create a new batch accumulator.
    pub fn new(batch_size: usize, drop_last: bool) -> Self {
        Self {
            batch_size: batch_size.max(1),
            drop_last,
            buffer: Vec::new(),
        }
    }
}

impl StatefulTransform for BatchAccumulator {
    fn push(&mut self, sample: Sample) -> Vec<Sample> {
        self.buffer.push(sample);
        if self.buffer.len() >= self.batch_size {
            std::mem::replace(&mut self.buffer, Vec::with_capacity(self.batch_size))
        } else {
            Vec::new()
        }
    }

    fn finish(&mut self) -> Vec<Sample> {
        if self.drop_last && self.buffer.len() < self.batch_size {
            self.buffer.clear();
            Vec::new()
        } else {
            std::mem::take(&mut self.buffer)
        }
    }

    fn name(&self) -> &str {
        "batch"
    }
}