konjure_sdk/data/

tensor.rs

//! Immutable, bounded, type-erased tensors.

use alloc::{string::ToString, sync::Arc, vec, vec::Vec};
use core::{fmt, marker::PhantomData};
use serde::{
    Deserialize, Deserializer, Serialize, Serializer,
    de::{Error as _, SeqAccess, Visitor},
    ser::{SerializeSeq, SerializeStruct},
};

use super::{BinaryOp, DType, DataError, Scalar};

/// Maximum number of elements admitted by one tensor.
pub const MAX_ELEMENTS: usize = 16 * 1024 * 1024;
/// Maximum number of axes admitted by one tensor.
pub const MAX_RANK: usize = 32;

/// One item in a tensor indexing expression.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AxisIndex {
    /// Select one axis element, removing that axis from the result.
    Index(isize),
    /// Select an axis interval using Python's half-open slicing rules.
    Slice {
        /// Inclusive starting position, or the direction-dependent default.
        start: Option<isize>,
        /// Exclusive stopping position, or the direction-dependent default.
        stop: Option<isize>,
        /// Non-zero distance between successive selected elements.
        step: isize,
    },
    /// Insert a length-one axis into the result.
    NewAxis,
    /// Expand to enough full slices to consume the remaining input axes.
    Ellipsis,
}

/// An immutable, strided tensor whose packed backing storage is shared by views.
///
/// Construction accepts at most [`MAX_RANK`] axes and [`MAX_ELEMENTS`] elements.
/// Slices and transposes retain the same storage; reshaping a noncontiguous view
/// materializes row-major bytes. Binary operations require matching types and use
/// trailing-axis broadcasting. Serialization is `{dtype, shape, data}`, where
/// every `data` entry is an exact decimal scalar string.
#[derive(Clone, Debug)]
pub struct Tensor {
    dtype: DType,
    shape: Arc<[usize]>,
    strides: Arc<[isize]>,
    offset: usize,
    data: Arc<[u8]>,
}

impl PartialEq for Tensor {
    fn eq(&self, other: &Self) -> bool {
        self.dtype == other.dtype
            && self.shape == other.shape
            && self.len() == other.len()
            && (0..self.len())
                .all(|index| self.scalar_at_linear(index) == other.scalar_at_linear(index))
    }
}

impl Eq for Tensor {}

impl Tensor {
    /// Creates a contiguous tensor from scalars of `dtype`.
    pub fn from_scalars(
        dtype: DType,
        shape: Vec<usize>,
        values: Vec<Scalar>,
    ) -> Result<Self, DataError> {
        let count = checked_count(&shape)?;
        if count != values.len() {
            return Err(DataError::ShapeMismatch);
        }
        if values.iter().any(|value| value.dtype() != dtype) {
            return Err(DataError::TypeMismatch);
        }
        let bytes = values_to_bytes(dtype, &values)?;
        Self::from_le_bytes(dtype, shape, bytes)
    }

    /// Creates a contiguous tensor from packed little-endian scalar bytes.
    pub fn from_le_bytes(
        dtype: DType,
        shape: Vec<usize>,
        bytes: Vec<u8>,
    ) -> Result<Self, DataError> {
        let count = checked_count(&shape)?;
        let width = dtype.bytes();
        if width == 0
            || width > 16
            || count.checked_mul(width).ok_or(DataError::AllocationLimit)? != bytes.len()
        {
            return Err(DataError::InvalidMetadata);
        }
        // Decode once: this validates the representation, including non-finite floats.
        for index in 0..count {
            let _ = scalar_at_bytes(dtype, &bytes, index)?;
        }
        let strides = contiguous_strides(&shape)?;
        Ok(Self {
            dtype,
            shape: Arc::from(shape),
            strides: Arc::from(strides),
            offset: 0,
            data: Arc::from(bytes),
        })
    }

    /// Creates a zero-filled contiguous tensor.
    pub fn zeros(dtype: DType, shape: Vec<usize>) -> Result<Self, DataError> {
        let count = checked_count(&shape)?;
        let bytes = count
            .checked_mul(dtype.bytes())
            .ok_or(DataError::AllocationLimit)?;
        Self::from_le_bytes(dtype, shape, vec![0; bytes])
    }

    /// Returns the scalar type.
    pub fn dtype(&self) -> DType {
        self.dtype
    }
    /// Returns the dimensions in logical axis order.
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
    /// Returns element strides for each logical axis.
    pub fn strides(&self) -> &[isize] {
        &self.strides
    }
    /// Returns the logical element count.
    pub fn len(&self) -> usize {
        self.shape.iter().copied().product()
    }
    /// Reports whether the tensor has no logical elements.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
    /// Returns the packed byte count of the logical tensor.
    pub fn byte_len(&self) -> usize {
        self.len().saturating_mul(self.dtype.bytes())
    }
    /// Returns the physical byte length of the shared backing allocation.
    pub fn storage_byte_len(&self) -> usize {
        self.data.len()
    }

    /// Reads one scalar using Python-style negative indices.
    pub fn get(&self, indices: &[isize]) -> Result<Scalar, DataError> {
        if indices.len() != self.shape.len() {
            return Err(DataError::Bounds);
        }
        let mut position = self.offset as isize;
        for ((&index, &dimension), &stride) in indices
            .iter()
            .zip(self.shape.iter())
            .zip(self.strides.iter())
        {
            let index = normalize_index(index, dimension)?;
            position = position
                .checked_add(index.checked_mul(stride).ok_or(DataError::Bounds)?)
                .ok_or(DataError::Bounds)?;
        }
        scalar_at_bytes(
            self.dtype,
            &self.data,
            usize::try_from(position).map_err(|_| DataError::Bounds)?,
        )
    }

