move the dimensioned indexing out of spirv_backend and into coremem_types

this allows us to use it from the CPU implementation.

crates/coremem/src/sim/spirv/cpu.rs

@@ -1,5 +1,4 @@
-use std::ops::Index;
+use coremem_types::dim::DimensionedSlice;
-
 use coremem_types::mat::Material;
 use coremem_types::real::Real;
 use coremem_types::step::{SimMeta, StepEContext, StepHContext, VolumeSampleNeg, VolumeSamplePos};
@@ -65,17 +64,27 @@ fn step_e_cell<R: Real, M: Material<R>>(
     h: &[Vec3<R>],
     m: &[Vec3<R>],
 ) {
-    let flat_idx = flat_idx(meta.dim, idx);
+    let dim = meta.dim;
-    let step_e_context = StepEContext {
+    let stim_e_matrix = DimensionedSlice::new(dim, stim_e);
-        inv_feature_size: meta.inv_feature_size(),
+    let mat_matrix = DimensionedSlice::new(dim, mat);
-        time_step: meta.time_step(),
+    let mut e_matrix = DimensionedSlice::new(dim, e);
-        stim_e: stim_e[flat_idx].cast(),
+    let h_matrix = DimensionedSlice::new(dim, h);
-        mat: &mat[flat_idx],
+
-        in_h: sample_neg(h, meta.dim(), idx),
+    let stim_e = stim_e_matrix[idx];
-        in_e: e[flat_idx],
+    let mat = &mat_matrix[idx];
+    let in_e = e_matrix[idx];
+    let in_h = VolumeSampleNeg::from_indexable(&h_matrix, idx);
+
+    let update_state = StepEContext {
+        inv_feature_size: meta.inv_feature_size,
+        time_step: meta.time_step,
+        stim_e: stim_e.cast(),
+        mat,
+        in_h,
+        in_e,
     };
-    let new_e = step_e_context.step_e();
+    let new_e = update_state.step_e();
-    e[flat_idx] = new_e;
+    e_matrix[idx] = new_e;
 }
 fn step_h_cell<R: Real, M: Material<R>>(
     idx: Vec3u,
@@ -86,19 +95,32 @@ fn step_h_cell<R: Real, M: Material<R>>(
     h: &mut [Vec3<R>],
     m: &mut [Vec3<R>],
 ) {
-    let flat_idx = flat_idx(meta.dim, idx);
+    let dim = meta.dim;
-    let step_h_context = StepHContext {
+
-        inv_feature_size: meta.inv_feature_size(),
+    let stim_h_matrix = DimensionedSlice::new(dim, stim_h);
-        time_step: meta.time_step(),
+    let mat_matrix = DimensionedSlice::new(dim, mat);
-        stim_h: stim_h[flat_idx].cast(),
+    let e_matrix = DimensionedSlice::new(dim, e);
-        mat: &mat[flat_idx],
+    let mut h_matrix = DimensionedSlice::new(dim, h);
-        in_e: sample_pos(e, meta.dim(), idx),
+    let mut m_matrix = DimensionedSlice::new(dim, m);
-        in_h: h[flat_idx],
+
-        in_m: m[flat_idx],
+    let stim_h = stim_h_matrix[idx];
+    let mat = &mat_matrix[idx];
+    let in_e = VolumeSamplePos::from_indexable(&e_matrix, dim, idx);
+    let in_h = h_matrix[idx];
+    let in_m = m_matrix[idx];
+
+    let update_state = StepHContext {
+        inv_feature_size: meta.inv_feature_size,
+        time_step: meta.time_step,
+        stim_h: stim_h.cast(),
+        mat,
+        in_e,
+        in_h,
+        in_m,
     };
-    let (new_h, new_m) = step_h_context.step_h();
+    let (new_h, new_m) = update_state.step_h();
-    h[flat_idx] = new_h;
+    h_matrix[idx] = new_h;
-    m[flat_idx] = new_m;
+    m_matrix[idx] = new_m;
 }
 
 fn apply_all_cells<F: FnMut(Vec3u)>(dim: Vec3u, mut f: F) {
@@ -111,51 +133,3 @@ fn apply_all_cells<F: FnMut(Vec3u)>(dim: Vec3u, mut f: F) {
     }
 }
 
-/// provides (x,y,z)-based indexing into a flat memory view
-struct DimIndexer<'a, T> {
-    items: &'a [T],
-    dim: Vec3u,
-}
-
-impl<'a, T> DimIndexer<'a, T> {
-    fn new(items: &'a [T], dim: Vec3u) -> Self {
-        Self { items, dim }
-    }
-}
-
-impl<'a, T> Index<Vec3u> for DimIndexer<'a, T> {
-    type Output = T;
-    fn index(&self, idx: Vec3u) -> &Self::Output {
-        let idx = flat_idx(self.dim, idx);
-        &self.items[idx]
-    }
-}
-
-fn flat_idx(dim: Vec3u, idx: Vec3u) -> usize {
-    let dx = dim.x() as usize;
-    let dy = dim.y() as usize;
-    // let dz = dim.z() as usize;
-    let ix = idx.x() as usize;
-    let iy = idx.y() as usize;
-    let iz = idx.z() as usize;
-    let effective_y = iz * dy + iy;
-    let effective_x = effective_y * dx + ix;
-    effective_x
-}
-
-fn sample_pos<R: Copy + Default>(arr: &[Vec3<R>], dim: Vec3u, idx: Vec3u) -> VolumeSamplePos<R> {
-    VolumeSamplePos::from_indexable(&DimIndexer::new(arr, dim), dim, idx)
-}
-
-fn sample_neg<R: Copy + Default>(arr: &[Vec3<R>], dim: Vec3u, idx: Vec3u) -> VolumeSampleNeg<R> {
-    VolumeSampleNeg::from_indexable(&DimIndexer::new(arr, dim), idx)
-}
-
-#[cfg(test)]
-mod test {
-    #[test]
-    fn flat_idx_() {
-        // TODO: test all these helper methods!
-        unimplemented!()
-    }
-}

crates/spirv_backend/src/support.rs

@@ -1,8 +1,9 @@
 use core::ops::{Index, IndexMut};
 
-use spirv_std::RuntimeArray;
+// TODO: remove this re-export
+pub use coremem_types::dim::DimensionedSlice;
 
-use coremem_types::vec::Vec3u;
+use spirv_std::RuntimeArray;
 
 pub struct SizedArray<T> {
     items: T,
@@ -45,62 +46,3 @@ impl<'a, T> IndexMut<usize> for SizedArray<&'a mut RuntimeArray<T>> {
     }
 }
 
-pub struct DimensionedSlice<T> {
-    dim: Vec3u,
-    items: T,
-}
-
-impl<T> DimensionedSlice<T> {
-    pub fn new(dim: Vec3u, items: T) -> Self {
-        Self { dim, items }
-    }
-}
-
-impl<'a, T: Index<usize> + ?Sized> Index<Vec3u> for DimensionedSlice<&'a T> {
-    type Output=T::Output;
-
-    fn index(&self, idx: Vec3u) -> &Self::Output {
-        let idx = index(idx, self.dim);
-        &self.items[idx]
-    }
-}
-
-impl<'a, T: Index<usize> + ?Sized> Index<Vec3u> for DimensionedSlice<&'a mut T> {
-    type Output=T::Output;
-
-    fn index(&self, idx: Vec3u) -> &Self::Output {
-        let idx = index(idx, self.dim);
-        &self.items[idx]
-    }
-}
-
-impl<'a, T: IndexMut<usize> + ?Sized> IndexMut<Vec3u> for DimensionedSlice<&'a mut T> {
-    fn index_mut(&mut self, idx: Vec3u) -> &mut Self::Output {
-        let idx = index(idx, self.dim);
-        &mut self.items[idx]
-    }
-}
-
-fn index(loc: Vec3u, dim: Vec3u) -> usize {
-    ((loc.z()*dim.y() + loc.y())*dim.x() + loc.x()) as usize
-}
-
-
-#[cfg(test)]
-mod test {
-    use super::*;
-    #[test]
-    fn test_index() {
-        let dim = Vec3u::new(2, 3, 7);
-        assert_eq!(index(Vec3u::new(0, 0, 0), dim), 0);
-        assert_eq!(index(Vec3u::new(1, 0, 0), dim), 1);
-        assert_eq!(index(Vec3u::new(0, 1, 0), dim), 2);
-        assert_eq!(index(Vec3u::new(1, 1, 0), dim), 3);
-        assert_eq!(index(Vec3u::new(0, 2, 0), dim), 4);
-        assert_eq!(index(Vec3u::new(0, 0, 1), dim), 6);
-        assert_eq!(index(Vec3u::new(1, 0, 1), dim), 7);
-        assert_eq!(index(Vec3u::new(0, 1, 1), dim), 8);
-        assert_eq!(index(Vec3u::new(1, 2, 1), dim), 11);
-        assert_eq!(index(Vec3u::new(1, 2, 2), dim), 17);
-    }
-}

crates/types/src/dim.rs

@@ -0,0 +1,66 @@
+use core::ops::{Index, IndexMut};
+
+use crate::vec::Vec3u;
+
+/// use this to wrap a flat region of memory into something which can be indexed by coordinates in
+/// 3d space.
+pub struct DimensionedSlice<T> {
+    dim: Vec3u,
+    items: T,
+}
+
+impl<T> DimensionedSlice<T> {
+    pub fn new(dim: Vec3u, items: T) -> Self {
+        Self { dim, items }
+    }
+}
+
+impl<'a, T: Index<usize> + ?Sized> Index<Vec3u> for DimensionedSlice<&'a T> {
+    type Output=T::Output;
+
+    fn index(&self, idx: Vec3u) -> &Self::Output {
+        let idx = index(idx, self.dim);
+        &self.items[idx]
+    }
+}
+
+impl<'a, T: Index<usize> + ?Sized> Index<Vec3u> for DimensionedSlice<&'a mut T> {
+    type Output=T::Output;
+
+    fn index(&self, idx: Vec3u) -> &Self::Output {
+        let idx = index(idx, self.dim);
+        &self.items[idx]
+    }
+}
+
+impl<'a, T: IndexMut<usize> + ?Sized> IndexMut<Vec3u> for DimensionedSlice<&'a mut T> {
+    fn index_mut(&mut self, idx: Vec3u) -> &mut Self::Output {
+        let idx = index(idx, self.dim);
+        &mut self.items[idx]
+    }
+}
+
+fn index(loc: Vec3u, dim: Vec3u) -> usize {
+    ((loc.z()*dim.y() + loc.y())*dim.x() + loc.x()) as usize
+}
+
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    #[test]
+    fn test_index() {
+        let dim = Vec3u::new(2, 3, 7);
+        assert_eq!(index(Vec3u::new(0, 0, 0), dim), 0);
+        assert_eq!(index(Vec3u::new(1, 0, 0), dim), 1);
+        assert_eq!(index(Vec3u::new(0, 1, 0), dim), 2);
+        assert_eq!(index(Vec3u::new(1, 1, 0), dim), 3);
+        assert_eq!(index(Vec3u::new(0, 2, 0), dim), 4);
+        assert_eq!(index(Vec3u::new(0, 0, 1), dim), 6);
+        assert_eq!(index(Vec3u::new(1, 0, 1), dim), 7);
+        assert_eq!(index(Vec3u::new(0, 1, 1), dim), 8);
+        assert_eq!(index(Vec3u::new(1, 2, 1), dim), 11);
+        assert_eq!(index(Vec3u::new(1, 2, 2), dim), 17);
+    }
+}
+

crates/types/src/lib.rs

@@ -2,6 +2,7 @@
 #![feature(core_intrinsics)]
 
 pub mod compound;
+pub mod dim;
 pub mod mat;
 pub mod real;
 pub mod step;