rten_tensor/

layout.rs

//! Layouts which describe the shape and strides of a tensor.

use std::iter::repeat_n;
use std::ops::Range;

use smallvec::{SmallVec, smallvec};

use crate::errors::{DimensionError, ExpandError, FromDataError, ReshapeError, SliceError};
use crate::index_iterator::{DynIndices, NdIndices};
use crate::overlap::{is_contiguous, may_have_internal_overlap};
use crate::slice_range::{IntoSliceItems, SliceItem};
use crate::type_num::{OptionalUInt, U0, U1, U2, U3, U4, U5, Unknown};

/// Return true if `permutation` is a valid permutation of dimensions for
/// a tensor of rank `ndim`.
pub fn is_valid_permutation(ndim: usize, permutation: &[usize]) -> bool {
    permutation.len() == ndim
        && (0..ndim).all(|dim| permutation.iter().filter(|d| **d == dim).count() == 1)
}

/// Merge dimensions of a layout where possible.
///
/// Two dimensions with sizes N1, N2 and strides S1, S2 can be merged into a
/// single dimension with size N1 * N2 and stride S2 if S1 = N1 * S2;
///
/// Returns a vector of `(size, stride)` tuples for the merged dimensions.
pub(crate) fn merge_axes<S: SizeArray, St: SizeArray>(
    shape: &S,
    strides: &St,
) -> SmallVec<[(usize, usize); 4]> {
    let last_dim = shape.len().saturating_sub(1);
    let (Some(prev_size), Some(prev_stride)) = (shape.get(last_dim), strides.get(last_dim)) else {
        return SmallVec::new();
    };

    let mut merged: SmallVec<[(usize, usize); 4]> = SmallVec::with_capacity(shape.len());
    merged.push((prev_size, prev_stride));

    for (outer_size, outer_stride) in shape.iter().zip(strides.iter()).rev().skip(1) {
        let (inner_size, inner_stride) = merged.last_mut().unwrap();
        let can_merge = outer_size == 1 || (outer_stride == *inner_stride * *inner_size);
        if can_merge {
            *inner_size *= outer_size;
        } else {
            merged.push((outer_size, outer_stride));
        }
    }

    merged.reverse();

    merged
}

/// Generate debug assertion that a dimension index is valid for a layout.
macro_rules! debug_assert_dim_valid {
    ($layout:ident, $dim:expr) => {
        debug_assert!(
            $dim < $layout.ndim(),
            "dim {} out of bounds for tensor with {} dims",
            $dim,
            $layout.ndim()
        )
    };
}

/// Array of dimension sizes or strides.
///
/// This trait is an abstraction around slices and arrays of [`usize`]s. Unlike
/// `AsRef<[usize]>` this does not require the data to actually be stored in
/// memory. This allows for strides which are computed on-demand or stored in a
/// smaller data type (eg. `u32`) and expanded to `usize`.
pub trait SizeArray: Clone + std::fmt::Debug + PartialEq<Self> {
    /// Return the length of the array.
    fn len(&self) -> usize;

    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Return the `index`'th entry.
    fn get(&self, index: usize) -> Option<usize>;

    /// Return an iterator over the entries.
    fn iter(&self) -> impl ExactSizeIterator<Item = usize> + DoubleEndedIterator;
}

impl<T: Clone + AsRef<[usize]> + std::fmt::Debug + PartialEq<Self>> SizeArray for T {
    fn len(&self) -> usize {
        self.as_ref().len()
    }

    fn get(&self, index: usize) -> Option<usize> {
        self.as_ref().get(index).copied()
    }

    fn iter(&self) -> impl ExactSizeIterator<Item = usize> + DoubleEndedIterator {
        self.as_ref().iter().copied()
    }
}

/// Describes the shape and strides of a tensor.
///
/// The `Layout` trait provides methods to query the shape of a tensor, i.e. the
/// number of dimensions and size of each, and the strides which determine the
/// mapping between logical indices in the tensor and offsets in the data storage.
///
/// This trait is implemented for tensor types
/// ([`TensorBase`](crate::TensorBase)), as well as the underlying layout types
/// such as [`NdLayout`] (for tensors with a static number of dimensions) and
/// [`DynLayout`] (for tensors with a dynamic number of dimensions).
pub trait Layout {
    /// Type used to represent the layout's shape.
    type Shape<'a>: SizeArray
    where
        Self: 'a;

    /// Type used to represent the layout's strides.
    type Strides<'a>: SizeArray
    where
        Self: 'a;

    /// Type used to represent indices.
    type Index<'a>: AsRef<[usize]> + Clone + std::fmt::Debug + PartialEq<Self::Index<'a>>;

    /// Iterator over indices in this tensor.
    type Indices;

    /// Map an index to a storage offset, without checking if it is valid for
    /// the tensor's shape.
    ///
    /// This method is not itself unsafe, because it only computes a storage
    /// offset but does not access any data. Using the offset to index into
    /// storage without a bounds check is unsafe however.
    fn offset_unchecked(&self, index: Self::Index<'_>) -> usize {
        index
            .as_ref()
            .iter()
            .zip(self.strides().iter())
            .map(|(idx, stride)| *idx * stride)
            .sum()
    }

