feat: add Coordinate and DynamicStorage

Author: siliconlad

src/coordinate.rs

@@ -0,0 +1,61 @@
+use std::fmt;
+use std::ops::Index;
+use std::ops::IndexMut;
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct Coordinate {
+    indices: Vec<usize>,
+}
+
+impl Coordinate {
+    pub fn new(indices: Vec<usize>) -> Self {
+        Self { indices }
+    }
+
+    pub fn len(&self) -> usize {
+        self.indices.len()
+    }
+
+    pub fn iter(&self) -> std::slice::Iter<'_, usize> {
+        self.indices.iter()
+    }
+
+    pub fn insert(&self, index: usize, axis: usize) -> Self {
+        let mut new_indices = self.indices.clone();
+        new_indices.insert(index, axis);
+        Self { indices: new_indices }
+    }
+}
+
+impl Index<usize> for Coordinate {
+    type Output = usize;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.indices[index]
+    }
+}
+
+impl IndexMut<usize> for Coordinate {
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+        &mut self.indices[index]
+    }
+}
+
+impl fmt::Display for Coordinate {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use itertools::Itertools;
+        let idxs = self.indices.iter().map(|&x| format!("{}", x)).join(", ");
+        write!(f, "({})", idxs)
+    }
+}
+
+#[macro_export]
+macro_rules! coord {
+    ($($index:expr),*) => {
+        Coordinate::new(vec![$($index),*])
+    };
+
+    ($index:expr; $count:expr) => {
+        Coordinate::new(vec![$index; $count])
+    };
+}

src/iter.rs

@@ -1,31 +1,35 @@
+use crate::coord;
+use crate::coordinate::Coordinate;
+use crate::shape::Shape;
+
 pub struct IndexIterator {
-    shape: Vec<usize>,
-    current: Vec<usize>,
+    shape: Shape,
+    current: Coordinate,
     done: bool,
 }
 
 impl IndexIterator {
-    pub fn new(shape: &[usize]) -> Self {
-        let current = vec![0; shape.len()];
+    pub fn new(shape: &Shape) -> Self {
+        let current = coord![0; shape.order()];
         IndexIterator {
-            shape: shape.to_vec(),
+            shape: shape.clone(),
             current,
             done: false,
         }
     }
 }
 
 impl Iterator for IndexIterator {
-    type Item = Vec<usize>;
+    type Item = Coordinate;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if self.done || self.shape.len() == 0 {
+        if self.done || self.shape.order() == 0 {
             return None;
         }
 
         let result = self.current.clone();
 
-        for i in (0..self.shape.len()).rev() {
+        for i in (0..self.shape.order()).rev() {
             if self.current[i] + 1 < self.shape[i] {
                 self.current[i] += 1;
                 break;
@@ -45,36 +49,37 @@ impl Iterator for IndexIterator {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::shape;
 
     #[test]
     fn test_index_iterator() {
-        let shape = vec![2, 3];
+        let shape = shape![2, 3].unwrap();
         let mut iter = IndexIterator::new(&shape);
 
-        assert_eq!(iter.next(), Some(vec![0, 0]));
-        assert_eq!(iter.next(), Some(vec![0, 1]));
-        assert_eq!(iter.next(), Some(vec![0, 2]));
-        assert_eq!(iter.next(), Some(vec![1, 0]));
-        assert_eq!(iter.next(), Some(vec![1, 1]));
-        assert_eq!(iter.next(), Some(vec![1, 2]));
+        assert_eq!(iter.next(), Some(coord![0, 0]));
+        assert_eq!(iter.next(), Some(coord![0, 1]));
+        assert_eq!(iter.next(), Some(coord![0, 2]));
+        assert_eq!(iter.next(), Some(coord![1, 0]));
+        assert_eq!(iter.next(), Some(coord![1, 1]));
+        assert_eq!(iter.next(), Some(coord![1, 2]));
         assert_eq!(iter.next(), None);
     }
 
     #[test]
     fn test_index_iterator_single_dimension() {
-        let shape = vec![4];
+        let shape = shape![4].unwrap();
         let mut iter = IndexIterator::new(&shape);
 
-        assert_eq!(iter.next(), Some(vec![0]));
-        assert_eq!(iter.next(), Some(vec![1]));
-        assert_eq!(iter.next(), Some(vec![2]));
-        assert_eq!(iter.next(), Some(vec![3]));
+        assert_eq!(iter.next(), Some(coord![0]));
+        assert_eq!(iter.next(), Some(coord![1]));
+        assert_eq!(iter.next(), Some(coord![2]));
+        assert_eq!(iter.next(), Some(coord![3]));
         assert_eq!(iter.next(), None);
     }
 
     #[test]
     fn test_index_iterator_empty_tensor() {
-        let shape = vec![];
+        let shape = shape![].unwrap();
         let mut iter = IndexIterator::new(&shape);
 
         assert_eq!(iter.next(), None);