feat: Kill Vec allocation in iteration

src/scene/grid.rs

@@ -311,21 +311,19 @@ impl Integrate for Grid {
 impl Grid {
     fn apply_pos_stencils<T>(&mut self, func: T)
     where
-        T: Fn(&mut PosStencilMut<Cell>) + Send + Sync,
+        T: Fn(PosStencilMut<Cell>) + Send + Sync,
     {
         const OFFSETS: [Index2; 4] = [idx!(0, 0), idx!(1, 0), idx!(0, 1), idx!(1, 1)];
 
         for offset in OFFSETS.iter() {
-            let mut stencils: Vec<_> = positive_stencils_mut(
+            positive_stencils_mut(
                 self.cells.as_mut_slice(),
                 self.dim,
                 Some(idx!(1, 1)),
                 Some(self.dim - idx!(1, 1)),
                 Some(*offset),
             )
-            .collect();
-
-            stencils.par_iter_mut().for_each(&func);
+            .for_each(&func);
         }
     }
 
@@ -355,7 +353,7 @@ impl Grid {
         };
 
         debug!(log, "Distribute all 's' factors for total sum.");
-        self.apply_pos_stencils(|s: &mut PosStencilMut<Cell>| {
+        self.apply_pos_stencils(|s: PosStencilMut<Cell>| {
             // This cell (1: pos, 0: x)  <-- s from pos x-neighbor.
             s.cell.s_nbs[1][0] = s_factor(s.neighbors[0]);
             // Pos. x-neighbor (0: neg, 0: x) <-- s from this cell.
@@ -383,18 +381,20 @@ impl Grid {
             } else {
                 warn!(
                     log,
-                    "Cell with index: '{}' contains only fluid neighbors.", c.index()
+                    "Cell with index: '{}' contains only fluid neighbors.",
+                    c.index()
                 );
                 0.0
             };
         });
 
         for _iter in 0..iterations {
-            self.apply_pos_stencils(|s: &mut PosStencilMut<Cell>| {
+            self.apply_pos_stencils(|s: PosStencilMut<Cell>| {
                 if s.cell.s_tot_inv == 0.0 {
                     debug!(
                         log,
-                        "Cell with index: '{}' contains only fluid neighbors.", s.cell.index()
+                        "Cell with index: '{}' contains only fluid neighbors.",
+                        s.cell.index()
                     );
                     return;
                 }

src/scene/grid_stencil.rs

@@ -1,4 +1,6 @@
 use crate::types::*;
+use rayon::iter::{IndexedParallelIterator, ParallelIterator};
+use rayon::slice::ParallelSliceMut;
 
 pub struct PosStencilMut<'a, T>
 where
@@ -11,50 +13,16 @@ where
     pub neighbors: [&'a mut T; 2],
 }
 
-pub trait PosStencil<T>
-where
-    T: Send + Sync
-{
-    fn positive_stencils_mut(
-        &mut self,
-        dim: Index2,
-        min: Option<Index2>,
-        max: Option<Index2>,
-        offset: Option<Index2>,
-    ) -> Box<dyn Iterator<Item = PosStencilMut<T>> + '_>;
-}
-
-impl<T> PosStencil<T> for Vec<T>
-where
-    T: Send + Sync
-{
-    fn positive_stencils_mut(
-        &mut self,
-        dim: Index2,
-        min: Option<Index2>,
-        max: Option<Index2>,
-        offset: Option<Index2>,
-    ) -> Box<dyn Iterator<Item = PosStencilMut<T>> + '_> {
-        return Box::new(positive_stencils_mut(
-            self.as_mut_slice(),
-            dim,
-            min,
-            max,
-            offset,
-        ));
-    }
-}
-
 /// First dimension is stored first (column-major).
 pub fn positive_stencils_mut<T>(
     data: &mut [T],
     dim: Index2,
     min: Option<Index2>,
     max: Option<Index2>,
     offset: Option<Index2>, // Stencil offset added to min/max.
-) -> impl Iterator<Item = PosStencilMut<T>>
+) -> impl ParallelIterator<Item = PosStencilMut<T>>
 where
-    T: Send + Sync
+    T: Send + Sync,
 {
     assert!(
         dim > idx!(0, 0) && dim.iter().fold(1, std::ops::Mul::mul) == data.len(),
@@ -74,11 +42,11 @@ where
     let start_y = 0 + min.y * dim.x;
 
     let it = data[start_y..]
-        .chunks_exact_mut(2 * dim.x)
+        .par_chunks_exact_mut(2 * dim.x)
         .flat_map(move |row| {
             let (top, bot) = row.split_at_mut(dim.x);
-            let y0 = top[min.x..].chunks_exact_mut(2);
-            let y1 = bot[min.x..].chunks_exact_mut(2);
+            let y0 = top[min.x..].par_chunks_exact_mut(2);
+            let y1 = bot[min.x..].par_chunks_exact_mut(2);
 
             y0.zip(y1).map(|ys| match ys {
                 ([ref mut x0_y0, ref mut x1_y0], [ref mut x0_y1, _]) => PosStencilMut {
@@ -142,8 +110,6 @@ fn test_parallel() {
     assert!(v[(2, 0)] == 3);
     assert!(v[(2, 1)] == 6);
 
-    //positive_stencils_mut(v.as_mut_slice(), idx!(3, 2), None, None, None).par_iter();
-
 }
 
 #[test]