    /// Map an index to a storage offset, or return `None` if the index is out
    /// of bounds along any dimension.
    ///
    /// Offsets returned by this method must be less than the layout's minimum
    /// storage length reported by [`min_data_len`](Layout::min_data_len).
    /// If a layout also implements [`TrustedLayout`] then callers can rely
    /// on this to avoid subsequent bounds checks.
    fn offset(&self, index: Self::Index<'_>) -> Option<usize>;

    /// Return the number of dimensions.
    fn ndim(&self) -> usize;

    /// Returns the number of elements in the array.
    fn len(&self) -> usize;

    /// Return true if this layout describes a contiguous tensor, where the
    /// logical order of elements matches the order in which they are stored.
    fn is_contiguous(&self) -> bool {
        is_contiguous(&self.shape(), &self.strides())
    }

    /// Return true if iterating over elements in this layout will visit
    /// elements multiple times.
    fn is_broadcast(&self) -> bool {
        !self.is_empty() && self.strides().iter().any(|s| s == 0)
    }

    /// Returns true if the array has no elements.
    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns an array of the sizes of each dimension.
    fn shape(&self) -> Self::Shape<'_>;

    /// Returns the size of the dimension `dim`.
    fn size(&self, dim: usize) -> usize {
        self.shape().get(dim).expect("invalid dim index")
    }

    /// Returns an array of the strides of each dimension.
    fn strides(&self) -> Self::Strides<'_>;

    /// Returns the offset between adjacent indices along dimension `dim`.
    fn stride(&self, dim: usize) -> usize {
        self.strides().get(dim).expect("invalid dimension")
    }

    /// Return an iterator over all valid indices in this tensor.
    fn indices(&self) -> Self::Indices;

    /// Return true if this layout's shape can be broadcast to the given shape.
    fn can_broadcast_to(&self, target_shape: &[usize]) -> bool {
        if self.ndim() > target_shape.len() {
            return false;
        }

        // For two shapes to be compatible for broadcasting, each dimension must
        // either be the same or be 1.
        //
        // If the tensor has fewer dimensions, pretend that it was prefixed with
        // 1-length dimensions to make the dimension counts equal.
        let target_dims = target_shape[target_shape.len() - self.shape().len()..]
            .iter()
            .copied();

        self.shape()
            .iter()
            .zip(target_dims)
            .all(|(a, b)| a == b || a == 1)
    }

    /// Return true if the tensor/view can be broadcast with another tensor or
    /// view with a given `shape` as part of a binary operation.
    ///
    /// The shape of the result may be larger than either the current shape
    /// or `shape`. eg. If a tensor of shape `[1, 5]` is broadcast with one
    /// of size `[2, 1, 1]` the result has shape `[2, 1, 5]`.
    ///
    /// See <https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md> for
    /// conditions in which broadcasting is allowed.
    fn can_broadcast_with(&self, shape: &[usize]) -> bool {
        // For two shapes to be compatible for broadcasting, each dimension must
        // either be the same or be 1.
        //
        // If the tensor has fewer dimensions, pretend that it was prefixed with
        // 1-length dimensions to make the dimension counts equal.

        let a = self.shape();
        let b = shape;

        let a_pad = b.len().saturating_sub(a.len());
        let b_pad = a.len().saturating_sub(b.len());

        let a_iter = a.iter().rev().chain(repeat_n(1, a_pad));
        let b_iter = b.iter().copied().rev().chain(repeat_n(1, b_pad));

        a_iter.zip(b_iter).all(|(a, b)| a == b || a == 1 || b == 1)
    }

    /// Return the minimum length required for the element data buffer used
    /// with this layout.
    fn min_data_len(&self) -> usize {
        if self.shape().iter().any(|d| d == 0) {
            return 0;
        }
        let max_offset: usize = self
            .shape()
            .iter()
            .zip(self.strides().iter())
            .map(|(size, stride)| (size - 1) * stride)
            .sum();
        max_offset + 1
    }

    /// Return a new layout formed by reshaping this one to `shape`.
    ///
    /// This has the same requirements as
    /// [`reshaped_for_copy`](Layout::reshaped_for_copy) but also requires
    /// that the layout is contiguous.
    fn reshaped_for_view<S: IntoLayout>(&self, shape: S) -> Result<S::Layout, ReshapeError> {
        if !self.is_contiguous() {
            return Err(ReshapeError::NotContiguous);
        }
        self.reshaped_for_copy(shape)
    }

    /// Return a new layout formed by reshaping this one to `shape`.
    fn reshaped_for_copy<S: IntoLayout>(&self, shape: S) -> Result<S::Layout, ReshapeError> {
        let layout = shape.into_layout();
        if layout.len() != self.len() {
            return Err(ReshapeError::LengthMismatch);
        }
        Ok(layout)
    }
}

/// A layout which upholds guarantees on returned storage offsets.
///
/// # Safety
///
/// Layouts which implement this trait promise that any offsets returned by
/// [`offset`](Layout::offset) and
/// [`offset_unchecked`](Layout::offset_unchecked) are less than the than the
/// minimum required storage length reported by
/// [`min_data_len`](Layout::min_data_len). This promise means that the offsets
/// can be used to access elements in a buffer without a bounds check.
pub unsafe trait TrustedLayout: Layout {}

/// Extension methods for layouts.
///
/// These are separate from the [`Layout`] trait to prevent them from being
/// overridden.
pub(crate) trait LayoutExt: Layout {
    /// Return the offset for an index or panic if invalid.
    #[inline]
    fn must_offset(&self, index: Self::Index<'_>) -> usize {
        self.offset(index.clone()).unwrap_or_else(|| {
            panic!(
                "index {:?} out of bounds for shape {:?}",
                index.as_ref(),
                self.shape()
            )
        })
    }
}

impl<L: Layout> LayoutExt for L {}

/// Provides convenience methods for querying the shape and strides of a matrix.
pub trait MatrixLayout {
    fn rows(&self) -> usize;
    fn cols(&self) -> usize;
    fn row_stride(&self) -> usize;
    fn col_stride(&self) -> usize;
}

/// Specifies whether a tensor or view may have an overlapping layout.
///
/// An overlapping layout is one in which multiple valid indices map to the same
/// offset in storage. To comply with Rust's rules for mutable aliases, mutable
/// tensors/views must disallow overlap.
pub enum OverlapPolicy {
    AllowOverlap,
    DisallowOverlap,
}

/// Defines the valid indices for an N-dimensional array and how to map them
/// to offsets in a linear buffer, where N is known at compile time.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct NdLayout<const N: usize> {
    shape: [usize; N],
    strides: [usize; N],
}

impl<const N: usize> Layout for NdLayout<N> {
    type Shape<'a> = [usize; N];
    type Strides<'a> = [usize; N];
    type Index<'a> = [usize; N];
    type Indices = NdIndices<N>;

    fn ndim(&self) -> usize {
        N
    }

    fn len(&self) -> usize {
        self.shape.iter().product()
    }

    #[inline]
    fn offset(&self, index: [usize; N]) -> Option<usize> {
        if !self.index_valid(index) {
            return None;
        }
        Some(self.offset_unchecked(index))
    }

    #[inline]
    fn offset_unchecked(&self, index: [usize; N]) -> usize {
        let mut offset = 0;
        for i in 0..N {
            offset += index[i] * self.strides[i];
        }
        offset
    }

    #[inline]
    fn shape(&self) -> Self::Shape<'_> {
        self.shape
    }

    #[inline]
    fn strides(&self) -> [usize; N] {
        self.strides
    }

    fn indices(&self) -> Self::Indices {
        NdIndices::from_shape(self.shape)
    }
}

unsafe impl<const N: usize> TrustedLayout for NdLayout<N> {}

impl<L: Layout> Layout for &L {
    type Shape<'b>
        = L::Shape<'b>
    where
        Self: 'b;
    type Strides<'b>
        = L::Strides<'b>
    where
        Self: 'b;
    type Index<'b> = L::Index<'b>;
    type Indices = L::Indices;

    fn ndim(&self) -> usize {
        (*self).ndim()
    }

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

    fn offset(&self, index: Self::Index<'_>) -> Option<usize> {
        (*self).offset(index)
    }

    fn offset_unchecked(&self, index: Self::Index<'_>) -> usize {
        (*self).offset_unchecked(index)
    }

    fn shape(&self) -> Self::Shape<'_> {
        (*self).shape()
    }

    fn strides(&self) -> Self::Strides<'_> {
        (*self).strides()
    }

    fn indices(&self) -> Self::Indices {
        (*self).indices()
    }
}

// The `Layout` impl for references proxies to the target, so upholds invariants
// if the target does.
unsafe impl<L: TrustedLayout> TrustedLayout for &L {}

impl MatrixLayout for NdLayout<2> {
    #[inline]
    fn rows(&self) -> usize {
        self.size(0)
    }

    #[inline]
    fn cols(&self) -> usize {
        self.size(1)
    }

    #[inline]
    fn row_stride(&self) -> usize {
        self.stride(0)
    }

    #[inline]
    fn col_stride(&self) -> usize {
        self.stride(1)
    }
}

/// Compute the shape and strides of a layout after slicing with `range`.
///
/// Returns an `(ndim, offset)` tuple for the number of dimensions in the
/// slice and the offset of the first element in the parent view's data.
///
/// This function is generic to allow for specialized variants to be generated
/// when slicing with statically known input or output shape sizes.
fn slice_layout(
    in_shape: impl AsRef<[usize]>,
    in_strides: impl AsRef<[usize]>,
    mut out_shape: impl AsMut<[usize]>,
    mut out_strides: impl AsMut<[usize]>,
    range: &[SliceItem],
) -> Result<(usize, usize), SliceError> {
    let in_shape = in_shape.as_ref();
    let out_shape = out_shape.as_mut();
    let out_strides = out_strides.as_mut();

    let mut ndim = 0;
    let mut offset = 0;

    for (in_dim, (&size, &stride)) in in_shape.iter().zip(in_strides.as_ref()).enumerate() {
        let (offset_adjust, new_size_stride) = match range.get(in_dim) {
            Some(&SliceItem::Index(idx)) => {
                let pos_idx = if idx >= 0 { idx } else { idx + size as isize };
                if pos_idx < 0 || pos_idx >= size as isize {
                    return Err(SliceError::InvalidIndex {
                        axis: in_dim,
                        index: idx,
                        size,
                    });
                }
                (stride * pos_idx as usize, None)
            }
            Some(SliceItem::Range(range)) => {
                let resolved = range.resolve(size).ok_or(SliceError::InvalidRange {
                    axis: in_dim,
                    range: *range,
                    size,
                })?;
                let step: usize = range
                    .step()
                    .try_into()
                    .map_err(|_| SliceError::InvalidStep {
                        axis: in_dim,
                        step: range.step(),
                    })?;
                let new_size = if step == 1 {
                    // Fast path when no custom step is used.
                    resolved.end - resolved.start
                } else {
                    range.index_range(size).steps()
                };
                let new_stride = stride * step;
                (stride * resolved.start, Some((new_size, new_stride)))
            }
            None => (0, Some((size, stride))),
        };

        offset += offset_adjust;
        if let Some((new_size, new_stride)) = new_size_stride {
            out_shape[ndim] = new_size;
            out_strides[ndim] = new_stride;
            ndim += 1;
        }
    }

    if out_shape.contains(&0) {
        offset = 0;
    }

    Ok((ndim, offset))
}

