tenshift-core 0.1.2

use super::ops::{cast_numeric_slice, expected_byte_len, numeric_to_le_bytes};
use crate::error::{Error, Result};
use std::sync::{Arc, OnceLock};

#[cfg(feature = "uring")]
use wireshift::{Buffer, Completed};

/// A typed, shaped buffer of data.
///
/// Tensors hold raw bytes with a dtype and shape. Data is reference-counted
/// so cloning a tensor is cheap (no copy). The actual bytes are only copied
/// when mutated (copy-on-write).
#[derive(Debug, Clone)]
pub struct Tensor {
    /// The raw data.
    pub(crate) data: Arc<TensorData>,
    /// Element type.
    pub(crate) dtype: DType,
    /// Shape dimensions. Arc-shared to avoid per-clone Vec allocation.
    pub(crate) shape: Arc<[usize]>,
}

#[derive(Debug, Default)]
pub(crate) struct TensorData {
    pub(crate) bytes: TensorBytes,
    pub(crate) f32_cache: OnceLock<Vec<f32>>,
    pub(crate) f64_cache: OnceLock<Vec<f64>>,
    pub(crate) i32_cache: OnceLock<Vec<i32>>,
    pub(crate) i64_cache: OnceLock<Vec<i64>>,
}

#[derive(Debug)]
pub(crate) enum TensorBytes {
    Shared(Arc<[u8]>),
    #[cfg(feature = "uring")]
    Pooled(Buffer<Completed>),
}

impl Default for TensorBytes {
    fn default() -> Self {
        Self::Shared(Arc::from([]))
    }
}

impl TensorBytes {
    fn as_slice(&self) -> &[u8] {
        match self {
            Self::Shared(bytes) => bytes,
            #[cfg(feature = "uring")]
            Self::Pooled(buffer) => buffer.filled(),
        }
    }

    fn len(&self) -> usize {
        self.as_slice().len()
    }
}

/// Supported element types.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DType {
    /// 32-bit float.
    F32,
    /// 64-bit float.
    F64,
    /// 8-bit unsigned integer.
    U8,
    /// 32-bit signed integer.
    I32,
    /// 64-bit signed integer.
    I64,
    /// Raw bytes (untyped).
    Bytes,
}

impl DType {
    /// Size of one element in bytes.
    pub fn element_size(self) -> usize {
        match self {
            Self::F32 | Self::I32 => 4,
            Self::F64 | Self::I64 => 8,
            Self::U8 | Self::Bytes => 1,
        }
    }

    /// Human-readable name for this dtype.
    pub fn name(self) -> &'static str {
        match self {
            Self::F32 => "f32",
            Self::F64 => "f64",
            Self::U8 => "u8",
            Self::I32 => "i32",
            Self::I64 => "i64",
            Self::Bytes => "bytes",
        }
    }
}

impl std::fmt::Display for DType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.name())
    }
}

impl Tensor {
    /// Create a float32 tensor from a slice of `f32` values.
    pub fn f32(data: &[f32], shape: Vec<usize>) -> Self {
        Self::from_bytes_unchecked(numeric_to_le_bytes(data), DType::F32, shape)
    }

    /// Create a float64 tensor from a slice of `f64` values.
    pub fn f64(data: &[f64], shape: Vec<usize>) -> Self {
        Self::from_bytes_unchecked(numeric_to_le_bytes(data), DType::F64, shape)
    }

    /// Create a uint8 tensor (images, raw bytes).
    pub fn u8(data: Vec<u8>, shape: Vec<usize>) -> Self {
        Self::from_bytes_unchecked(data, DType::U8, shape)
    }

    /// Create an int32 tensor.
    pub fn i32(data: &[i32], shape: Vec<usize>) -> Self {
        Self::from_bytes_unchecked(numeric_to_le_bytes(data), DType::I32, shape)
    }

    /// Create an int64 tensor (labels, indices).
    pub fn i64(data: &[i64], shape: Vec<usize>) -> Self {
        Self::from_bytes_unchecked(numeric_to_le_bytes(data), DType::I64, shape)
    }

    /// Create a raw bytes tensor (untyped).
    pub fn bytes(data: Vec<u8>) -> Self {
        let len = data.len();
        Self::from_bytes_unchecked(data, DType::Bytes, vec![len])
    }