    /// Produces a shared-storage slice using Python indexing rules.
    pub fn slice(&self, indices: &[AxisIndex]) -> Result<Self, DataError> {
        let layout = slice_layout(&self.shape, &self.strides, self.offset, indices)?;
        Ok(Self {
            dtype: self.dtype,
            shape: Arc::from(layout.shape),
            strides: Arc::from(layout.strides),
            offset: layout.offset,
            data: self.data.clone(),
        })
    }

    /// Replaces a selection in a new contiguous tensor of the original shape.
    ///
    /// Replacement values broadcast across the selected shape using trailing
    /// axes. A rank-zero replacement fills the selection with a single scalar.
    /// Source and replacement may share storage: every replacement value is read
    /// from its original immutable allocation, so overlapping assignment is safe.
    /// Noncontiguous source views are copied in logical row-major order.
    ///
    /// # Errors
    /// Rejects invalid selectors, different element types, or replacement shapes
    /// that cannot broadcast into exactly the selection's shape. On error, source
    /// and replacement remain unchanged.
    ///
    /// ```
    /// use konjure_sdk::data::{AxisIndex, DType, Scalar, Tensor};
    /// let original = Tensor::zeros(DType::I16, vec![2, 3])?;
    /// let row = Tensor::from_scalars(DType::I16, vec![3],
    ///     ["1", "2", "3"].into_iter().map(|text| Scalar::parse(DType::I16, text))
    ///         .collect::<Result<Vec<_>, _>>()?)?;
    /// let updated = original.with_slice(&[AxisIndex::Index(-1)], &row)?;
    /// assert_eq!(updated.get(&[1, 2])?.to_string(), "3");
    /// assert_eq!(original.get(&[1, 2])?.to_string(), "0");
    /// # Ok::<(), konjure_sdk::data::DataError>(())
    /// ```
    pub fn with_slice(&self, indices: &[AxisIndex], replacement: &Self) -> Result<Self, DataError> {
        self.replace_selection(core::iter::once(indices), replacement)
    }

    /// Replaces a selection formed by successive indexing expressions.
    ///
    /// Each path entry indexes the preceding view. All entries are composed
    /// before one logical payload copy, including for a noncontiguous source.
    /// An empty path selects the entire tensor. Broadcasting and immutability
    /// follow [`Self::with_slice`].
    ///
    /// # Errors
    /// Rejects paths longer than [`MAX_RANK`], invalid selectors, mixed types,
    /// and replacement shapes that cannot broadcast into the final selection.
    pub fn with_index_path(
        &self,
        path: &[Vec<AxisIndex>],
        replacement: &Self,
    ) -> Result<Self, DataError> {
        self.replace_selection(path.iter().map(Vec::as_slice), replacement)
    }

