count neighbors differently depending on task

src/cell.rs

@@ -9,6 +9,52 @@ pub struct Cell {
     pub state: u32,
 }
 
+#[derive(PartialEq)]
+enum StateOp {
+    Types,
+    States,
+}
+
+impl Cell {
+    fn neighbor_count(&self, grid: &Vec<Vec<Cell>>, operation: StateOp) -> (u32, u32) {
+        let mut infected: u32 = 0;
+        let mut states: u32 = 0;
+        let mut sum: u32 = 0;
+        let mut ill: u32 = 0;
+
+        for j in -1..=1 {
+            for i in -1..=1 {
+                if i == 0 && j == 0 && operation == StateOp::Types {
+                    continue;
+                }
+                let nx = (self.x as i32 + i + WIDTH as i32) % WIDTH as i32;
+                let ny = (self.y as i32 + j + HEIGHT as i32) % HEIGHT as i32;
+                let next = &grid[ny as usize][nx as usize];
+                let state = next.state;
+
+                if operation == StateOp::Types {
+                    if state > 1 && state < Q - 1 {
+                        infected += 1;
+                    } else if state == Q - 1 {
+                        ill += 1;
+                    }
+                } else {
+                    sum += state;
+                    if state == 1 && next != self {
+                        states += 1;
+                    }
+                }
+            }
+        }
+
+        if operation == StateOp::Types {
+            return (infected, ill);
+        } else {
+            return (sum, states);
+        }
+    }
+}
+
 pub fn init_grid(option: Option<&Vec<Vec<Cell>>>) -> Vec<Vec<Cell>> {
     let mut grid: Vec<Vec<Cell>> = Vec::new();
     for j in 0..HEIGHT {