iterator fun

src/complexagent.rs

@@ -33,7 +33,7 @@ impl Move {
         } else if self.winner == Some(agent_color) {
             // results in a win for the agent
             self_value = i64::MAX;
-        } else if agent_color != self.color {
+        } else if self.color != agent_color {
             self_value = -self_value;
         }
 
@@ -52,25 +52,25 @@ struct FutureMoves {
 
 impl FutureMoves {
     pub fn generate(color: Piece, board: &Board, depth: usize) -> Self {
-        let mut layers: Vec<Vec<Move>> = Vec::with_capacity(depth);
-        layers.push(prob_space_idx(board, color, None).collect());
+        let initial_layer: Vec<Move> = prob_space_idx(board, color, None).collect();
 
-        for current_depth in 0..depth {
-            let current_layer = &layers[current_depth];
-            if current_layer.is_empty() {
-                break;
+        // std::iter::successors is super cool!
+        let layers: Vec<Vec<Move>> = std::iter::successors(Some(initial_layer), |prev_layer| {
+            if prev_layer.is_empty() {
+                return None;
             }
-
-            layers.push(
-                current_layer
+            Some(
+                prev_layer
                     .iter()
                     .enumerate()
                     .flat_map(|(parent_idx, parent_move)| {
                         prob_space_idx(&parent_move.board, !parent_move.color, Some(parent_idx))
                     })
                     .collect(),
-            );
-        }
+            )
+        })
+        .take(depth + 1)
+        .collect();
 
         let mut tmp = Self {
             inner: layers,
@@ -82,24 +82,29 @@ impl FutureMoves {
 
     fn compute_values(&mut self) {
         // could be overhauled via this: https://github.com/rust-lang/rust/issues/75027
-        for depth in (0..self.inner.len()).rev() {
-            let (curr, prev) = self.inner.split_at_mut(depth + 1);
-            let curr = curr.last_mut().unwrap();
-            let prev = prev.first_mut().map(|x| x.as_slice()).unwrap_or(&[]);
+        (0..self.inner.len()).rev().for_each(|depth| {
+            let (parents, children) = self.inner.split_at_mut(depth + 1);
 
-            curr.iter_mut().for_each(|mv| {
+            // SAFETY! `parents` will always be at index `depth` which will always be within range (0..self.inner.len())
+            let parents = unsafe { parents.last_mut().unwrap_unchecked() };
+
+            let children = children.first_mut().map(|x| x.as_slice()).unwrap_or(&[]);
+
+            parents.iter_mut().for_each(|mv| {
                 let self_value = mv.compute_self_value(self.color, depth + 1);
                 // calculate average value of each move
-                let children_value = prev
+                let children_value = children
                     .iter()
-                    .filter(|child| child.parent_index == Some(mv.parent_index.unwrap_or(0)))
+                    .filter(|child| {
+                        child.parent_index.is_some() && child.parent_index == mv.parent_index
+                    })
                     .map(|child| child.value)
                     .sum::<i64>()
-                    .checked_div(prev.len() as i64)
+                    .checked_div(children.len() as i64)
                     .unwrap_or(0);
                 mv.value = self_value + children_value;
             });
-        }
+        });
     }
 
     pub fn inner(&self) -> &[Vec<Move>] {
@@ -129,7 +134,9 @@ impl Agent for ComplexAgent {
             layers.inner().iter().map(Vec::len).sum::<usize>()
         );
 
-        layers.inner()[0]
+        layers
+            .inner()
+            .first()?
             .iter()
             .max_by_key(|m| m.value)
             .map(|m| (m.i, m.j))