    fn replace_selection<'a>(
        &self,
        path: impl IntoIterator<Item = &'a [AxisIndex]>,
        replacement: &Self,
    ) -> Result<Self, DataError> {
        if self.dtype != replacement.dtype {
            return Err(DataError::TypeMismatch);
        }
        let output_strides = contiguous_strides(&self.shape)?;
        let mut selection = SliceLayout {
            shape: self.shape.to_vec(),
            strides: output_strides.clone(),
            offset: 0,
        };
        for (depth, indices) in path.into_iter().enumerate() {
            if depth == MAX_RANK {
                return Err(DataError::InvalidSlice);
            }
            selection = slice_layout(
                &selection.shape,
                &selection.strides,
                selection.offset,
                indices,
            )?;
        }
        if broadcast_shape(&selection.shape, &replacement.shape)? != selection.shape {
            return Err(DataError::ShapeMismatch);
        }
        let prefix = selection.shape.len() - replacement.shape.len();
        let selected_count = checked_count(&selection.shape)?;
        let width = self.dtype.bytes();
        let mut bytes = self.to_le_bytes();
        for mut flat in 0..selected_count {
            let mut destination =
                isize::try_from(selection.offset).map_err(|_| DataError::Bounds)?;
            let mut source = isize::try_from(replacement.offset).map_err(|_| DataError::Bounds)?;
            for axis in (0..selection.shape.len()).rev() {
                let coordinate = flat % selection.shape[axis];
                flat /= selection.shape[axis];
                let coordinate = isize::try_from(coordinate).map_err(|_| DataError::Bounds)?;
                destination = destination
                    .checked_add(
                        coordinate
                            .checked_mul(selection.strides[axis])
                            .ok_or(DataError::Bounds)?,
                    )
                    .ok_or(DataError::Bounds)?;
                if axis >= prefix && replacement.shape[axis - prefix] != 1 {
                    source = source
                        .checked_add(
                            coordinate
                                .checked_mul(replacement.strides[axis - prefix])
                                .ok_or(DataError::Bounds)?,
                        )
                        .ok_or(DataError::Bounds)?;
                }
            }
            let value = scalar_at_bytes(
                self.dtype,
                &replacement.data,
                usize::try_from(source).map_err(|_| DataError::Bounds)?,
            )?;
            let start = usize::try_from(destination)
                .map_err(|_| DataError::Bounds)?
                .checked_mul(width)
                .ok_or(DataError::Bounds)?;
            let end = start.checked_add(width).ok_or(DataError::Bounds)?;
            bytes
                .get_mut(start..end)
                .ok_or(DataError::Bounds)?
                .copy_from_slice(&value.bits().to_le_bytes()[..width]);
        }
        // Every written scalar came from validated immutable storage; the output
        // uses the original logical shape with fresh contiguous byte ownership.
        Ok(Self {
            dtype: self.dtype,
            shape: self.shape.clone(),
            strides: Arc::from(output_strides),
            offset: 0,
            data: Arc::from(bytes),
        })
    }

    /// Changes dimensions while preserving logical row-major order.
    pub fn reshape(&self, shape: Vec<usize>) -> Result<Self, DataError> {
        if checked_count(&shape)? != self.len() {
            return Err(DataError::ShapeMismatch);
        }
        let data = if self.is_contiguous() {
            self.data.clone()
        } else {
            Arc::from(self.to_le_bytes())
        };
        let offset = if self.is_contiguous() { self.offset } else { 0 };
        let strides = contiguous_strides(&shape)?;
        Ok(Self {
            dtype: self.dtype,
            shape: Arc::from(shape),
            strides: Arc::from(strides),
            offset,
            data,
        })
    }

    /// Reorders axes; each old axis must occur exactly once.
    pub fn transpose(&self, axes: &[usize]) -> Result<Self, DataError> {
        if axes.len() != self.shape.len() {
            return Err(DataError::InvalidShape);
        }
        let mut seen = vec![false; axes.len()];
        let mut shape = Vec::with_capacity(axes.len());
        let mut strides = Vec::with_capacity(axes.len());
        for &axis in axes {
            if axis >= axes.len() || seen[axis] {
                return Err(DataError::InvalidShape);
            }
            seen[axis] = true;
            shape.push(self.shape[axis]);
            strides.push(self.strides[axis]);
        }
        Ok(Self {
            dtype: self.dtype,
            shape: Arc::from(shape),
            strides: Arc::from(strides),
            offset: self.offset,
            data: self.data.clone(),
        })
    }

    /// Applies a scalar operation with NumPy trailing-axis broadcasting.
    pub fn binary(&self, operation: BinaryOp, rhs: &Self) -> Result<Self, DataError> {
        if self.dtype != rhs.dtype {
            return Err(DataError::TypeMismatch);
        }
        let shape = broadcast_shape(&self.shape, &rhs.shape)?;
        let count = checked_count(&shape)?;
        let mut bytes = Vec::with_capacity(count * self.dtype.bytes());
        for mut flat in 0..count {
            let mut indices = vec![0isize; shape.len()];
            for axis in (0..shape.len()).rev() {
                indices[axis] = (flat % shape[axis]) as isize;
                flat /= shape[axis];
            }
            let value = self
                .value_broadcast(&indices, &shape)?
                .binary(operation, self_or_rhs(rhs, &indices, &shape)?)?;
            bytes.extend_from_slice(&value.bits().to_le_bytes()[..self.dtype.bytes()]);
        }
        Self::from_le_bytes(self.dtype, shape, bytes)
    }

    /// Applies a scalar operation to every element.
    pub fn scalar_binary(&self, operation: BinaryOp, rhs: Scalar) -> Result<Self, DataError> {
        if rhs.dtype() != self.dtype {
            return Err(DataError::TypeMismatch);
        }
        let mut bytes = Vec::with_capacity(self.byte_len());
        for index in 0..self.len() {
            let value = self.scalar_at_linear(index).binary(operation, rhs)?;
            bytes.extend_from_slice(&value.bits().to_le_bytes()[..self.dtype.bytes()]);
        }
        Self::from_le_bytes(self.dtype, self.shape.to_vec(), bytes)
    }

    /// Converts every scalar to `dtype`.
    pub fn convert(&self, dtype: DType) -> Result<Self, DataError> {
        let mut bytes = Vec::with_capacity(
            self.len()
                .checked_mul(dtype.bytes())
                .ok_or(DataError::AllocationLimit)?,
        );
        for index in 0..self.len() {
            let value = self.scalar_at_linear(index).convert(dtype)?;
            bytes.extend_from_slice(&value.bits().to_le_bytes()[..dtype.bytes()]);
        }
        Self::from_le_bytes(dtype, self.shape.to_vec(), bytes)
    }

    /// Returns the additive reduction, or zero for an empty tensor.
    pub fn sum(&self) -> Result<Scalar, DataError> {
        let mut total = Scalar::zero(self.dtype);
        for index in 0..self.len() {
            total = total.binary(BinaryOp::Add, self.scalar_at_linear(index))?;
        }
        Ok(total)
    }

    /// Returns the element-wise arithmetic negation.
    pub fn neg(&self) -> Result<Self, DataError> {
        let mut bytes = Vec::with_capacity(self.byte_len());
        for index in 0..self.len() {
            let value = self.scalar_at_linear(index).neg()?;
            bytes.extend_from_slice(&value.bits().to_le_bytes()[..self.dtype.bytes()]);
        }
        Self::from_le_bytes(self.dtype, self.shape.to_vec(), bytes)
    }

    /// Multiplies rank-two-or-greater tensors, broadcasting leading batch axes.
    pub fn matmul(&self, rhs: &Self) -> Result<Self, DataError> {
        if self.dtype != rhs.dtype {
            return Err(DataError::TypeMismatch);
        }
        if self.shape.len() < 2 || rhs.shape.len() < 2 {
            return Err(DataError::InvalidShape);
        }
        let (m, k) = (
            self.shape[self.shape.len() - 2],
            self.shape[self.shape.len() - 1],
        );
        if k != rhs.shape[rhs.shape.len() - 2] {
            return Err(DataError::ShapeMismatch);
        }
        let n = rhs.shape[rhs.shape.len() - 1];
        let batch = broadcast_shape(
            &self.shape[..self.shape.len() - 2],
            &rhs.shape[..rhs.shape.len() - 2],
        )?;
        let mut shape = batch.clone();
        shape.push(m);
        shape.push(n);
        let count = checked_count(&shape)?;
        let mut bytes = Vec::with_capacity(count * self.dtype.bytes());
        for flat in 0..count {
            let mut out = linear_indices(flat, &shape);
            let column = out.pop().unwrap_or(0);
            let row = out.pop().unwrap_or(0);
            let mut total = Scalar::zero(self.dtype);
            for inner in 0..k {
                let left = self.matmul_value(&out, row, inner)?;
                let right = rhs.matmul_value(&out, inner, column)?;
                total = total.binary(BinaryOp::Add, left.binary(BinaryOp::Mul, right)?)?;
            }
            bytes.extend_from_slice(&total.bits().to_le_bytes()[..self.dtype.bytes()]);
        }
        Self::from_le_bytes(self.dtype, shape, bytes)
    }

    /// Materializes logical elements in row-major order.
    pub fn to_scalars(&self) -> Vec<Scalar> {
        let mut values = Vec::with_capacity(self.len());
        for flat in 0..self.len() {
            values.push(self.scalar_at_linear(flat));
        }
        values
    }

    /// Materializes logical elements as packed little-endian scalar bytes.
    pub fn to_le_bytes(&self) -> Vec<u8> {
        let width = self.dtype.bytes();
        let mut bytes = Vec::with_capacity(self.byte_len());
        for index in 0..self.len() {
            bytes.extend_from_slice(&self.scalar_at_linear(index).bits().to_le_bytes()[..width]);
        }
        bytes
    }

    fn is_contiguous(&self) -> bool {
        contiguous_strides(&self.shape).is_ok_and(|strides| strides == self.strides.as_ref())
    }
    fn value_broadcast(
        &self,
        output: &[isize],
        output_shape: &[usize],
    ) -> Result<Scalar, DataError> {
        let prefix = output_shape
            .len()
            .checked_sub(self.shape.len())
            .ok_or(DataError::ShapeMismatch)?;
        let mut indices = Vec::with_capacity(self.shape.len());
        for (axis, &dimension) in self.shape.iter().enumerate() {
            indices.push(if dimension == 1 {
                0
            } else {
                output[prefix + axis]
            });
        }
        self.get(&indices)
    }
    fn matmul_value(
        &self,
        output_batch: &[usize],
        row: usize,
        column: usize,
    ) -> Result<Scalar, DataError> {
        let batch_rank = self.shape.len() - 2;
        let output_prefix = output_batch
            .len()
            .checked_sub(batch_rank)
            .ok_or(DataError::ShapeMismatch)?;
        let mut indices = Vec::with_capacity(self.shape.len());
        for axis in 0..batch_rank {
            indices.push(if self.shape[axis] == 1 {
                0
            } else {
                output_batch[output_prefix + axis] as isize
            });
        }
        indices.push(row as isize);
        indices.push(column as isize);
        self.get(&indices)
    }
    fn scalar_at_linear(&self, mut flat: usize) -> Scalar {
        let mut position = self.offset as isize;
        for axis in (0..self.shape.len()).rev() {
            let dimension = self.shape[axis];
            let index = flat % dimension;
            flat /= dimension;
            position += (index as isize) * self.strides[axis];
        }
        scalar_at_bytes(self.dtype, &self.data, position as usize)
            .expect("Tensor constructors and view operations preserve packed-data bounds")
    }
}

