1 files changed, 54 insertions, 24 deletions

diff --git a/src/sim/sim.rs b/src/sim/sim.rs
index 5cd065e..6e0029e 100644
--- a/src/sim/sim.rs
+++ b/src/sim/sim.rs
@@ -1,4 +1,6 @@
-use std::marker::PhantomData;
+use std::{collections::HashMap, marker::PhantomData};
+
+use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 
 use crate::{
     config::CHUNK_SIZE,
@@ -25,6 +27,8 @@ impl<'a> ChunkAccess<'a> {
     }
 }
 
+unsafe impl Sync for ChunkAccess<'_> {}
+
 pub struct UpdateCtx<'a, 'b, 'c> {
     pub chunks: &'a mut [Option<&'b mut Chunk>; 9],
     pub seqno: u64,
@@ -87,6 +91,12 @@ impl UpdateCtx<'_, '_, '_> {
         let x = self.x + dx;
         let y = self.y + dy;
         set_cell(self.chunks, x, y, cell);
+        // wake all the chunks
+        self.chunks.iter_mut().for_each(|c| {
+            if let Some(chunk) = c {
+                chunk.sleeping = false;
+            }
+        });
     }
 
     pub fn candidates_swap(&mut self, candidates: &[(i32, i32)]) -> bool {
@@ -155,37 +165,57 @@ const NEIGHBORHOOD_OFFSETS: [(i32, i32); 9] = [
     (1, 1),
 ];
 
-pub fn sim_tick(world: &mut World, seqno: u64) {
+pub fn sim_tick(world: &mut World, seqno: u64, use_threading: bool) {
     puffin::profile_function!();
-    let mut update_groups: [Vec<(i32, i32)>; 9] = Default::default();
 
-    // assign a color to each chunk s.t. every chunk is surrounded by <= 8 chunks of different colors
+    let mut columns: HashMap<i32, Vec<i32>> = HashMap::new();
+    for &(cx, cy) in world.chunk_position_to_chunk_idx.keys() {
+        columns.entry(cx).or_default().push(cy);
+    }
+
+    // color columns s.t. columns of same color are separated by two columns
+    // and sort the column bottom-to-top
     // --------------------
     // | 0, 1, 2, 0, 1, 2 |
-    // | 3, 4, 5, 3, 4, 5 |
-    // | 6, 7, 8, 6, 7, 8 |
     // | 0, 1, 2, 0, 1, 2 |
-    // | 3, 4, 5, 3, 4, 5 |
-    // | 6, 7, 8, 6, 7, 8 |
+    // | 0, 1, 2, 0, 1, 2 |
     // --------------------
-
-    for (&(cx, cy), _) in &world.chunk_position_to_chunk_idx {
-        let color = (cx.rem_euclid(3) * 3 + cy.rem_euclid(3)) as usize;
-        update_groups[color].push((cx, cy));
+    let mut columns_by_color: [Vec<(i32, Vec<i32>)>; 3] = Default::default();
+    for (cx, mut cys) in columns {
+        cys.sort_unstable_by(|a, b| b.cmp(a));
+        columns_by_color[cx.rem_euclid(3) as usize].push((cx, cys));
     }
 
-    for group in &update_groups {
-        let access = ChunkAccess::new(&mut world.chunks);
-        // TODO this can be parallelized since they will never share neighbours
-        for &(cx, cy) in group {
-            let mut chunks: [Option<&mut Chunk>; 9] = NEIGHBORHOOD_OFFSETS.map(|(dx, dy)| {
-                world
-                    .chunk_position_to_chunk_idx
-                    .get(&(cx + dx, cy + dy))
-                    .map(|&idx| unsafe { access.get(idx) })
-            });
+    let access = ChunkAccess::new(&mut world.chunks);
 
-            sim_tick_chunk(&mut chunks, seqno);
-        }
+    for color in &columns_by_color {
+        puffin::profile_scope!("chunk_color");
+
+        let chunk_closure = |(cx, cys): &(i32, Vec<i32>)| {
+            let cx = *cx;
+            for &cy in cys {
+                let mut chunks: [Option<&mut Chunk>; 9] = NEIGHBORHOOD_OFFSETS.map(|(dx, dy)| {
+                    world
+                        .chunk_position_to_chunk_idx
+                        .get(&(cx + dx, cy + dy))
+                        .map(|&idx| unsafe { access.get(idx) })
+                });
+
+                if let Some(target) = &mut chunks[4] {
+                    if target.sleeping {
+                        continue;
+                    }
+                    target.sleeping = true;
+                }
+
+                sim_tick_chunk(&mut chunks, seqno);
+            }
+        };
+
+        if use_threading {
+            color.par_iter().for_each(chunk_closure);
+        } else {
+            color.iter().for_each(chunk_closure);
+        };
     }
 }