/// Return an iterator over the strides of a layout that broadcasts a view
/// with shape `from_shape` and strides `from_strides` to `to_shape`.
fn broadcast_strides<'a>(
    from_shape: &'a [usize],
    from_strides: &'a [usize],
    to_shape: &'a [usize],
) -> impl Iterator<Item = usize> + 'a {
    let pad = to_shape.len() - from_shape.len();
    repeat_n(0, pad).chain(from_shape.iter().zip(from_strides).enumerate().map(
        move |(i, (size, stride))| {
            if *size == 1 && to_shape[i + pad] > 1 {
                0
            } else {
                *stride
            }
        },
    ))
}

impl<const N: usize> NdLayout<N> {
    /// Convert this layout to one with a dynamic rank.
    pub fn as_dyn(&self) -> DynLayout {
        self.into()
    }

    /// Return true if all components of `index` are in-bounds.
    fn index_valid(&self, index: [usize; N]) -> bool {
        let mut valid = true;
        for i in 0..N {
            valid = valid && index[i] < self.shape[i]
        }
        valid
    }

    /// Return the strides that a contiguous layout with a given shape would
    /// have.
    fn contiguous_strides(shape: [usize; N]) -> [usize; N] {
        let mut strides = [0; N];
        for i in 0..N {
            strides[i] = shape[i + 1..].iter().product();
        }
        strides
    }
}

impl<'a, const N: usize> TryFrom<&'a DynLayout> for NdLayout<N> {
    type Error = DimensionError;

    /// Convert a dynamic layout into a static layout with N dims. Fails if
    /// `value.ndim() != N`.
    fn try_from(value: &'a DynLayout) -> Result<NdLayout<N>, DimensionError> {
        let shape = value.shape();
        let shape: [usize; N] = shape.try_into().map_err(|_| DimensionError {
            actual: shape.len(),
            expected: N,
        })?;
        let strides = value.strides();
        let strides: [usize; N] = strides.try_into().map_err(|_| DimensionError {
            actual: strides.len(),
            expected: N,
        })?;
        Ok(NdLayout { shape, strides })
    }
}

/// Defines the valid indices for an N-dimensional array and how to map them
/// to offsets in a linear buffer, where N can be varied at runtime.
///
/// The layout specifies the size of each dimension of the tensor (the _shape_)
/// and the stride (gap) between offsets in each dimension.
#[derive(Debug, PartialEq)]
pub struct DynLayout {
    /// Array of dimension sizes followed by the corresponding dimension strides.
    ///
    /// Since we always have the same number of stride and shape dims, these
    /// are combined into one array to avoid redundantly storing separate
    /// lengths for each.
    shape_and_strides: SmallVec<[usize; 8]>,
}

impl Clone for DynLayout {
    fn clone(&self) -> DynLayout {
        DynLayout {
            // We implement `Clone` manually here so we can clone
            // `shape_and_strides` using `SmallVec::from_slice` instead of
            // `SmallVec::from`. This is faster for `Copy` types.
            shape_and_strides: SmallVec::from_slice(self.shape_and_strides.as_slice()),
        }
    }
}

impl Layout for DynLayout {
    type Shape<'a> = &'a [usize];
    type Strides<'a> = &'a [usize];
    type Index<'a> = &'a [usize];
    type Indices = DynIndices;

    /// Return the number of elements in the tensor shape described by this layout.
    fn len(&self) -> usize {
        self.shape().iter().product()
    }

    #[inline]
    fn offset(&self, index: Self::Index<'_>) -> Option<usize> {
        let shape = self.shape();
        let strides = self.strides();
        let mut valid = index.as_ref().len() == shape.len();
        let mut offset = 0;
        for (idx, (size, stride)) in index.as_ref().iter().zip(shape.iter().zip(strides)) {
            valid = valid && idx < size;
            offset += idx * stride;
        }
        valid.then_some(offset)
    }

    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Return the number of dimensions.
    #[inline]
    fn ndim(&self) -> usize {
        self.shape_and_strides.len() / 2
    }

    /// Return the sizes of each dimension.
    #[inline]
    fn shape(&self) -> &[usize] {
        &self.shape_and_strides[0..self.ndim()]
    }

    /// Returns the size of the dimension `dim`.
    #[inline]
    fn size(&self, dim: usize) -> usize {
        debug_assert_dim_valid!(self, dim);
        self.shape_and_strides[dim]
    }

    /// Return the stride (offset between elements) in the tensor's element array.
    #[inline]
    fn strides(&self) -> Self::Strides<'_> {
        &self.shape_and_strides[self.ndim()..]
    }

    /// Return the stride for a specific dimension.
    #[inline]
    fn stride(&self, dim: usize) -> usize {
        debug_assert_dim_valid!(self, dim);
        self.shape_and_strides[self.ndim() + dim]
    }

    fn indices(&self) -> DynIndices {
        DynIndices::from_shape(self.shape())
    }
}

unsafe impl TrustedLayout for DynLayout {}

impl DynLayout {
    pub fn make_contiguous(&mut self) {
        self.shape_and_strides = Self::contiguous_shape_and_strides(self.shape());
    }

    fn permute_iter<I: Clone + Iterator<Item = usize>>(&mut self, dims: I) {
        let shape_iter = dims.clone().map(|dim| self.size(dim));
        let stride_iter = dims.map(|dim| self.stride(dim));
        self.shape_and_strides = shape_iter.chain(stride_iter).collect();
    }

    /// Swap the order of dimensions in this layout to the order described by
    /// `dims`.
    fn permute(&mut self, dims: &[usize]) {
        assert!(
            is_valid_permutation(self.ndim(), dims),
            "permutation is invalid"
        );
        self.permute_iter(dims.iter().copied());
    }

    /// Reverse the order of dimensions in this layout.
    fn transpose(&mut self) {
        self.permute_iter((0..self.ndim()).rev());
    }

    /// Create a shape-and-strides array for a contiguous layout.
    fn contiguous_shape_and_strides(shape: &[usize]) -> SmallVec<[usize; 8]> {
        let mut strides_and_shape: SmallVec<[usize; 8]> = SmallVec::from_slice(shape);
        strides_and_shape.resize(shape.len() * 2, 0);
        let mut stride = 1;
        for i in (0..shape.len()).rev() {
            strides_and_shape[shape.len() + i] = stride;
            stride *= shape[i];
        }
        strides_and_shape
    }
}

impl<L: Layout> From<&L> for DynLayout {
    fn from(layout: &L) -> DynLayout {
        let shape: SmallVec<[usize; 5]> = layout.shape().iter().collect();
        let strides: SmallVec<[usize; 5]> = layout.strides().iter().collect();
        DynLayout::from_shape_and_strides(&shape, &strides, OverlapPolicy::AllowOverlap)
            .expect("invalid layout")
    }
}

impl<const N: usize> From<NdLayout<N>> for DynLayout {
    fn from(value: NdLayout<N>) -> DynLayout {
        Self::from(&value)
    }
}

/// MutLayout extends [`Layout`] with methods for creating, modifying and
/// transforming layouts.
///
/// ## Strides and internal overlap
///
/// Rust requires that only one mutable reference can exist for any value. When
/// creating mutable tensor views or iterators, it is therefore important to
/// know whether multiple elements in the layout may map to the same offset.
///
/// Accurately checking this for arbitrary shape and strides is non-trivial. See
/// notes in `mem_overlap.c` in the NumPy source. RTen handles this by using
/// a conservative check for internal overlap when constructing a layout from
/// arbitrary strides. Specifically it sorts dimensions by decreasing stride and
/// then verifies that each dimension fully "steps over" the next one. This
/// allows for layouts which are transposed or have been sliced, but disallows
/// some more complex non-overlapping constructions.
///
/// When constructing a layout via
/// [`from_shape_and_strides`](MutLayout::from_shape_and_strides) the intended
/// usage is specified via an [`OverlapPolicy`].
pub trait MutLayout: Layout + Clone {
    /// Create a layout with custom strides.
    ///
    /// The strides specify the offset gap between successive entries along a
    /// given axis. `overlap` controls whether the layout is allowed to map
    /// multiple indices to the same element. This can be true for immutable
    /// views, but must be false for tensors or views that are mutable.
    fn from_shape_and_strides(
        shape: Self::Shape<'_>,
        strides: Self::Strides<'_>,
        overlap: OverlapPolicy,
    ) -> Result<Self, FromDataError>;

