refactor: add a trait to describe shapes

src/layout/dyn_repr.rs

@@ -0,0 +1,188 @@
+use alloc::boxed::Box;
+use alloc::vec::Vec;
+use core::iter::Cloned;
+use core::ops::{Add, AddAssign, Mul, MulAssign, Sub, SubAssign};
+use core::ops::{Deref, DerefMut, Index, IndexMut};
+use core::slice::Iter;
+
+use crate::layout::rank::DynRank;
+use crate::layout::ranked::Ranked;
+use crate::layout::shape::Shape;
+use crate::Axis;
+
+const CAP: usize = 4;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum DynAxesRepr<T>
+{
+    Inline(usize, [T; CAP]),
+    Alloc(Box<[T]>),
+}
+
+/// An array shape with a dynamic rank.
+pub type DShape = DynAxesRepr<usize>;
+
+impl<T> Deref for DynAxesRepr<T>
+{
+    type Target = [T];
+
+    fn deref(&self) -> &Self::Target
+    {
+        match self {
+            DynAxesRepr::Inline(len, arr) => {
+                debug_assert!(*len <= arr.len());
+                unsafe { arr.get_unchecked(..*len) }
+            }
+            DynAxesRepr::Alloc(items) => items,
+        }
+    }
+}
+
+impl<T> DerefMut for DynAxesRepr<T>
+{
+    fn deref_mut(&mut self) -> &mut Self::Target
+    {
+        match self {
+            DynAxesRepr::Inline(len, arr) => {
+                debug_assert!(*len <= arr.len());
+                unsafe { arr.get_unchecked_mut(..*len) }
+            }
+            DynAxesRepr::Alloc(items) => items,
+        }
+    }
+}
+
+impl<T> Index<usize> for DynAxesRepr<T>
+{
+    type Output = T;
+
+    fn index(&self, index: usize) -> &Self::Output
+    {
+        &(**self)[index]
+    }
+}
+
+impl<T> IndexMut<usize> for DynAxesRepr<T>
+{
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output
+    {
+        &mut (**self)[index]
+    }
+}
+
+impl<T> Index<Axis> for DynAxesRepr<T>
+{
+    type Output = T;
+
+    fn index(&self, index: Axis) -> &Self::Output
+    {
+        self.index(index.0)
+    }
+}
+
+impl<T> IndexMut<Axis> for DynAxesRepr<T>
+{
+    fn index_mut(&mut self, index: Axis) -> &mut Self::Output
+    {
+        self.index_mut(index.0)
+    }
+}
+
+impl<Rhs, T> From<Rhs> for DynAxesRepr<T>
+where
+    Rhs: AsRef<[T]>,
+    T: Default + Copy,
+{
+    fn from(value: Rhs) -> Self
+    {
+        let value = value.as_ref();
+        let n = value.len();
+        if n <= CAP {
+            let mut inline = [T::default(); CAP];
+            inline.split_at_mut(n).0.copy_from_slice(value);
+            Self::Inline(n, inline)
+        } else {
+            Self::Alloc(value.into())
+        }
+    }
+}
+
+impl<T> Ranked for DynAxesRepr<T>
+{
+    type NDim = DynRank;
+
+    fn ndim(&self) -> usize
+    {
+        match self {
+            DynAxesRepr::Inline(d, _) => d.clone(),
+            DynAxesRepr::Alloc(items) => items.len(),
+        }
+    }
+}
+
+macro_rules! impl_op {
+    ($op_trait:ty, $op_fn:ident, $op_assign_trait:ty, $op_assign_fn:ident) => {
+        /// *Panics* if the two dimensionalities are different
+        impl<Rhs> $op_trait for DShape
+        where
+            Rhs: Into<Self>,
+        {
+            type Output = DShape;
+
+            fn $op_fn(self, rhs: Rhs) -> <Self as $op_trait>::Output {
+                let mut output = self.clone();
+                output.$op_assign_fn(rhs);
+                output
+            }
+        }
+
+        /// *Panics* if the two dimensionalities are different
+        impl<Rhs> $op_assign_trait for DShape
+        where
+            Rhs: Into<Self>,
+        {
+            fn $op_assign_fn(&mut self, rhs: Rhs) {
+                let other = rhs.into();
+                for i in 0..self.ndim().max(other.ndim()) {
+                    self[i].$op_assign_fn(other[i]);
+                }
+            }
+        }
+    };
+}
+
+impl_op!(Add<Rhs>, add, AddAssign<Rhs>, add_assign);
+impl_op!(Sub<Rhs>, sub, SubAssign<Rhs>, sub_assign);
+impl_op!(Mul<Rhs>, mul, MulAssign<Rhs>, mul_assign);
+
+impl<T> IntoIterator for DynAxesRepr<T>
+where T: Clone
+{
+    type Item = T;
+    type IntoIter = alloc::vec::IntoIter<T>;
+
+    fn into_iter(self) -> Self::IntoIter
+    {
+        match self {
+            DynAxesRepr::Inline(len, arr) => Vec::from(arr[..len].to_vec()).into_iter(),
+            DynAxesRepr::Alloc(b) => b.into_vec().into_iter(),
+        }
+    }
+}
+
+impl Shape for DynAxesRepr<usize>
+{
+    type Iter<'a>
+        = Cloned<Iter<'a, usize>>
+    where Self: 'a;
+
+    fn axis_len(&self, axis: usize) -> usize
+    {
+        self[axis]
+    }
+
+    fn iter(&self) -> Self::Iter<'_>
+    {
+        (**self).iter().cloned()
+    }
+}

src/layout/mod.rs

@@ -11,6 +11,46 @@
 mod bitset;
 pub mod rank;
 pub mod ranked;
+mod shape;
+mod n_repr;
+mod dyn_repr;
+
+use core::any::type_name;
+use core::error::Error;
+use core::fmt::{Debug, Display};
+use core::marker::PhantomData;
+
+use crate::layout::ranked::Ranked;
 
 #[allow(deprecated)]
 pub use bitset::{Layout, LayoutBitset};
+pub use dyn_repr::DShape;
+pub use n_repr::NShape;
+pub use shape::Shape;
+
+/// The error type for dealing with shapes and strides
+#[derive(Debug, Clone, Copy)]
+pub enum ShapeStrideError<S>
+{
+    /// Out of bounds; specifically, using an index that is larger than the dimensionality of the shape or strides `S`.
+    OutOfBounds(PhantomData<S>, usize),
+    /// The error when trying to construct or mutate a shape or strides with the wrong dimensionality value.
+    RankMismatch(PhantomData<S>, usize),
+    /// The desired shape would represent an array with more elements than `isize::MAX`
+    ShapeOverflow,
+}
+
+impl<S: Ranked> Display for ShapeStrideError<S>
+{
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+    {
+        match self {
+            ShapeStrideError::OutOfBounds(_, idx) =>
+                write!(f, "Index {idx} is larger than the dimensionality of {}", type_name::<S>()),
+            ShapeStrideError::RankMismatch(_, rank) => write!(f, "{} has a rank of {}, which is incompatible with requested rank of {rank}", type_name::<S>(), type_name::<S::NDim>()),
+            ShapeStrideError::ShapeOverflow => write!(f, "The desired shape would represent an array with more elements than `usize::MAX`")
+        }
+    }
+}
+
+impl<S: Debug + Ranked> Error for ShapeStrideError<S> {}