    /// Create a raw bytes tensor from a completed `wireshift` buffer.
    #[cfg(feature = "uring")]
    pub fn bytes_from_completed(buffer: Buffer<Completed>) -> Self {
        let len = buffer.filled_len();
        Self::from_storage_unchecked(TensorBytes::Pooled(buffer), DType::Bytes, vec![len])
    }

    /// Create a tensor from raw bytes, dtype, and shape.
    ///
    /// # Errors
    /// Returns an error if the provided byte length does not match the product of shape elements
    /// multiplied by the element size for the given `dtype`.
    pub fn from_bytes(data: Vec<u8>, dtype: DType, shape: Vec<usize>) -> Result<Self> {
        let expected = expected_byte_len(dtype, &shape)?;
        if data.len() != expected {
            return Err(Error::InvalidConfig {
                reason: format!(
                    "tensor byte length {} does not match dtype {:?} and shape {:?} (expected {expected})",
                    data.len(),
                    dtype,
                    shape
                ),
            });
        }

        Ok(Self::from_bytes_unchecked(data, dtype, shape))
    }

    /// The element type of this tensor.
    pub fn dtype(&self) -> DType {
        self.dtype
    }

    /// The shape of this tensor.
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    /// Total number of elements.
    ///
    /// Returns `None` if the element count overflows `usize`.
    pub fn num_elements(&self) -> Option<usize> {
        self.shape
            .iter()
            .try_fold(1_usize, |acc, &dim| acc.checked_mul(dim))
    }

    /// Raw byte data (for zero-copy transfer to numpy/torch).
    pub fn as_bytes(&self) -> &[u8] {
        self.data.bytes.as_slice()
    }

    /// Interpret data as f32 slice.
    ///
    /// # Errors
    /// Returns an error if the tensor's dtype is not `f32`.
    pub fn try_as_f32(&self) -> Result<&[f32]> {
        if self.dtype != DType::F32 {
            return Err(Error::InvalidConfig {
                reason: format!("expected dtype f32, got {}", self.dtype),
            });
        }
        cast_numeric_slice(
            self.data.bytes.as_slice(),
            &self.data.f32_cache,
            f32::from_le_bytes,
        )
    }

    /// Interpret data as f64 slice.
    ///
    /// # Errors
    /// Returns an error if the tensor's dtype is not `f64`.
    pub fn try_as_f64(&self) -> Result<&[f64]> {
        if self.dtype != DType::F64 {
            return Err(Error::InvalidConfig {
                reason: format!("expected dtype f64, got {}", self.dtype),
            });
        }
        cast_numeric_slice(
            self.data.bytes.as_slice(),
            &self.data.f64_cache,
            f64::from_le_bytes,
        )
    }

    /// Interpret data as i32 slice.
    ///
    /// # Errors
    /// Returns an error if the tensor's dtype is not `i32`.
    pub fn try_as_i32(&self) -> Result<&[i32]> {
        if self.dtype != DType::I32 {
            return Err(Error::InvalidConfig {
                reason: format!("expected dtype i32, got {}", self.dtype),
            });
        }
        cast_numeric_slice(
            self.data.bytes.as_slice(),
            &self.data.i32_cache,
            i32::from_le_bytes,
        )
    }

    /// Interpret data as i64 slice.
    ///
    /// # Errors
    /// Returns an error if the tensor's dtype is not `i64`.
    pub fn try_as_i64(&self) -> Result<&[i64]> {
        if self.dtype != DType::I64 {
            return Err(Error::InvalidConfig {
                reason: format!("expected dtype i64, got {}", self.dtype),
            });
        }
        cast_numeric_slice(
            self.data.bytes.as_slice(),
            &self.data.i64_cache,
            i64::from_le_bytes,
        )
    }

    /// Byte length of the data buffer.
    pub fn byte_len(&self) -> usize {
        self.data.bytes.len()
    }

    pub(crate) fn from_bytes_unchecked(data: Vec<u8>, dtype: DType, shape: Vec<usize>) -> Self {
        Self::from_storage_unchecked(TensorBytes::Shared(Arc::from(data)), dtype, shape)
    }

    pub(crate) fn from_storage_unchecked(
        data: TensorBytes,
        dtype: DType,
        shape: Vec<usize>,
    ) -> Self {
        Self {
            data: Arc::new(TensorData {
                bytes: data,
                f32_cache: OnceLock::new(),
                f64_cache: OnceLock::new(),
                i32_cache: OnceLock::new(),
                i64_cache: OnceLock::new(),
            }),
            dtype,
            shape: Arc::from(shape),
        }
    }
}