    /// Slice a layout by selecting a single entry from a given axis.
    ///
    /// Returns an `(offset_range, layout)` tuple for the sliced layout.
    fn index_axis(&self, axis: usize, index: usize) -> (Range<usize>, <Self as RemoveDim>::Output)
    where
        Self: RemoveDim,
    {
        assert!(axis < self.ndim());
        assert!(index < self.size(axis));

        let layout = self.remove_dim(axis);

        let range = if layout.is_empty() {
            0..0
        } else {
            let start_offset = self.stride(axis) * index;
            start_offset..start_offset + layout.min_data_len()
        };

        (range, layout)
    }

    /// Move the axis at position `from` to `to` by swapping their strides.
    fn move_axis(&mut self, from: usize, to: usize);

    /// Return a layout with the axes permuted according to the given order.
    fn permuted(&self, order: Self::Index<'_>) -> Self;

    // Modify the size of a dimension. This does not alter the strides.
    fn resize_dim(&mut self, dim: usize, size: usize);

    /// Reverse the order of dimensions. This is equivalent to
    /// `self.permuted([N-1, N-2, ... 0])`.
    fn transposed(&self) -> Self;

    /// Slice the layout and return a static-rank layout.
    ///
    /// Returns a tuple of `(offset_range, sliced_layout)`.
    fn slice<const M: usize>(
        &self,
        range: &[SliceItem],
    ) -> Result<(Range<usize>, NdLayout<M>), SliceError>;

    /// Slice the layout and return a dynamic rank layout.
    ///
    /// Returns a tuple of `(offset_range, sliced_layout)`.
    fn slice_dyn(&self, range: &[SliceItem]) -> Result<(Range<usize>, DynLayout), SliceError>;

    /// Slice the layout along a given axis.
    ///
    /// Returns a tuple of `(offset_range, sliced_layout)`.
    fn slice_axis(
        &self,
        axis: usize,
        range: Range<usize>,
    ) -> Result<(Range<usize>, Self), SliceError> {
        if axis >= self.ndim() {
            return Err(SliceError::InvalidAxis { axis });
        }
        if range.end < range.start || range.end > self.size(axis) {
            return Err(SliceError::InvalidRange {
                axis,
                range: range.into(),
                size: self.size(axis),
            });
        }

        let mut sliced_layout = self.clone();
        sliced_layout.resize_dim(axis, range.len());
        let range = if sliced_layout.is_empty() {
            0..0
        } else {
            let start_offset = range.start * sliced_layout.stride(axis);
            let end_offset = start_offset + sliced_layout.min_data_len();
            start_offset..end_offset
        };
        Ok((range, sliced_layout))
    }

    /// Return a layout with all size-one dimensions removed.
    fn squeezed(&self) -> DynLayout;

    /// Split the layout along the given axis into two.
    ///
    /// Returns a tuple of `(left, right)` where each item is an `(offset_range,
    /// layout)` tuple.
    fn split(&self, axis: usize, mid: usize) -> ((Range<usize>, Self), (Range<usize>, Self));
}

/// Trait for creating a layout from a shape.
pub trait FromShape: Layout {
    /// Create a new contiguous layout with a given shape.
    fn from_shape(shape: Self::Shape<'_>) -> Self;
}

/// Trait for broadcasting a layout from one shape to another.
pub trait BroadcastLayout<L: Layout> {
    /// Broadcast the `self` layout to a given shape.
    fn broadcast<S: IntoLayout<Layout = L>>(&self, shape: S) -> Result<L, ExpandError>;
}

impl<const N: usize, const M: usize> BroadcastLayout<NdLayout<M>> for NdLayout<N> {
    fn broadcast<S: IntoLayout<Layout = NdLayout<M>>>(
        &self,
        shape: S,
    ) -> Result<NdLayout<M>, ExpandError> {
        let shape: [usize; M] = shape.as_ref().try_into().unwrap();
        if !self.can_broadcast_to(&shape) {
            return Err(ExpandError::ShapeMismatch);
        }
        let mut strides = [0usize; M];
        for (i, stride) in broadcast_strides(&self.shape(), &self.strides(), &shape).enumerate() {
            strides[i] = stride;
        }

        Ok(NdLayout { shape, strides })
    }
}

impl<const N: usize> BroadcastLayout<DynLayout> for NdLayout<N> {
    fn broadcast<S: IntoLayout<Layout = DynLayout>>(
        &self,
        shape: S,
    ) -> Result<DynLayout, ExpandError> {
        let dyn_layout: DynLayout = self.into();
        dyn_layout.broadcast(shape.as_ref())
    }
}

impl BroadcastLayout<DynLayout> for DynLayout {
    fn broadcast<S: IntoLayout<Layout = DynLayout>>(
        &self,
        shape: S,
    ) -> Result<DynLayout, ExpandError> {
        let to_shape = shape.as_ref();

        if !self.can_broadcast_to(to_shape) {
            return Err(ExpandError::ShapeMismatch);
        }

        let mut shape_and_strides = SmallVec::with_capacity(to_shape.len() * 2);
        shape_and_strides.extend(to_shape.iter().copied());
        shape_and_strides.extend(broadcast_strides(self.shape(), self.strides(), to_shape));

        Ok(DynLayout { shape_and_strides })
    }
}

impl<const N: usize> BroadcastLayout<NdLayout<N>> for DynLayout {
    fn broadcast<S: IntoLayout<Layout = NdLayout<N>>>(
        &self,
        shape: S,
    ) -> Result<NdLayout<N>, ExpandError> {
        let dyn_broadcast = self.broadcast(shape.as_ref())?;
        let layout = (&dyn_broadcast)
            .try_into()
            .map_err(|_| ExpandError::ShapeMismatch)?;
        Ok(layout)
    }
}

impl<const N: usize> FromShape for NdLayout<N> {
    fn from_shape(shape: [usize; N]) -> Self {
        Self {
            shape,
            strides: Self::contiguous_strides(shape),
        }
    }
}

impl<const N: usize> MutLayout for NdLayout<N> {
    fn from_shape_and_strides(
        shape: Self::Shape<'_>,
        strides: Self::Strides<'_>,
        overlap: OverlapPolicy,
    ) -> Result<Self, FromDataError> {
        let layout = NdLayout { shape, strides };

        match overlap {
            OverlapPolicy::DisallowOverlap => {
                if may_have_internal_overlap(layout.shape, layout.strides) {
                    return Err(FromDataError::MayOverlap);
                }
            }
            OverlapPolicy::AllowOverlap => {}
        }

        Ok(layout)
    }

    fn move_axis(&mut self, from: usize, to: usize) {
        assert!(from < N && to < N);
        let mut dyn_layout = self.as_dyn();
        dyn_layout.move_axis(from, to);
        *self = NdLayout::try_from(&dyn_layout).unwrap();
    }

    fn permuted(&self, dims: [usize; N]) -> NdLayout<N> {
        assert!(is_valid_permutation(N, &dims), "permutation is invalid");
        let mut shape = [0; N];
        let mut strides = [0; N];
        for i in 0..N {
            shape[i] = self.shape[dims[i]];
            strides[i] = self.strides[dims[i]];
        }
        NdLayout { shape, strides }
    }

    fn resize_dim(&mut self, dim: usize, size: usize) {
        self.shape[dim] = size;
    }

    fn transposed(&self) -> NdLayout<N> {
        let dims = std::array::from_fn(|i| N - i - 1);
        self.permuted(dims)
    }

    fn slice<const M: usize>(
        &self,
        range: &[SliceItem],
    ) -> Result<(Range<usize>, NdLayout<M>), SliceError> {
        if self.ndim() < range.len() {
            return Err(SliceError::TooManyDims {
                ndim: self.ndim(),
                range_ndim: range.len(),
            });
        }

        let mut shape: [usize; M] = [0; M];
        let mut strides: [usize; M] = [0; M];

        let (ndim, offset) =
            slice_layout(self.shape, self.strides, &mut shape, &mut strides, range)?;

        if ndim != M {
            return Err(SliceError::OutputDimsMismatch {
                actual: ndim,
                expected: M,
            });
        }

        let layout = NdLayout { shape, strides };
        Ok((offset..offset + layout.min_data_len(), layout))
    }

    fn slice_dyn(&self, range: &[SliceItem]) -> Result<(Range<usize>, DynLayout), SliceError> {
        self.as_dyn().slice_dyn(range)
    }

    fn squeezed(&self) -> DynLayout {
        self.as_dyn().squeezed()
    }

