Add dynamic shape

Also add a Zero impl for NShape

Author: akern40

src/layout/dyn_repr.rs

@@ -1,13 +1,13 @@
 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 num_traits::Zero;
 
 use crate::layout::rank::DynRank;
 use crate::layout::ranked::Ranked;
-use crate::layout::shape::{IntoShape, Shape};
+use crate::layout::shape::Shape;
 use crate::Axis;
 
 const CAP: usize = 4;
@@ -19,19 +19,9 @@ pub enum DynAxesRepr<T>
     Alloc(Box<[T]>),
 }
 
+/// An array shape with a dynamic rank.
 pub type DShape = DynAxesRepr<usize>;
 
-impl<T> DynAxesRepr<T>
-{
-    fn ndim(&self) -> usize
-    {
-        match self {
-            DynAxesRepr::Inline(len, _) => *len,
-            DynAxesRepr::Alloc(items) => items.len(),
-        }
-    }
-}
-
 impl<T> Deref for DynAxesRepr<T>
 {
     type Target = [T];
@@ -52,7 +42,13 @@ impl<T> DerefMut for DynAxesRepr<T>
 {
     fn deref_mut(&mut self) -> &mut Self::Target
     {
-        todo!()
+        match self {
+            DynAxesRepr::Inline(len, arr) => {
+                debug_assert!(*len <= arr.len());
+                unsafe { arr.get_unchecked_mut(..*len) }
+            }
+            DynAxesRepr::Alloc(items) => items,
+        }
     }
 }
 
@@ -92,33 +88,6 @@ impl<T> IndexMut<Axis> for DynAxesRepr<T>
     }
 }
 
-impl<T> Default for DynAxesRepr<T>
-where T: Zero + Copy
-{
-    fn default() -> Self
-    {
-        Self::Inline(1, [T::zero(); 4])
-    }
-}
-
-impl Zero for DShape
-{
-    fn zero() -> Self
-    {
-        Self::Inline(1, [0usize; CAP])
-    }
-
-    fn is_zero(&self) -> bool
-    {
-        self.iter().all(|e| e.is_zero())
-    }
-
-    fn set_zero(&mut self)
-    {
-        self.iter_mut().for_each(|e| e.set_zero());
-    }
-}
-
 impl<Rhs, T> From<Rhs> for DynAxesRepr<T>
 where
     Rhs: AsRef<[T]>,
@@ -156,7 +125,7 @@ macro_rules! impl_op {
         /// *Panics* if the two dimensionalities are different
         impl<Rhs> $op_trait for DShape
         where
-            Rhs: IntoShape<NDim = DynRank>,
+            Rhs: Into<Self>,
         {
             type Output = DShape;
 
@@ -170,10 +139,10 @@ macro_rules! impl_op {
         /// *Panics* if the two dimensionalities are different
         impl<Rhs> $op_assign_trait for DShape
         where
-            Rhs: IntoShape<NDim = DynRank>,
+            Rhs: Into<Self>,
         {
             fn $op_assign_fn(&mut self, rhs: Rhs) {
-                let other = rhs.into_shape();
+                let other = rhs.into();
                 for i in 0..self.ndim().max(other.ndim()) {
                     self[i].$op_assign_fn(other[i]);
                 }
@@ -204,11 +173,16 @@ where T: Clone
 impl Shape for DynAxesRepr<usize>
 {
     type Iter<'a>
-        = Iter<'a, usize>
+        = Cloned<Iter<'a, usize>>
     where Self: 'a;
 
+    fn axis_len(&self, axis: usize) -> usize
+    {
+        self[axis]
+    }
+
     fn iter(&self) -> Self::Iter<'_>
     {
-        (**self).iter()
+        (**self).iter().cloned()
     }
 }

src/layout/mod.rs

@@ -13,6 +13,7 @@ pub mod rank;
 pub mod ranked;
 mod shape;
 mod n_repr;
+mod dyn_repr;
 
 use core::any::type_name;
 use core::error::Error;
@@ -23,6 +24,7 @@ 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;
 

src/layout/n_repr.rs

@@ -171,6 +171,19 @@ where ConstRank<N>: Rank
     }
 }
 
+impl<const N: usize> Zero for NShape<N>
+{
+    fn zero() -> Self
+    {
+        Self([usize::zero(); N])
+    }
+
+    fn is_zero(&self) -> bool
+    {
+        self.0.iter().all(|e| *e == 0usize)
+    }
+}
+
 macro_rules! impl_op {
     ($op_trait:ty, $op_fn:ident, $op_assign_trait:ty, $op_assign_fn:ident) => {
         impl<Rhs, const N: usize> $op_trait for NShape<N>
@@ -270,6 +283,6 @@ impl<const N: usize> Default for NShape<N>
 {
     fn default() -> Self
     {
-        Self([0; N])
+        Self::zero()
     }
 }