tenshift-core 0.1.3

//! Collation utilities and modes for the collector stage.
//!
//! Collation runs after worker-side transforms and turns batches of samples
//! into the final consumer-facing structure, preserving the architecture split
//! between parallel stateless work and ordered collector work.

#![allow(clippy::module_name_repetitions)]

use super::{panic_message, CollateMode};
use crate::error::{Error, Result};
use crate::sample::{Sample, Tensor};
use std::panic::{catch_unwind, AssertUnwindSafe};

pub(crate) fn collate_batch(batch: Vec<Sample>, collate_mode: &CollateMode) -> Result<Vec<Sample>> {
    match collate_mode {
        CollateMode::Disabled => Ok(batch),
        CollateMode::Default => default_collate(batch).map(|sample| vec![sample]),
        CollateMode::Custom(collate) => match catch_unwind(AssertUnwindSafe(|| collate(batch))) {
            Ok(result) => result.map(|sample| vec![sample]),
            Err(payload) => Err(Error::CollateFailed {
                reason: format!("custom collate panicked: {}", panic_message(payload)),
            }),
        },
    }
}

pub(crate) fn default_collate(batch: Vec<Sample>) -> Result<Sample> {
    let Some(first) = batch.first() else {
        return Ok(Sample::new());
    };

    let field_names: Vec<String> = first.field_names().map(ToOwned::to_owned).collect();
    let mut output = Sample::new();

    for sample in &batch[1..] {
        let sample_fields: Vec<&str> = sample.field_names().collect();
        if sample_fields.len() != field_names.len() {
            return Err(Error::CollateFailed {
                reason: format!(
                    "batch contains inconsistent field counts: expected {}, got {}",
                    field_names.len(),
                    sample_fields.len()
                ),
            });
        }

        for field_name in &field_names {
            if !sample_fields
                .iter()
                .any(|candidate| candidate == field_name)
            {
                return Err(Error::CollateFailed {
                    reason: format!(
                        "batch contains inconsistent field names; missing '{field_name}'"
                    ),
                });
            }
        }
    }

    for field_name in field_names {
        let first_tensor = first.get(&field_name).ok_or_else(|| Error::CollateFailed {
            reason: format!("field '{field_name}' is missing from the first sample"),
        })?;
        const MAX_COLLATE_BYTES: usize = 1024 * 1024 * 1024; // 1 GiB per field
        let total_bytes = first_tensor
            .byte_len()
            .checked_mul(batch.len())
            .ok_or_else(|| Error::CollateFailed {
                reason: format!(
                    "field '{field_name}' would allocate {} * {} bytes, which overflows usize",
                    first_tensor.byte_len(),
                    batch.len()
                ),
            })?;
        if total_bytes > MAX_COLLATE_BYTES {
            return Err(Error::CollateFailed {
                reason: format!(
                    "field '{field_name}' collation would allocate {total_bytes} bytes, exceeding 1 GiB limit. fix: reduce batch size or tensor size."
                ),
            });
        }
        let mut bytes = Vec::with_capacity(total_bytes);

        for sample in &batch {
            let tensor = sample
                .get(&field_name)
                .ok_or_else(|| Error::CollateFailed {
                    reason: format!("field '{field_name}' is missing from one or more samples"),
                })?;
            if tensor.dtype() != first_tensor.dtype() {
                return Err(Error::CollateFailed {
                    reason: format!("field '{field_name}' has mixed dtypes within the batch"),
                });
            }
            if tensor.shape() != first_tensor.shape() {
                return Err(Error::CollateFailed {
                    reason: format!("field '{field_name}' has mixed shapes within the batch"),
                });
            }
            bytes.extend_from_slice(tensor.as_bytes());
        }

        let mut shape = Vec::with_capacity(first_tensor.shape().len() + 1);
        shape.push(batch.len());
        shape.extend_from_slice(first_tensor.shape());

        let stacked = Tensor::from_bytes(bytes, first_tensor.dtype(), shape)?;
        output.insert(field_name, stacked);
    }

    Ok(output)
}