struct SliceLayout {
    shape: Vec<usize>,
    strides: Vec<isize>,
    offset: usize,
}

fn slice_layout(
    input_shape: &[usize],
    input_strides: &[isize],
    initial_offset: usize,
    indices: &[AxisIndex],
) -> Result<SliceLayout, DataError> {
    if indices.len() > 2 * MAX_RANK + 1 {
        return Err(DataError::InvalidSlice);
    }
    let ellipses = indices
        .iter()
        .filter(|item| matches!(item, AxisIndex::Ellipsis))
        .count();
    if ellipses > 1 {
        return Err(DataError::InvalidSlice);
    }
    let consumed = indices
        .iter()
        .filter(|item| !matches!(item, AxisIndex::NewAxis | AxisIndex::Ellipsis))
        .count();
    if consumed > input_shape.len() {
        return Err(DataError::InvalidSlice);
    }
    let fill = input_shape.len() - consumed;
    let mut expanded = Vec::new();
    for item in indices {
        if matches!(item, AxisIndex::Ellipsis) {
            expanded.extend((0..fill).map(|_| AxisIndex::Slice {
                start: None,
                stop: None,
                step: 1,
            }));
        } else {
            expanded.push(item.clone());
        }
    }
    if ellipses == 0 {
        expanded.extend((0..fill).map(|_| AxisIndex::Slice {
            start: None,
            stop: None,
            step: 1,
        }));
    }
    let mut source_axis = 0usize;
    let mut offset = initial_offset as isize;
    let mut shape = Vec::new();
    let mut strides = Vec::new();
    for item in expanded {
        match item {
            AxisIndex::NewAxis => {
                shape.push(1);
                strides.push(0);
            }
            AxisIndex::Index(index) => {
                let index = normalize_index(index, input_shape[source_axis])?;
                offset = offset
                    .checked_add(
                        index
                            .checked_mul(input_strides[source_axis])
                            .ok_or(DataError::Bounds)?,
                    )
                    .ok_or(DataError::Bounds)?;
                source_axis += 1;
            }
            AxisIndex::Slice { start, stop, step } => {
                let (start, length) = normalize_slice(start, stop, step, input_shape[source_axis])?;
                offset = offset
                    .checked_add(
                        start
                            .checked_mul(input_strides[source_axis])
                            .ok_or(DataError::Bounds)?,
                    )
                    .ok_or(DataError::Bounds)?;
                shape.push(length);
                strides.push(
                    input_strides[source_axis]
                        .checked_mul(step)
                        .ok_or(DataError::Bounds)?,
                );
                source_axis += 1;
            }
            AxisIndex::Ellipsis => return Err(DataError::InvalidSlice),
        }
    }
    let count = checked_count(&shape)?;
    Ok(SliceLayout {
        shape,
        strides,
        offset: if count == 0 {
            0
        } else {
            usize::try_from(offset).map_err(|_| DataError::Bounds)?
        },
    })
}