    fn split(&self, axis: usize, mid: usize) -> ((Range<usize>, Self), (Range<usize>, Self)) {
        assert!(axis < self.ndim());
        assert!(mid <= self.size(axis));

        let left_shape = std::array::from_fn(|i| if i == axis { mid } else { self.shape[i] });
        let right_shape = std::array::from_fn(|i| {
            if i == axis {
                self.size(axis) - mid
            } else {
                self.shape[i]
            }
        });

        let left = NdLayout {
            shape: left_shape,
            strides: self.strides,
        };
        let right = NdLayout {
            shape: right_shape,
            strides: self.strides,
        };

        let mid_offset = mid * self.strides[axis];
        let left_offsets = 0..left.min_data_len();
        let end_offset = self.min_data_len();

        let right_offsets = if right.is_empty() {
            end_offset..end_offset
        } else {
            mid_offset..end_offset
        };

        ((left_offsets, left), (right_offsets, right))
    }
}

impl FromShape for DynLayout {
    fn from_shape(shape: &[usize]) -> Self {
        DynLayout {
            shape_and_strides: Self::contiguous_shape_and_strides(shape),
        }
    }
}

impl MutLayout for DynLayout {
    fn from_shape_and_strides(
        shape: &[usize],
        strides: &[usize],
        overlap: OverlapPolicy,
    ) -> Result<Self, FromDataError> {
        let mut shape_and_strides = SmallVec::with_capacity(shape.len() + strides.len());
        shape_and_strides.extend_from_slice(shape);
        shape_and_strides.extend_from_slice(strides);
        let layout = DynLayout { shape_and_strides };

        match overlap {
            OverlapPolicy::DisallowOverlap => {
                if may_have_internal_overlap(layout.shape(), layout.strides()) {
                    return Err(FromDataError::MayOverlap);
                }
            }
            OverlapPolicy::AllowOverlap => {}
        }

        Ok(layout)
    }

    fn move_axis(&mut self, from: usize, to: usize) {
        let ndim = self.ndim();
        assert!(from < ndim && to < ndim);

        let size = self.shape_and_strides.remove(from);
        let stride = self.shape_and_strides.remove(ndim - 1 + from);
        self.shape_and_strides.insert(to, size);
        self.shape_and_strides.insert(ndim + to, stride);
    }

    fn permuted(&self, order: &[usize]) -> DynLayout {
        let mut permuted = self.clone();
        permuted.permute(order);
        permuted
    }

    fn resize_dim(&mut self, dim: usize, size: usize) {
        self.shape_and_strides[dim] = size;
    }

    fn transposed(&self) -> DynLayout {
        let mut transposed = self.clone();
        transposed.transpose();
        transposed
    }

    fn slice<const M: usize>(
        &self,
        range: &[SliceItem],
    ) -> Result<(Range<usize>, NdLayout<M>), SliceError> {
        let (offset_range, dyn_layout) = self.slice_dyn(range)?;
        let nd_layout =
            NdLayout::try_from(&dyn_layout).map_err(|_| SliceError::OutputDimsMismatch {
                actual: dyn_layout.ndim(),
                expected: M,
            })?;
        Ok((offset_range, nd_layout))
    }

    fn slice_dyn(&self, range: &[SliceItem]) -> Result<(Range<usize>, DynLayout), SliceError> {
        if self.ndim() < range.len() {
            return Err(SliceError::TooManyDims {
                ndim: self.ndim(),
                range_ndim: range.len(),
            });
        }

        let out_dims = self.ndim()
            - range
                .iter()
                .filter(|item| matches!(item, SliceItem::Index(_)))
                .count();
        let mut shape_and_strides = smallvec![0; out_dims * 2];
        let (out_shape, out_strides) = shape_and_strides.as_mut_slice().split_at_mut(out_dims);

        let (_ndim, offset) =
            slice_layout(self.shape(), self.strides(), out_shape, out_strides, range)?;

        let layout = Self { shape_and_strides };
        Ok((offset..offset + layout.min_data_len(), layout))
    }

    fn squeezed(&self) -> DynLayout {
        let shape = self.shape().iter().copied().filter(|&size| size != 1);
        let strides = self
            .shape()
            .iter()
            .zip(self.strides())
            .filter_map(|(&size, &stride)| if size != 1 { Some(stride) } else { None });
        DynLayout {
            shape_and_strides: shape.chain(strides).collect(),
        }
    }

    fn split(&self, axis: usize, mid: usize) -> ((Range<usize>, Self), (Range<usize>, Self)) {
        assert!(axis < self.ndim());
        assert!(mid <= self.size(axis));

        let mut left_shape_strides: SmallVec<[usize; 8]> = (0..self.ndim())
            .map(|i| if i == axis { mid } else { self.size(i) })
            .collect();
        left_shape_strides.extend_from_slice(self.strides());

        let mut right_shape_strides: SmallVec<[usize; 8]> = (0..self.ndim())
            .map(|i| {
                if i == axis {
                    self.size(axis) - mid
                } else {
                    self.size(i)
                }
            })
            .collect();
        right_shape_strides.extend_from_slice(self.strides());

        let left = DynLayout {
            shape_and_strides: left_shape_strides,
        };
        let right = DynLayout {
            shape_and_strides: right_shape_strides,
        };

        let mid_offset = mid * self.stride(axis);
        let left_offsets = 0..left.min_data_len();
        let end_offset = self.min_data_len();

        let right_offsets = if right.is_empty() {
            end_offset..end_offset
        } else {
            mid_offset..end_offset
        };

        ((left_offsets, left), (right_offsets, right))
    }
}

/// Trait for shapes which can be used to create a contiguous layout.
///
/// This is implemented for `[usize; N]` for creating static-rank layouts from
/// arrays, and `&[usize]` for creating dynamic-rank layouts from slices.
pub trait IntoLayout: AsRef<[usize]> + std::fmt::Debug {
    /// The type of layout produced from this shape.
    type Layout: MutLayout;

    /// Convert this shape into a contiguous layout.
    fn into_layout(self) -> Self::Layout;
}

impl<const N: usize> IntoLayout for [usize; N] {
    type Layout = NdLayout<N>;

    #[inline]
    fn into_layout(self) -> NdLayout<N> {
        NdLayout::from_shape(self)
    }
}

impl IntoLayout for &[usize] {
    type Layout = DynLayout;

    #[inline]
    fn into_layout(self) -> DynLayout {
        DynLayout::from_shape(self)
    }
}

/// Trait which extends [`MutLayout`] with support for changing the number of
/// dimensions in-place.
///
/// This is only implemented for [`DynLayout`], since layouts that have a static
/// rank cannot change their dimension count at runtime.
pub trait ResizeLayout: MutLayout {
    /// Insert a size-one axis at the given index in the shape. This will have
    /// the same stride as the dimension that follows it.
    fn insert_axis(&mut self, index: usize);

    /// Remove a size-1 axis at the given index.
    ///
    /// Since the axis has size one, this does not alter the number of elements
    /// in the layout or the order in which they are visited.
    ///
    /// Panics if the axis does not have a size of 1.
    #[track_caller]
    fn remove_axis(&mut self, index: usize) {
        assert!(
            self.size(index) == 1,
            "cannot remove axis of size {}",
            self.size(index)
        );
        self.remove_axis_of_any_size(index)
    }

    /// Remove an axis that may have any size.
    ///
    /// If the size of the axis is not one, this will "remove" elements from
    /// the layout.
    fn remove_axis_of_any_size(&mut self, index: usize);

    /// Merge consecutive axes where possible.
    ///
    /// This "simplifies" the layout by minimizing the number of dimensions
    /// while preserving the iteration order.
    fn merge_axes(&mut self);
}

impl ResizeLayout for DynLayout {
    fn insert_axis(&mut self, index: usize) {
        let ndim = self.ndim();
        let new_size = 1;

        // Choose stride for new dimension as if we were inserting it at the
        // beginning. If `dim != 0` then the result is as if we inserted the
        // dim at the start and then permuted the layout.
        let (max_stride, size_for_max_stride) = self
            .strides()
            .iter()
            .copied()
            .zip(self.shape().iter().copied())
            .max_by_key(|(stride, _size)| *stride)
            .unwrap_or((1, 1));
        let new_stride = max_stride * size_for_max_stride;

        self.shape_and_strides.insert(index, new_size);
        self.shape_and_strides.insert(ndim + 1 + index, new_stride);
    }

    fn remove_axis_of_any_size(&mut self, index: usize) {
        self.shape_and_strides.remove(index);
        self.shape_and_strides.remove(self.ndim() + index);
    }

