flense 0.4.0

//! Concrete dimension storage, used both for sizes and strides.

use core::{
    fmt::Debug,
    iter::Sum,
    ops::{
        Add,
        AddAssign,
        Deref,
        Index,
        Mul,
        MulAssign,
        RangeBounds,
        Sub,
        SubAssign,
    },
    slice::SliceIndex,
};

/// Concrete dimension description.
///
/// Used either for representing strides or representing sizes.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(transparent)]
pub struct Dim<T, const N: usize>(pub(crate) [T; N]);

impl<T, const N: usize> Dim<T, N> {
    /// Get the dimension at `IDX` infallibly.
    ///
    /// Fails to compile if `IDX >= N`.
    #[inline]
    #[must_use]
    pub fn axis<const IDX: usize>(&self) -> &T {
        const {
            assert!(IDX < N, "flense: must get an in-bounds dimension");
        }
        // SAFETY: We just checked that IDX is in bounds for an array of length N.
        unsafe { self.0.get_unchecked(IDX) }
    }

    /// Returns a reference to the element or subslice depending on the type of
    /// index.
    ///
    /// See [`core::slice::get`].
    #[inline]
    #[must_use]
    pub fn get<I>(&self, index: I) -> Option<&I::Output>
    where
        I: SliceIndex<[T]>,
    {
        self.0.get(index)
    }

    /// Returns a reference to the element or subslice depending on the type of
    /// index.
    ///
    /// See [`core::slice::get_unchecked`].
    ///
    /// # Safety
    /// 1. `index` must be in bounds. See the safety documentation on
    ///    [`core::slice::get_unchecked`].
    #[inline]
    #[must_use]
    pub unsafe fn get_unchecked<I>(&self, index: I) -> &I::Output
    where
        I: SliceIndex<[T]>,
    {
        // SAFETY: This function has the same preconditions as this wrapper.
        unsafe { self.0.get_unchecked(index) }
    }
}

impl<T, const N: usize> Index<usize> for Dim<T, N> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        #[expect(clippy::indexing_slicing, reason = "panicking is allowable in Index")]
        &self.0[index]
    }
}

impl<'a, T, const N: usize> Dim<T, N>
where
    T: 'a + Sum<&'a T>,
{
    /// Sum all the values in `self`.
    ///
    /// # Panics
    /// Panics if overflow checks are enabled and the computation overflows.
    #[inline]
    #[must_use]
    pub fn sum(&'a self) -> T {
        self.0.iter().sum()
    }
}

impl<T, const N: usize> Dim<T, N>
where
    T: PartialOrd,
{
    /// Whether `self`'s values are strictly greater than `other`'s in each
    /// dimension.
    ///
    /// This is useful when `self` represents lengths, and `other` represents
    /// indices.
    #[inline]
    #[must_use]
    pub fn contains(&self, other: &Self) -> bool {
        self.0
            .iter()
            .zip(other.0.iter())
            .all(|(lhs, rhs)| lhs > rhs)
    }

    /// Whether `self` has dimensions less than or equal to `other` in each
    /// dimension.
    ///
    /// This is useful when either:
    /// 1. Both `self` and `other` represent lengths, is used to compute whether
    ///    a [`LensBase::try_join`](crate::lens::LensBase::try_join) is possible
    ///    between two lenses.
    /// 2. `self` and `other` are checking for range bound compatibility in
    ///    [`LensBase::slice`](crate::lens::LensBase::slice).
    #[inline]
    #[must_use]
    pub fn is_compatible_with(&self, other: &Self) -> bool {
        self.0
            .iter()
            .zip(other.0.iter())
            .all(|(lhs, rhs)| lhs <= rhs)
    }
}

impl<T, const N: usize> Dim<T, N>
where
    T: Copy + PartialOrd + Ord,
{
    /// Compute the minimum values between `self` and `other` in each dimension.
    #[inline]
    #[must_use]
    pub fn min(&self, other: &Self) -> Self {
        let mut result = self.0;
        for (lhs, rhs) in result.iter_mut().zip(other.0.iter()) {
            *lhs = (*lhs).min(*rhs);
        }
        Self(result)
    }
}

impl<R, const N: usize> Dim<R, N>
where
    R: RangeBounds<usize>,
{
    /// Convert a dimension-of-ranges into two dimensions-of-indices.
    ///
    /// The returned indices are `[inclusive, exclusive)`. Unbounded range edges
    /// convert to `[0, dimensions)`.
    ///
    /// Returns `None` if converting a bound overflows `usize` (an
    /// `Excluded(usize::MAX)` start or `Included(usize::MAX)` end), as those
    /// have no representable exclusive/inclusive index.
    ///
    /// Does not make any attempt to bound "out-of-bounds" values, like `0..100`
    /// where `dimensions` is 5.
    #[inline]
    #[must_use]
    pub fn to_indices(&self, dimensions: Dim<usize, N>) -> Option<(Dim<usize, N>, Dim<usize, N>)> {
        use core::ops::Bound::{
            Excluded,
            Included,
            Unbounded,
        };

        let mut starts = [0; N];
        let mut ends = [0; N];

        for (((range, len), start_out), end_out) in self
            .0
            .iter()
            .zip(dimensions.0)
            .zip(starts.iter_mut())
            .zip(ends.iter_mut())
        {
            *start_out = match range.start_bound() {
                Included(start) => *start,
                Excluded(start) => start.checked_add(1)?,
                Unbounded => 0,
            };
            *end_out = match range.end_bound() {
                Included(end) => end.checked_add(1)?,
                Excluded(end) => *end,
                Unbounded => len,
            };
        }

        Some((Dim(starts), Dim(ends)))
    }
}

impl<T> From<T> for Dim<T, 1> {
    #[inline]
    fn from(value: T) -> Self {
        Self([value])
    }
}

impl<T> Deref for Dim<T, 1> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.0[0]
    }
}

impl<T, const N: usize> From<[T; N]> for Dim<T, N> {
    #[inline]
    fn from(value: [T; N]) -> Self {
        Self(value)
    }
}

impl<const N: usize> TryFrom<Dim<usize, N>> for Dim<isize, N> {
    type Error = <isize as TryFrom<usize>>::Error;

    #[inline]
    fn try_from(value: Dim<usize, N>) -> Result<Self, Self::Error> {
        let mut result = [0; N];
        for (elt, value) in result.iter_mut().zip(value.0) {
            *elt = value.try_into()?;
        }
        Ok(Self(result))
    }
}

macro_rules! impl_dim_op {
    ($trait:ident, $method:ident, $op:tt, $traitassign:ident, $methodassign:ident, $opassign:tt) => {
        impl<T, const N: usize> $trait for Dim<T, N>
        where
            T: Copy + $trait<Output = T>,
        {
            type Output = Self;

            #[inline]
            fn $method(mut self, rhs: Self) -> Self::Output {
                for (lhs, rhs) in self.0.iter_mut().zip(rhs.0) {
                    *lhs = *lhs $op rhs;
                }
                self
            }
        }

        impl<T, const N: usize> $traitassign for Dim<T, N>
        where
            T: $traitassign
        {
            #[inline]
            fn $methodassign(self: &mut Self, rhs: Self) {
                for (lhs, rhs) in self.0.iter_mut().zip(rhs.0) {
                    *lhs $opassign rhs;
                }
            }
        }
    }
}

impl_dim_op!(Add, add, +, AddAssign, add_assign, +=);
impl_dim_op!(Sub, sub, -, SubAssign, sub_assign, -=);
impl_dim_op!(Mul, mul, *, MulAssign, mul_assign, *=);