fn self_or_rhs(tensor: &Tensor, indices: &[isize], shape: &[usize]) -> Result<Scalar, DataError> {
    tensor.value_broadcast(indices, shape)
}
fn checked_count(shape: &[usize]) -> Result<usize, DataError> {
    if shape.len() > MAX_RANK {
        return Err(DataError::InvalidShape);
    }
    let count = shape.iter().try_fold(1usize, |a, &d| {
        if d > isize::MAX as usize {
            return Err(DataError::InvalidShape);
        }
        a.checked_mul(d).ok_or(DataError::AllocationLimit)
    })?;
    if count > MAX_ELEMENTS {
        Err(DataError::AllocationLimit)
    } else {
        Ok(count)
    }
}
fn contiguous_strides(shape: &[usize]) -> Result<Vec<isize>, DataError> {
    checked_count(shape)?;
    let mut result = vec![0; shape.len()];
    let mut stride = 1usize;
    for axis in (0..shape.len()).rev() {
        result[axis] = isize::try_from(stride).map_err(|_| DataError::AllocationLimit)?;
        stride = stride
            .checked_mul(shape[axis])
            .ok_or(DataError::AllocationLimit)?;
    }
    Ok(result)
}
fn normalize_index(index: isize, len: usize) -> Result<isize, DataError> {
    let len = isize::try_from(len).map_err(|_| DataError::Bounds)?;
    let index = if index < 0 {
        index.checked_add(len).ok_or(DataError::Bounds)?
    } else {
        index
    };
    if index < 0 || index >= len {
        Err(DataError::Bounds)
    } else {
        Ok(index)
    }
}
fn normalize_slice(
    start: Option<isize>,
    stop: Option<isize>,
    step: isize,
    len: usize,
) -> Result<(isize, usize), DataError> {
    if step == 0 {
        return Err(DataError::InvalidSlice);
    }
    let len = isize::try_from(len).map_err(|_| DataError::InvalidSlice)?;
    if len == 0 {
        return Ok((0, 0));
    }
    if step > 0 {
        let mut a = start.unwrap_or(0);
        let mut b = stop.unwrap_or(len);
        if a < 0 {
            a += len;
        }
        if b < 0 {
            b += len;
        }
        a = a.clamp(0, len);
        b = b.clamp(0, len);
        let n = if b <= a { 0 } else { (b - a - 1) / step + 1 };
        Ok((a, n as usize))
    } else {
        let mut a = start.unwrap_or(len - 1);
        let mut b = stop.unwrap_or(-1);
        if start.is_some() && a < 0 {
            a += len;
        }
        if stop.is_some() && b < 0 {
            b += len;
        }
        a = a.clamp(-1, len - 1);
        b = b.clamp(-1, len - 1);
        let n = if a <= b {
            0
        } else {
            usize::try_from(a - b - 1).map_err(|_| DataError::InvalidSlice)? / step.unsigned_abs()
                + 1
        };
        Ok((a, n))
    }
}
fn broadcast_shape(left: &[usize], right: &[usize]) -> Result<Vec<usize>, DataError> {
    let rank = left.len().max(right.len());
    if rank > MAX_RANK {
        return Err(DataError::InvalidShape);
    }
    let mut shape = vec![1; rank];
    for out in 0..rank {
        let a = if out < rank - left.len() {
            1
        } else {
            left[out - (rank - left.len())]
        };
        let b = if out < rank - right.len() {
            1
        } else {
            right[out - (rank - right.len())]
        };
        shape[out] = if a == b {
            a
        } else if a == 1 {
            b
        } else if b == 1 {
            a
        } else {
            return Err(DataError::ShapeMismatch);
        };
    }
    checked_count(&shape)?;
    Ok(shape)
}
fn linear_indices(mut flat: usize, shape: &[usize]) -> Vec<usize> {
    let mut indices = vec![0; shape.len()];
    for axis in (0..shape.len()).rev() {
        if shape[axis] != 0 {
            indices[axis] = flat % shape[axis];
            flat /= shape[axis];
        }
    }
    indices
}
fn scalar_at_bytes(dtype: DType, bytes: &[u8], index: usize) -> Result<Scalar, DataError> {
    let width = dtype.bytes();
    let start = index.checked_mul(width).ok_or(DataError::Bounds)?;
    let slice = bytes
        .get(start..start.checked_add(width).ok_or(DataError::Bounds)?)
        .ok_or(DataError::Bounds)?;
    let mut raw = [0u8; 16];
    raw[..width].copy_from_slice(slice);
    Scalar::from_bits(dtype, u128::from_le_bytes(raw))
}
fn values_to_bytes(dtype: DType, values: &[Scalar]) -> Result<Vec<u8>, DataError> {
    let width = dtype.bytes();
    let mut bytes = Vec::with_capacity(
        values
            .len()
            .checked_mul(width)
            .ok_or(DataError::AllocationLimit)?,
    );
    for value in values {
        if value.dtype() != dtype {
            return Err(DataError::TypeMismatch);
        }
        bytes.extend_from_slice(&value.bits().to_le_bytes()[..width]);
    }
    Ok(bytes)
}

impl Serialize for Tensor {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        // Keep this streaming: a tensor may contain MAX_ELEMENTS values. The
        // generated `wire::Tensor` declares this exact object shape for hosts.
        let mut state = serializer.serialize_struct("Tensor", 3)?;
        state.serialize_field("dtype", &self.dtype)?;
        state.serialize_field("shape", &self.shape.as_ref())?;
        state.serialize_field("data", &DecimalValues(self))?;
        state.end()
    }
}

struct DecimalValues<'a>(&'a Tensor);

impl Serialize for DecimalValues<'_> {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let mut sequence = serializer.serialize_seq(Some(self.0.len()))?;
        for index in 0..self.0.len() {
            sequence.serialize_element(&self.0.scalar_at_linear(index).to_string())?;
        }
        sequence.end()
    }
}

impl<'de> Deserialize<'de> for Tensor {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        #[derive(Deserialize)]
        #[serde(deny_unknown_fields)]
        struct Wire {
            dtype: DType,
            #[serde(deserialize_with = "deserialize_shape")]
            shape: Vec<usize>,
            #[serde(deserialize_with = "deserialize_decimal_data")]
            data: Vec<alloc::string::String>,
        }
        let wire = Wire::deserialize(deserializer)?;
        let count = checked_count(&wire.shape).map_err(D::Error::custom)?;
        if wire.data.len() != count {
            return Err(D::Error::custom(DataError::ShapeMismatch));
        }
        let mut values = Vec::with_capacity(count);
        for text in wire.data {
            values.push(Scalar::parse(wire.dtype, &text).map_err(D::Error::custom)?);
        }
        Tensor::from_scalars(wire.dtype, wire.shape, values).map_err(D::Error::custom)
    }
}