    fn merge_axes(&mut self) {
        let merged = merge_axes(&self.shape(), &self.strides());
        self.shape_and_strides = merged
            .iter()
            .map(|dim| dim.0)
            .chain(merged.iter().map(|dim| dim.1))
            .collect();
    }
}

/// Trait for converting types into indices for use with a given layout.
///
/// Static-rank tensors can be indexed with `[usize; N]` arrays. Dynamic-rank
/// tensors can be indexed with any type that can be converted to an `&[usize]`
/// slice.
pub trait AsIndex<L: Layout> {
    /// Convert `self` into an index for use the layout `L`.
    fn as_index(&self) -> L::Index<'_>;
}

impl<T: AsRef<[usize]>> AsIndex<DynLayout> for T {
    fn as_index(&self) -> &[usize] {
        self.as_ref()
    }
}

impl<const N: usize> AsIndex<NdLayout<N>> for [usize; N] {
    fn as_index(&self) -> [usize; N] {
        *self
    }
}

impl AsIndex<NdLayout<1>> for usize {
    fn as_index(&self) -> [usize; 1] {
        [*self]
    }
}

/// Trait for converting types into shapes for a given layout.
pub trait AsShape<L: Layout> {
    /// Convert `self` into an index for use the layout `L`.
    fn as_shape(&self) -> L::Shape<'_>;
}

impl<T: AsRef<[usize]>> AsShape<DynLayout> for T {
    fn as_shape(&self) -> &[usize] {
        self.as_ref()
    }
}

impl<const N: usize> AsShape<NdLayout<N>> for [usize; N] {
    fn as_shape(&self) -> [usize; N] {
        *self
    }
}

impl AsShape<NdLayout<1>> for usize {
    fn as_shape(&self) -> [usize; 1] {
        [*self]
    }
}

/// Trait that removes one dimension from a layout.
pub trait RemoveDim {
    type Output: MutLayout;

    /// Return a copy of this layout with the dimension at index `dim` removed.
    fn remove_dim(&self, dim: usize) -> Self::Output;
}

impl<R: RemoveDim> RemoveDim for &R {
    type Output = R::Output;

    fn remove_dim(&self, dim: usize) -> Self::Output {
        (*self).remove_dim(dim)
    }
}

impl RemoveDim for DynLayout {
    type Output = DynLayout;

    fn remove_dim(&self, dim: usize) -> DynLayout {
        let ndim = self.ndim();
        assert!(ndim > 0, "cannot remove axis from tensor with 0 dims");

        let shape = (0..ndim - 1).map(|i| {
            if i < dim {
                self.size(i)
            } else {
                self.size(i + 1)
            }
        });
        let strides = (0..ndim - 1).map(|i| {
            if i < dim {
                self.stride(i)
            } else {
                self.stride(i + 1)
            }
        });
        DynLayout {
            shape_and_strides: shape.chain(strides).collect(),
        }
    }
}

macro_rules! impl_remove_dim {
    ($in_dims:expr, $out_dims:expr) => {
        impl RemoveDim for NdLayout<$in_dims> {
            type Output = NdLayout<$out_dims>;

            fn remove_dim(&self, dim: usize) -> Self::Output {
                let shape = std::array::from_fn(|i| {
                    if i < dim {
                        self.shape[i]
                    } else {
                        self.shape[i + 1]
                    }
                });
                let strides = std::array::from_fn(|i| {
                    if i < dim {
                        self.strides[i]
                    } else {
                        self.strides[i + 1]
                    }
                });
                NdLayout { shape, strides }
            }
        }
    };
}

impl_remove_dim!(1, 0);
impl_remove_dim!(2, 1);
impl_remove_dim!(3, 2);
impl_remove_dim!(4, 3);
impl_remove_dim!(5, 4);

/// Trait that inserts one dimension into a layout.
pub trait InsertDim {
    type Output: MutLayout;

    /// Return a copy of this layout with a size-1 dimension inserted at `dim`.
    fn insert_dim(&self, dim: usize) -> Self::Output;
}

impl InsertDim for DynLayout {
    type Output = DynLayout;

    fn insert_dim(&self, dim: usize) -> DynLayout {
        let mut clone = self.clone();
        clone.insert_axis(dim);
        clone
    }
}

macro_rules! impl_insert_dim {
    ($in_dims:expr, $out_dims:expr) => {
        impl InsertDim for NdLayout<$in_dims> {
            type Output = NdLayout<$out_dims>;

            fn insert_dim(&self, dim: usize) -> Self::Output {
                let new_stride = self
                    .shape
                    .iter()
                    .zip(self.strides.iter())
                    .map(|(s, st)| s * st)
                    .max()
                    .unwrap_or(1);

                let shape = std::array::from_fn(|i| {
                    if i < dim {
                        self.shape[i]
                    } else if i == dim {
                        1
                    } else {
                        self.shape[i - 1]
                    }
                });
                let strides = std::array::from_fn(|i| {
                    if i < dim {
                        self.strides[i]
                    } else if i == dim {
                        new_stride
                    } else {
                        self.strides[i - 1]
                    }
                });
                NdLayout { shape, strides }
            }
        }
    };
}

impl_insert_dim!(0, 1);
impl_insert_dim!(1, 2);
impl_insert_dim!(2, 3);
impl_insert_dim!(3, 4);
impl_insert_dim!(4, 5);

/// Trait for slicing a layout with a range.
///
/// `R` is the type of the slice range. `IdxCount` is a marker type indicating
/// the number of items in `R` that are indices, as opposed to ranges.
pub trait SliceWith<R: IntoSliceItems, IdxCount: OptionalUInt> {
    /// The layout produced after slicing.
    type Layout: MutLayout;

    /// Slice the layout with a range.
    ///
    /// Returns a tuple of `(offset_range, sliced_layout)` where `offset_range`
    /// is the range of data from the original view that is used by the slice
    /// and `sliced_layout` is the layout of the sliced view.
    fn slice_with(&self, range: R) -> Result<(Range<usize>, Self::Layout), SliceError>;
}

impl<R: IntoSliceItems, L: MutLayout> SliceWith<R, Unknown> for L {
    type Layout = DynLayout;

    fn slice_with(&self, range: R) -> Result<(Range<usize>, Self::Layout), SliceError> {
        self.slice_dyn(range.into_slice_items().as_ref())
    }
}

impl<R: IntoSliceItems, const N: usize> SliceWith<R, U0> for NdLayout<N> {
    type Layout = NdLayout<N>;

    fn slice_with(&self, range: R) -> Result<(Range<usize>, Self::Layout), SliceError> {
        self.slice(range.into_slice_items().as_ref())
    }
}

macro_rules! impl_slice_with_dynlayout {
    ($range_ndim:ty) => {
        impl<R: IntoSliceItems> SliceWith<R, $range_ndim> for DynLayout {
            type Layout = DynLayout;

            fn slice_with(&self, range: R) -> Result<(Range<usize>, Self::Layout), SliceError> {
                self.slice_dyn(range.into_slice_items().as_ref())
            }
        }
    };
}

impl_slice_with_dynlayout!(U0);
impl_slice_with_dynlayout!(U1);
impl_slice_with_dynlayout!(U2);
impl_slice_with_dynlayout!(U3);
impl_slice_with_dynlayout!(U4);
impl_slice_with_dynlayout!(U5);

macro_rules! impl_slice_with {
    ($ndim:literal, $range_ndim:ty, $out_ndim:literal) => {
        impl<R: IntoSliceItems> SliceWith<R, $range_ndim> for NdLayout<$ndim> {
            type Layout = NdLayout<$out_ndim>;

            fn slice_with(&self, range: R) -> Result<(Range<usize>, Self::Layout), SliceError> {
                self.slice(range.into_slice_items().as_ref())
            }
        }
    };
}

impl_slice_with!(1, U1, 0);
impl_slice_with!(2, U1, 1);
impl_slice_with!(2, U2, 0);
impl_slice_with!(3, U1, 2);
impl_slice_with!(3, U2, 1);
impl_slice_with!(3, U3, 0);
impl_slice_with!(4, U1, 3);
impl_slice_with!(4, U2, 2);
impl_slice_with!(4, U3, 1);
impl_slice_with!(4, U4, 0);
impl_slice_with!(5, U1, 4);
impl_slice_with!(5, U2, 3);
impl_slice_with!(5, U3, 2);
impl_slice_with!(5, U4, 1);
impl_slice_with!(5, U5, 0);

#[cfg(test)]
mod tests {
    use rten_testing::TestCases;

