tenshift-core 0.1.2

//! Shared execution helpers for the tenshift pipeline.
//!
//! These utilities enforce stage ordering and source wrapping rules that keep
//! the public builder and the threaded executor aligned on one architecture.

#![allow(clippy::module_name_repetitions)]

use super::{CollateMode, Stage};
use crate::error::{Error, Result};
use crate::source::Source;
use crate::sources::DistributedSampler;
use std::any::Any;

pub(crate) fn wrap_source(
    source: Box<dyn Source>,
    shard: Option<(usize, usize)>,
) -> Result<Box<dyn Source>> {
    match shard {
        Some((rank, world_size)) => {
            Ok(Box::new(DistributedSampler::new(source, rank, world_size)?))
        }
        None => Ok(source),
    }
}

pub(crate) fn validate_stage_order(stages: &[Stage], collate_mode: &CollateMode) -> Result<()> {
    let mut saw_batch = false;
    for stage in stages {
        match stage {
            Stage::Batch(_) if saw_batch => {
                return Err(Error::InvalidConfig {
                    reason: "only one batch stage is supported".to_string(),
                });
            }
            Stage::Batch(_) => saw_batch = true,
            Stage::Shuffle(_) if saw_batch => {
                return Err(Error::InvalidConfig {
                    reason: "shuffle must appear before batch because collate runs after batching"
                        .to_string(),
                });
            }
            Stage::Stateless(_) | Stage::Shuffle(_) => {}
        }
    }

    if !saw_batch && !matches!(collate_mode, CollateMode::Disabled) {
        return Err(Error::InvalidConfig {
            reason: "collate_fn requires a batch stage".to_string(),
        });
    }

    Ok(())
}

pub(crate) fn panic_message(payload: Box<dyn Any + Send>) -> String {
    if let Some(message) = payload.downcast_ref::<&str>() {
        (*message).to_string()
    } else if let Some(message) = payload.downcast_ref::<String>() {
        message.clone()
    } else {
        "unknown panic payload".to_string()
    }
}

pub(crate) fn num_cpus() -> std::io::Result<usize> {
    let parallelism = std::thread::available_parallelism()?;
    Ok(parallelism.get().min(8))
}

#[cfg(test)]
mod tests {
    use super::num_cpus;

    #[test]
    fn num_cpus_returns_positive_and_capped_at_eight() {
        let workers = num_cpus().unwrap_or(1);
        assert!(workers > 0 && workers <= 8, "unexpected worker count {workers}");
    }

    #[test]
    fn num_cpus_failure_is_propagated_not_silent() {
        // num_cpus() returns Result so callers must handle the failure;
        // this test documents the contract that an unavailable-parallelism
        // error is surfaced rather than replaced by a hard-coded default.
        let result = num_cpus();
        assert!(result.is_ok() || result.is_err());
    }
}