fn deserialize_shape<'de, D>(deserializer: D) -> Result<Vec<usize>, D::Error>
where
    D: Deserializer<'de>,
{
    deserialize_limited_sequence(deserializer, MAX_RANK, PhantomData)
}

fn deserialize_decimal_data<'de, D>(deserializer: D) -> Result<Vec<alloc::string::String>, D::Error>
where
    D: Deserializer<'de>,
{
    deserialize_limited_sequence(deserializer, MAX_ELEMENTS, PhantomData)
}

fn deserialize_limited_sequence<'de, D, T>(
    deserializer: D,
    limit: usize,
    marker: PhantomData<T>,
) -> Result<Vec<T>, D::Error>
where
    D: Deserializer<'de>,
    T: Deserialize<'de>,
{
    struct Limited<T> {
        limit: usize,
        marker: PhantomData<T>,
    }
    impl<'de, T> Visitor<'de> for Limited<T>
    where
        T: Deserialize<'de>,
    {
        type Value = Vec<T>;

        fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(formatter, "an array with at most {} values", self.limit)
        }

        fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
        where
            A: SeqAccess<'de>,
        {
            let mut values = Vec::new();
            while let Some(value) = sequence.next_element()? {
                if values.len() == self.limit {
                    return Err(A::Error::custom(DataError::AllocationLimit));
                }
                values.push(value);
            }
            Ok(values)
        }
    }
    deserializer.deserialize_seq(Limited { limit, marker })
}