    use std::ops::Range;

    use super::OverlapPolicy;
    use crate::SliceItem;
    use crate::errors::{ReshapeError, SliceError};
    use crate::layout::{DynLayout, FromShape, Layout, MutLayout, NdLayout, ResizeLayout};

    fn layout_with_strides<const N: usize>(shape: [usize; N], strides: [usize; N]) -> NdLayout<N> {
        NdLayout::from_shape_and_strides(shape, strides, OverlapPolicy::AllowOverlap).unwrap()
    }

    #[test]
    fn test_is_broadcast() {
        // Non-empty, contiguous layout
        let layout = DynLayout::from_shape(&[5, 5]);
        assert!(!layout.is_broadcast());

        // Empty layout
        let layout = DynLayout::from_shape(&[5, 0]);
        assert!(!layout.is_broadcast());

        // Broadcasting layout
        let layout =
            DynLayout::from_shape_and_strides(&[5, 5], &[0, 0], OverlapPolicy::AllowOverlap)
                .unwrap();
        assert!(layout.is_broadcast());
    }

    #[test]
    fn test_from_shape_and_strides() {
        #[derive(Debug)]
        struct Case<'a> {
            shape: &'a [usize],
            strides: &'a [usize],
        }

        let cases = [
            // Contiguous layout
            Case {
                shape: &[10, 10],
                strides: &[10, 1],
            },
            // Broadcasting layout
            Case {
                shape: &[10, 10],
                strides: &[10, 0],
            },
        ];

        cases.test_each(|case| {
            let layout = DynLayout::from_shape_and_strides(
                case.shape,
                case.strides,
                OverlapPolicy::AllowOverlap,
            )
            .unwrap();
            assert_eq!(layout.shape(), case.shape);
            assert_eq!(layout.strides(), case.strides);
        })
    }

    #[test]
    fn test_index_axis() {
        #[derive(Debug)]
        struct Case {
            layout: NdLayout<2>,
            axis: usize,
            index: usize,
            expected: (usize, NdLayout<1>), // (start offset, sliced layout)
        }

        let cases = [
            Case {
                layout: NdLayout::from_shape([3, 4]),
                axis: 0,
                index: 1,
                expected: (4, layout_with_strides([4], [1])),
            },
            Case {
                layout: NdLayout::from_shape([3, 4]),
                axis: 1,
                index: 2,
                expected: (2, layout_with_strides([3], [4])),
            },
            // Empty result. The tensor's storage is empty, so the offset must
            // be zero.
            Case {
                layout: NdLayout::from_shape([0, 3]),
                axis: 1,
                index: 1,
                expected: (0, layout_with_strides([0], [3])),
            },
            Case {
                layout: layout_with_strides([3, 0], [1, 3]),
                axis: 0,
                index: 1,
                expected: (0, layout_with_strides([0], [3])),
            },
        ];

        cases.test_each(|case| {
            let Case {
                layout,
                axis,
                index,
                expected,
            } = case;

            let (expected_start, expected_layout) = expected;

            let (offsets, sliced_layout) = layout.index_axis(*axis, *index);
            assert_eq!(sliced_layout, *expected_layout);
            assert_eq!(offsets.start, *expected_start);
            assert_eq!(offsets.len(), expected_layout.min_data_len());

            let (_, sliced_layout_dyn) = layout.as_dyn().index_axis(*axis, *index);
            assert_eq!(sliced_layout_dyn, expected_layout.as_dyn());
        })
    }

    #[test]
    #[should_panic(expected = "axis < self.ndim()")]
    fn test_index_axis_invalid_axis() {
        NdLayout::from_shape([2, 3]).index_axis(2, 0);
    }

    #[test]
    #[should_panic(expected = "index < self.size(axis)")]
    fn test_index_axis_invalid_index() {
        NdLayout::from_shape([2, 3]).index_axis(0, 3);
    }

    #[test]
    fn test_move_axis() {
        let mut layout = DynLayout::from_shape(&[2, 4, 8]);
        assert_eq!(layout.strides(), [32, 8, 1]);

        layout.move_axis(1, 0);
        assert_eq!(layout.shape(), [4, 2, 8]);
        assert_eq!(layout.strides(), [8, 32, 1]);

        layout.move_axis(0, 1);
        assert_eq!(layout.shape(), [2, 4, 8]);
        assert_eq!(layout.strides(), [32, 8, 1]);

        layout.move_axis(2, 1);
        assert_eq!(layout.shape(), [2, 8, 4]);
        assert_eq!(layout.strides(), [32, 1, 8]);
    }

    #[test]
    #[should_panic]
    fn test_move_axis_invalid_from() {
        let mut layout = DynLayout::from_shape(&[2, 4, 8]);
        layout.move_axis(3, 0);
    }

    #[test]
    #[should_panic]
    fn test_move_axis_invalid_to() {
        let mut layout = DynLayout::from_shape(&[2, 4, 8]);
        layout.move_axis(0, 3);
    }

    #[test]
    #[should_panic(expected = "permutation is invalid")]
    fn test_permute_invalid_len() {
        let mut layout = DynLayout::from_shape(&[5, 5]);
        layout.permute(&[1, 0, 3]);
    }

    #[test]
    #[should_panic(expected = "permutation is invalid")]
    fn test_permute_too_few_dims() {
        let mut layout = DynLayout::from_shape(&[5, 5]);
        layout.permute(&[1]);
    }

    #[test]
    #[should_panic(expected = "permutation is invalid")]
    fn test_permute_repeated_dims() {
        let mut layout = DynLayout::from_shape(&[5, 5]);
        layout.permute(&[1, 1]);
    }

    #[test]
    fn test_remove_axis_of_any_size() {
        let shape = [1, 2, 3, 4];
        for d in 0..shape.len() {
            let mut layout = DynLayout::from_shape(&shape);
            let (expected_shape, expected_strides): (Vec<usize>, Vec<usize>) = layout
                .shape()
                .iter()
                .zip(layout.strides())
                .enumerate()
                .filter_map(|(i, (size, stride))| if i != d { Some((size, stride)) } else { None })
                .unzip();

            layout.remove_axis_of_any_size(d);

            assert_eq!(layout.shape(), expected_shape);
            assert_eq!(layout.strides(), expected_strides);
        }
    }

    #[test]
    fn test_reshaped() {
        #[derive(Debug)]
        struct Case<'a> {
            layout: DynLayout,
            new_shape: &'a [usize],
            for_copy: bool,
            error: Option<ReshapeError>,
        }

        let cases = [
            // Reshapes that don't allow copying.
            Case {
                layout: DynLayout::from_shape(&[2, 2]),
                new_shape: &[4],
                for_copy: false,
                error: None,
            },
            Case {
                layout: DynLayout::from_shape(&[2, 2]).transposed(),
                new_shape: &[4],
                for_copy: false,
                error: Some(ReshapeError::NotContiguous),
            },
            Case {
                layout: DynLayout::from_shape(&[2, 2]),
                new_shape: &[3],
                for_copy: false,
                error: Some(ReshapeError::LengthMismatch),
            },
            // Reshapes that do allow copying.
            Case {
                layout: DynLayout::from_shape(&[2, 2]).transposed(),
                new_shape: &[4],
                for_copy: true,
                error: None,
            },
            Case {
                layout: DynLayout::from_shape(&[2, 2]),
                new_shape: &[3],
                for_copy: false,
                error: Some(ReshapeError::LengthMismatch),
            },
        ];

        cases.test_each(|case| {
            let Case {
                layout,
                new_shape,
                for_copy,
                error,
            } = case;

            let reshaped = if *for_copy {
                layout.reshaped_for_copy(*new_shape)
            } else {
                layout.reshaped_for_view(*new_shape)
            };

            assert_eq!(reshaped.as_ref().err(), error.as_ref());
            if let Ok(new_layout) = reshaped {
                assert_eq!(new_layout.shape(), *new_shape);
            }
        })
    }

    #[test]
    fn test_squeezed() {
        let layout = DynLayout::from_shape(&[1, 1, 10, 20]);
        let squeezed = layout.squeezed();
        assert_eq!(squeezed.shape(), &[10, 20]);
        assert_eq!(squeezed.strides(), &[20, 1]);
    }