impl fmt::Display for Tensor {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "Tensor<{:?}, {:?}>", self.dtype, self.shape)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn negative_slice_reverses() {
        let t = Tensor::from_scalars(
            DType::I32,
            vec![4],
            vec![
                Scalar::parse(DType::I32, "1").unwrap(),
                Scalar::parse(DType::I32, "2").unwrap(),
                Scalar::parse(DType::I32, "3").unwrap(),
                Scalar::parse(DType::I32, "4").unwrap(),
            ],
        )
        .unwrap();
        let r = t
            .slice(&[AxisIndex::Slice {
                start: None,
                stop: None,
                step: -1,
            }])
            .unwrap();
        assert_eq!(r.get(&[0]).unwrap().to_string(), "4");
        assert_eq!(r.get(&[3]).unwrap().to_string(), "1");
        let minimum_step = t
            .slice(&[AxisIndex::Slice {
                start: None,
                stop: None,
                step: isize::MIN,
            }])
            .unwrap();
        assert_eq!(strings(&minimum_step), ["4"]);
    }
    #[test]
    fn broadcast_adds_trailing_axes() {
        let a = Tensor::from_scalars(
            DType::I32,
            vec![2, 1],
            vec![
                Scalar::parse(DType::I32, "1").unwrap(),
                Scalar::parse(DType::I32, "2").unwrap(),
            ],
        )
        .unwrap();
        let b = Tensor::from_scalars(
            DType::I32,
            vec![3],
            (3..6)
                .map(|n| Scalar::parse(DType::I32, &n.to_string()).unwrap())
                .collect(),
        )
        .unwrap();
        let r = a.binary(BinaryOp::Add, &b).unwrap();
        assert_eq!(r.shape(), &[2, 3]);
        assert_eq!(r.get(&[1, 2]).unwrap().to_string(), "7");
    }

    #[test]
    fn zero_axis_broadcast_has_no_division_or_elements() {
        let empty = Tensor::zeros(DType::I8, vec![0, 3]).unwrap();
        let row = Tensor::zeros(DType::I8, vec![1, 3]).unwrap();
        let result = empty.binary(BinaryOp::Add, &row).unwrap();
        assert_eq!(result.shape(), &[0, 3]);
        assert!(result.is_empty());
    }

    #[test]
    fn noncontiguous_negative_slice_materializes_in_logical_order() {
        let source = ints(DType::I16, &[2, 3], &[1, 2, 3, 4, 5, 6]);
        let reverse_columns = source
            .slice(&[
                AxisIndex::Slice {
                    start: None,
                    stop: None,
                    step: 1,
                },
                AxisIndex::Slice {
                    start: None,
                    stop: None,
                    step: -1,
                },
            ])
            .unwrap();
        let reshaped = reverse_columns.reshape(vec![3, 2]).unwrap();
        assert_eq!(strings(&reshaped), ["3", "2", "1", "6", "5", "4"]);
        assert_eq!(
            reverse_columns,
            ints(DType::I16, &[2, 3], &[3, 2, 1, 6, 5, 4])
        );
        let empty = Tensor::zeros(DType::I16, vec![0]).unwrap();
        assert!(
            empty
                .slice(&[AxisIndex::Slice {
                    start: None,
                    stop: None,
                    step: -1,
                }])
                .unwrap()
                .is_empty()
        );
    }

    #[test]
    fn failed_i8_arithmetic_does_not_mutate_shared_view() {
        let source = ints(DType::I8, &[2], &[126, 127]);
        let view = source
            .slice(&[AxisIndex::Slice {
                start: Some(1),
                stop: None,
                step: 1,
            }])
            .unwrap();
        let one = Scalar::parse(DType::I8, "1").unwrap();
        assert_eq!(
            view.scalar_binary(BinaryOp::Add, one),
            Err(DataError::Overflow)
        );
        assert_eq!(view.get(&[0]).unwrap().to_string(), "127");
        assert_eq!(source.get(&[1]).unwrap().to_string(), "127");
    }

    #[test]
    fn reductions_and_matmul_report_checked_overflow() {
        let values = ints(DType::I16, &[3], &[1, 2, 3]);
        assert_eq!(values.sum().unwrap().to_string(), "6");
        let left = ints(DType::I8, &[1, 1], &[127]);
        let right = ints(DType::I8, &[1, 1], &[2]);
        assert_eq!(left.matmul(&right), Err(DataError::Overflow));
    }

    #[test]
    fn matmul_handles_values_batches_and_empty_inner_axes() {
        let left = ints(DType::I16, &[2, 3], &[1, 2, 3, 4, 5, 6]);
        let right = ints(DType::I16, &[3, 2], &[7, 8, 9, 10, 11, 12]);
        assert_eq!(
            strings(&left.matmul(&right).unwrap()),
            ["58", "64", "139", "154"]
        );

        let batched = ints(
            DType::I16,
            &[2, 2, 3],
            &[1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6],
        );
        let result = batched.matmul(&right).unwrap();
        assert_eq!(result.shape(), &[2, 2, 2]);
        assert_eq!(
            strings(&result),
            ["58", "64", "139", "154", "58", "64", "139", "154"]
        );

        let empty_left = Tensor::zeros(DType::I16, vec![2, 0]).unwrap();
        let empty_right = Tensor::zeros(DType::I16, vec![0, 3]).unwrap();
        assert_eq!(
            strings(&empty_left.matmul(&empty_right).unwrap()),
            ["0", "0", "0", "0", "0", "0"]
        );
    }

    #[test]
    fn views_retain_storage_and_support_extended_indexing() {
        let source = ints(DType::I16, &[2, 3], &[1, 2, 3, 4, 5, 6]);
        let transposed = source.transpose(&[1, 0]).unwrap();
        let view = transposed
            .slice(&[
                AxisIndex::NewAxis,
                AxisIndex::Ellipsis,
                AxisIndex::Index(-1),
            ])
            .unwrap();
        assert_eq!(view.shape(), &[1, 3]);
        assert_eq!(view.get(&[-1, -1]).unwrap().to_string(), "6");
        assert_eq!(view.storage_byte_len(), source.storage_byte_len());
        assert!(view.byte_len() < view.storage_byte_len());
        assert_eq!(source.transpose(&[0, 0]), Err(DataError::InvalidShape));
        assert_eq!(source.transpose(&[0]), Err(DataError::InvalidShape));
    }

    #[test]
    fn slice_replacement_broadcasts_into_reversed_columns_without_mutation() {
        let original = ints(DType::I16, &[2, 3], &[1, 2, 3, 4, 5, 6]);
        let alias = original.clone();
        let row = ints(DType::I16, &[3], &[10, 20, 30]);
        let updated = original
            .with_slice(&[AxisIndex::Ellipsis, reverse()], &row)
            .unwrap();
        assert_eq!(strings(&updated), ["30", "20", "10", "30", "20", "10"]);
        assert_eq!(updated.shape(), original.shape());
        assert_eq!(original, alias);
        assert_eq!(strings(&original), ["1", "2", "3", "4", "5", "6"]);
        assert!(!Arc::ptr_eq(&updated.data, &original.data));

        let scalar = ints(DType::I16, &[], &[99]);
        let indexed = original
            .with_slice(&[AxisIndex::Index(-1), AxisIndex::Index(-1)], &scalar)
            .unwrap();
        assert_eq!(strings(&indexed), ["1", "2", "3", "4", "5", "99"]);
        let expanded = original
            .with_slice(&[AxisIndex::NewAxis, AxisIndex::Ellipsis], &row)
            .unwrap();
        assert_eq!(strings(&expanded), ["10", "20", "30", "10", "20", "30"]);
    }

    #[test]
    fn overlapping_and_noncontiguous_replacements_read_original_storage() {
        let original = ints(DType::I16, &[5], &[1, 2, 3, 4, 5]);
        let replacement = original
            .slice(&[AxisIndex::Slice {
                start: None,
                stop: Some(4),
                step: 1,
            }])
            .unwrap();
        let updated = original
            .with_slice(
                &[AxisIndex::Slice {
                    start: Some(1),
                    stop: None,
                    step: 1,
                }],
                &replacement,
            )
            .unwrap();
        assert_eq!(strings(&updated), ["1", "1", "2", "3", "4"]);
        assert_eq!(strings(&original), ["1", "2", "3", "4", "5"]);
        assert_eq!(strings(&replacement), ["1", "2", "3", "4"]);
        let reversed = original.slice(&[reverse()]).unwrap();
        assert_eq!(
            strings(&original.with_slice(&[], &reversed).unwrap()),
            ["5", "4", "3", "2", "1"]
        );

        let matrix = ints(DType::I16, &[2, 3], &[1, 2, 3, 4, 5, 6]);
        let transposed = matrix.transpose(&[1, 0]).unwrap();
        let row = ints(DType::I16, &[2], &[8, 9]);
        let changed = transposed.with_slice(&[AxisIndex::Index(1)], &row).unwrap();
        assert_eq!(changed.shape(), &[3, 2]);
        assert_eq!(strings(&changed), ["1", "4", "8", "9", "3", "6"]);
        assert_eq!(strings(&transposed), ["1", "4", "2", "5", "3", "6"]);
        assert_eq!(strings(&matrix), ["1", "2", "3", "4", "5", "6"]);
    }

    #[test]
    fn replacing_a_subview_materializes_only_its_logical_tensor() {
        let original = ints(DType::I16, &[2, 4], &[1, 2, 3, 4, 5, 6, 7, 8]);
        let view = original
            .slice(&[
                AxisIndex::Index(1),
                AxisIndex::Slice {
                    start: None,
                    stop: None,
                    step: 2,
                },
            ])
            .unwrap();
        let updated = view
            .with_slice(&[AxisIndex::Index(0)], &ints(DType::I16, &[], &[99]))
            .unwrap();
        assert_eq!(strings(&updated), ["99", "7"]);
        assert_eq!(updated.shape(), &[2]);
        assert_eq!(updated.storage_byte_len(), 4);
        assert_eq!(view.storage_byte_len(), 16);
        assert_eq!(strings(&view), ["5", "7"]);
        assert_eq!(original.get(&[1, 0]).unwrap().to_string(), "5");
    }

    #[test]
    fn nested_index_paths_compose_before_replacement() {
        let original = ints(DType::I16, &[2, 3], &[1, 2, 3, 4, 5, 6]);
        let path = vec![
            vec![AxisIndex::Ellipsis, reverse()],
            vec![AxisIndex::Index(1)],
            vec![AxisIndex::Slice {
                start: None,
                stop: Some(2),
                step: 1,
            }],
        ];
        let replacement = ints(DType::I16, &[2], &[8, 9]);
        let updated = original.with_index_path(&path, &replacement).unwrap();
        assert_eq!(strings(&updated), ["1", "2", "3", "4", "9", "8"]);
        assert_eq!(strings(&original), ["1", "2", "3", "4", "5", "6"]);
        let transposed = original.transpose(&[1, 0]).unwrap();
        let updated = transposed
            .with_index_path(&[vec![reverse()], vec![AxisIndex::Index(0)]], &replacement)
            .unwrap();
        assert_eq!(strings(&updated), ["1", "4", "2", "5", "8", "9"]);
        assert_eq!(
            original.with_index_path(&vec![vec![]; MAX_RANK + 1], &replacement),
            Err(DataError::InvalidSlice)
        );
        let scalar = ints(DType::I16, &[], &[7]);
        assert_eq!(
            strings(&original.with_index_path(&[], &scalar).unwrap()),
            ["7", "7", "7", "7", "7", "7"]
        );
    }

    #[test]
    fn slice_replacement_checks_empty_shapes_scalars_types_and_bounds() {
        let empty = Tensor::zeros(DType::I16, vec![0, 3]).unwrap();
        let row = ints(DType::I16, &[3], &[1, 2, 3]);
        assert!(empty.with_slice(&[], &row).unwrap().is_empty());
        assert!(
            empty
                .with_slice(&[reverse()], &ints(DType::I16, &[], &[1]))
                .unwrap()
                .is_empty()
        );
        assert_eq!(
            empty.with_slice(&[], &ints(DType::I16, &[2], &[1, 2])),
            Err(DataError::ShapeMismatch)
        );
        assert_eq!(
            empty.with_slice(&[], &ints(DType::I32, &[3], &[1, 2, 3])),
            Err(DataError::TypeMismatch)
        );
        assert_eq!(
            empty.with_slice(&[AxisIndex::Index(0)], &row),
            Err(DataError::Bounds)
        );

        let scalar = ints(DType::I16, &[], &[7]);
        assert_eq!(
            strings(
                &scalar
                    .with_slice(&[], &ints(DType::I16, &[], &[9]))
                    .unwrap()
            ),
            ["9"]
        );
        assert_eq!(
            scalar.with_slice(&[], &ints(DType::I16, &[1], &[9])),
            Err(DataError::ShapeMismatch)
        );
        let original = ints(DType::I16, &[3], &[1, 2, 3]);
        assert_eq!(
            original.with_slice(&[AxisIndex::Ellipsis, AxisIndex::Ellipsis], &scalar),
            Err(DataError::InvalidSlice)
        );
        assert_eq!(
            original.with_slice(
                &[AxisIndex::Slice {
                    start: None,
                    stop: None,
                    step: 0
                }],
                &scalar
            ),
            Err(DataError::InvalidSlice)
        );
        assert_eq!(
            original.with_slice(&[], &ints(DType::I16, &[1, 3], &[1, 2, 3])),
            Err(DataError::ShapeMismatch)
        );
        let minimum_step = original
            .with_slice(
                &[AxisIndex::Slice {
                    start: None,
                    stop: None,
                    step: isize::MIN,
                }],
                &scalar,
            )
            .unwrap();
        assert_eq!(strings(&minimum_step), ["1", "2", "7"]);
        assert_eq!(strings(&original), ["1", "2", "3"]);
    }

    fn reverse() -> AxisIndex {
        AxisIndex::Slice {
            start: None,
            stop: None,
            step: -1,
        }
    }

    #[test]
    fn serde_rejects_malformed_counts_nonfinite_values_and_large_shapes() {
        let count = serde_json::from_str::<Tensor>(r#"{"dtype":"i8","shape":[2],"data":["1"]}"#);
        assert!(count.is_err());
        let nonfinite =
            serde_json::from_str::<Tensor>(r#"{"dtype":"f32","shape":[1],"data":["NaN"]}"#);
        assert!(nonfinite.is_err());
        let large =
            serde_json::from_str::<Tensor>(r#"{"dtype":"u8","shape":[16777217],"data":[]}"#);
        assert!(large.is_err());
    }

    #[test]
    fn streaming_serde_matches_generated_tensor_wire_shape() {
        let tensor = ints(DType::I128, &[2], &[1, i128::MAX]);
        let expected = serde_json::to_value(crate::data::wire::Tensor {
            dtype: tensor.dtype(),
            shape: tensor.shape().to_vec(),
            data: strings(&tensor),
        })
        .unwrap();
        assert_eq!(serde_json::to_value(&tensor).unwrap(), expected);
    }

    fn ints(dtype: DType, shape: &[usize], values: &[i128]) -> Tensor {
        Tensor::from_scalars(
            dtype,
            shape.to_vec(),
            values
                .iter()
                .map(|value| Scalar::parse(dtype, &value.to_string()).unwrap())
                .collect(),
        )
        .unwrap()
    }

    fn strings(tensor: &Tensor) -> Vec<alloc::string::String> {
        tensor
            .to_scalars()
            .iter()
            .map(ToString::to_string)
            .collect()
    }
}