    #[test]
    fn test_slice_axis() {
        #[derive(Clone, Debug)]
        struct Case<'a> {
            shape: &'a [usize],
            axis: usize,
            range: Range<usize>,
            sliced_shape: &'a [usize],
            offsets: Range<usize>,
        }

        let cases = [Case {
            shape: &[3, 5],
            axis: 1,
            range: 2..4,
            sliced_shape: &[3, 2],
            offsets: 2..14,
        }];

        cases.test_each_clone(|case| {
            let Case {
                shape,
                axis,
                range,
                sliced_shape,
                offsets,
            } = case;

            let layout = DynLayout::from_shape(shape);
            let (offset_range, sliced_layout) = layout.slice_axis(axis, range).unwrap();
            assert_eq!(sliced_layout.shape(), sliced_shape);
            assert_eq!(sliced_layout.strides(), layout.strides());
            assert_eq!(offset_range, offsets);
        })
    }

    #[test]
    fn test_slice_axis_invalid() {
        #[derive(Debug)]
        struct Case<'a> {
            shape: &'a [usize],
            axis: usize,
            range: Range<usize>,
            expected: SliceError,
        }

        let cases = [
            Case {
                shape: &[1, 2, 3],
                axis: 4,
                range: 0..1,
                expected: SliceError::InvalidAxis { axis: 4 },
            },
            Case {
                shape: &[1, 2, 3],
                axis: 0,
                range: 0..2,
                expected: SliceError::InvalidRange {
                    axis: 0,
                    range: (0..2).into(),
                    size: 1,
                },
            },
        ];

        cases.test_each(|case| {
            let layout = DynLayout::from_shape(case.shape);
            let result = layout.slice_axis(case.axis, case.range.clone());
            assert_eq!(result, Err(case.expected.clone()));
        })
    }

    #[test]
    fn test_slice_invalid() {
        #[derive(Debug)]
        struct Case<'a> {
            layout: DynLayout,
            ranges: &'a [SliceItem],
            expected: SliceError,
        }

        let cases = [
            Case {
                layout: DynLayout::from_shape(&[3, 5]),
                ranges: &[SliceItem::Index(4), SliceItem::Index(0)],
                expected: SliceError::InvalidIndex {
                    axis: 0,
                    index: 4,
                    size: 3,
                },
            },
            Case {
                layout: DynLayout::from_shape(&[3, 5]),
                ranges: &[SliceItem::Range((1..4).into()), SliceItem::Index(0)],
                expected: SliceError::InvalidRange {
                    axis: 0,
                    range: (1..4).into(),
                    size: 3,
                },
            },
            Case {
                layout: DynLayout::from_shape(&[3, 5]),
                ranges: &[SliceItem::Index(-4)],
                expected: SliceError::InvalidIndex {
                    axis: 0,
                    index: -4,
                    size: 3,
                },
            },
            Case {
                layout: DynLayout::from_shape(&[3, 5]),
                ranges: &[SliceItem::Range((4..).into()), SliceItem::Index(0)],
                expected: SliceError::InvalidRange {
                    axis: 0,
                    range: (4..).into(),
                    size: 3,
                },
            },
            Case {
                layout: DynLayout::from_shape(&[3, 5]),
                ranges: &[SliceItem::full_range(), SliceItem::range(0, None, -1)],
                expected: SliceError::InvalidStep { axis: 1, step: -1 },
            },
        ];

        cases.test_each(|case| {
            let result = case.layout.slice_dyn(case.ranges);
            assert_eq!(result, Err(case.expected.clone()));
        })
    }

    #[test]
    fn test_size_stride() {
        let layout = DynLayout::from_shape(&[10, 20, 30]);
        for (dim, (&size, &stride)) in layout.shape().iter().zip(layout.strides()).enumerate() {
            assert_eq!(layout.size(dim), size);
            assert_eq!(layout.stride(dim), stride);
        }
    }

    #[test]
    fn test_split() {
        #[derive(Debug)]
        struct Case {
            shape: [usize; 2],
            strides: Option<[usize; 2]>,
            axis: usize,
            mid: usize,
        }

        let mut cases = Vec::new();

        // All combinations of (axis, mid) for a small shape.
        let shape = [4, 2];
        for axis in 0..shape.len() {
            for mid in 0..shape[axis] {
                cases.push(Case {
                    shape,
                    axis,
                    mid,
                    strides: None,
                });
            }
        }

        // Empty layout
        cases.push(Case {
            shape: [0, 0],
            strides: None,
            axis: 0,
            mid: 0,
        });

        // Case where we are splitting a 1-sized dimension with `mid=1` and
        // the stride is larger than the minimum storage length for the layout.
        cases.push(Case {
            shape: [1, 4],
            strides: Some([10, 0]),
            axis: 0,
            mid: 1,
        });

        fn check_split<L: MutLayout>(layout: L, axis: usize, mid: usize) {
            let (left, right) = layout.split(axis, mid);
            let (left_offsets, left_layout) = left;
            let (right_offsets, right_layout) = right;

            assert_eq!(left_layout.strides(), layout.strides());
            assert_eq!(right_layout.strides(), layout.strides());

            assert_eq!(left_offsets.len(), left_layout.min_data_len());
            assert_eq!(right_offsets.len(), right_layout.min_data_len());

            let orig_len = layout.min_data_len();
            assert!(left_offsets.start <= orig_len && left_offsets.end <= orig_len);
            assert!(right_offsets.start <= orig_len && right_offsets.end <= orig_len);

            for i in 0..layout.ndim() {
                assert_eq!(
                    left_layout.size(i),
                    if i == axis { mid } else { layout.size(i) }
                );
                assert_eq!(
                    right_layout.size(i),
                    if i == axis {
                        layout.size(i) - mid
                    } else {
                        layout.size(i)
                    }
                );
            }
        }

        cases.test_each(|case| {
            let Case {
                shape,
                strides,
                axis,
                mid,
            } = case;

            let layout = if let Some(strides) = strides {
                NdLayout::from_shape_and_strides(*shape, *strides, OverlapPolicy::AllowOverlap)
                    .unwrap()
            } else {
                NdLayout::from_shape(*shape)
            };
            let dyn_layout = if let Some(strides) = strides {
                DynLayout::from_shape_and_strides(
                    shape.as_slice(),
                    strides.as_slice(),
                    OverlapPolicy::AllowOverlap,
                )
                .unwrap()
            } else {
                DynLayout::from_shape(shape.as_slice())
            };

            check_split(layout, *axis, *mid);
            check_split(dyn_layout, *axis, *mid);
        })
    }

    #[test]
    fn test_merge_axes() {
        #[derive(Debug)]
        struct Case<'a> {
            shape: &'a [usize],
            strides: &'a [usize],
            merged_shape: &'a [usize],
            merged_strides: &'a [usize],
        }

        let cases = [
            // Empty shape
            Case {
                shape: &[],
                strides: &[],
                merged_shape: &[],
                merged_strides: &[],
            },
            // Vector
            Case {
                shape: &[10],
                strides: &[2],
                merged_shape: &[10],
                merged_strides: &[2],
            },
            // Simple contiguous layout
            Case {
                shape: &[10, 10],
                strides: &[10, 1],
                merged_shape: &[100],
                merged_strides: &[1],
            },
            // Transposed matrix
            Case {
                shape: &[10, 10],
                strides: &[1, 10],
                merged_shape: &[10, 10],
                merged_strides: &[1, 10],
            },
            // Leading 1-sized dims
            Case {
                shape: &[1, 10, 10],
                strides: &[10, 1, 10],
                merged_shape: &[10, 10],
                merged_strides: &[1, 10],
            },
            // Inner 1-sized dims
            Case {
                shape: &[2, 1, 1, 2],
                strides: &[2, 2, 2, 1],
                merged_shape: &[4],
                merged_strides: &[1],
            },
            // Inner 1-sized dims that have been shifted over from the left,
            // ie. where the 1-sized dims where inserted at the left and then
            // shifted over to the middle.
            Case {
                shape: &[2, 1, 1, 2],
                strides: &[2, 4, 4, 1],
                merged_shape: &[4],
                merged_strides: &[1],
            },
        ];

        cases.test_each(|case| {
            let mut layout = DynLayout::from_shape_and_strides(
                case.shape,
                case.strides,
                OverlapPolicy::AllowOverlap,
            )
            .unwrap();
            layout.merge_axes();
            assert_eq!(layout.shape(), case.merged_shape);
            assert_eq!(layout.strides(), case.merged_strides);
        })
    }
}