simplify code

src/complexagent.rs

@@ -62,7 +62,7 @@ struct FutureMoves {
     current_root: Option<usize>,
 
     /// Current generated depth of the Arena
-    current_depth: isize,
+    current_depth: usize,
 
     /// Target depth of children to generate
     max_depth: usize,
@@ -89,7 +89,7 @@ impl FutureMoves {
             .filter(|&idx| self.is_connected_to_root(idx)) // put here so this will not extend needlessly before prunes
             .collect();
 
-        for _ in self.current_depth..=(self.max_depth as isize) {
+        for _ in self.current_depth..=self.max_depth {
             // TODO! use `i` in order to prune along-the-way
             // i.e, every 4 moves of color `self.agent_color`, do a pruning step,
             // only keeping the top [`Move`]
@@ -100,16 +100,11 @@ impl FutureMoves {
                     ProgressStyle::with_template("Generating children: ({pos}/{len}) {per_sec}")
                         .unwrap(),
                 )
-                .flat_map(|node_idx| {
+                .flat_map(|node_idx| self.generate_children(node_idx))
-                    self.generate_children(
+                .flatten()
-                        node_idx,
-                        &self.arena[node_idx].board.clone(),
-                        !self.arena[node_idx].color,
-                    )
-                })
                 .collect();
         }
-        self.current_depth = self.max_depth as isize;
+        self.current_depth = self.max_depth;
     }
 
     /// Determines if a [`Move`] at index `idx` is connected to `self.current_root`
@@ -124,26 +119,27 @@ impl FutureMoves {
                 current = self.arena[parent_idx].parent;
             }
         }
+
         false
     }
 
     /// Creates children for a parent (`parent`), returns an iterator it's children's indexes
-    fn generate_children(
+    fn generate_children(&mut self, parent_idx: usize) -> Option<impl Iterator<Item = usize>> {
-        &mut self,
+        // early-exit if a winner for the parent already exists
-        parent_idx: usize,
+        if self.arena[parent_idx].winner != Winner::None {
-        board: &Board,
+            return None;
-        color: Piece,
+        }
-    ) -> impl Iterator<Item = usize> {
+
         let mut new: Vec<Move> =
             // use [`Board::all_positions`] here instead of [`Board::possible_moves`]
             // because we use [`Board::what_if`] later and we want to reduce calls to [`Board::propegate_from_dry`]
             Board::all_positions()
-                .flat_map(|(i, j)| board.what_if(i, j, !self.arena[parent_idx].color).map(|x| (i, j, x)))
+                .flat_map(|(i, j)| self.arena[parent_idx].board.what_if(i, j, !self.arena[parent_idx].color).map(|x| (i, j, x)))
                 .map(|(i, j, new_board)| Move {
                     i,
                     j,
                     board: new_board,
-                    winner: new_board.game_winner(color),
+                    winner: new_board.game_winner(!self.arena[parent_idx].color),
                     parent: Some(parent_idx),
                     children: Vec::new(),
                     value: 0,
@@ -154,7 +150,7 @@ impl FutureMoves {
         const TOP_K_CHILDREN: usize = 1;
 
         // we want to keep only the best move of the agent
-        if color == self.agent_color && new.len() > TOP_K_CHILDREN {
+        if !self.arena[parent_idx].color == self.agent_color && new.len() > TOP_K_CHILDREN {
             // TODO! Move this to `extend_layers` so we can prune based on recursive [`Move`] value
             // negative, because we want the max value to be at the first index
             new.sort_by_key(|x| -x.compute_self_value(self.agent_color));
@@ -167,7 +163,7 @@ impl FutureMoves {
 
         self.arena[parent_idx].children.extend(new_indices.clone());
 
-        new_indices
+        Some(new_indices)
     }
 
     /// Given an index from `self.arena`, what depth is it at? 1-indexed (ROOT IS AT INDEX 1)
@@ -182,24 +178,21 @@ impl FutureMoves {
         depth
     }
 
-    fn compute_values(&mut self) {
+    fn compute_values(&mut self, indexes: impl Iterator<Item = usize>) {
         // PERF! pre-organize all indexes based on what depth they're at
         // previously, I did a lookup map based on if a node was visited, still resulted in a full
         // O(n) iteration each depth
         let mut by_depth: Vec<Vec<usize>> = (0..=self.max_depth + 2).map(|_| Vec::new()).collect();
-        for idx in 0..self.arena.len() {
+        for idx in indexes {
             let depth = self.depth_of(idx);
-            // -1 because `depth_of` is zero-indexed
+            // -1 because `depth_of` is one-indexed
             by_depth[depth - 1].push(idx);
         }
 
         for (depth, nodes) in by_depth.into_iter().enumerate().rev() {
-            // because enum is 0-index, inc depth
-            let depth = depth + 1;
-
             for idx in nodes {
                 let self_value =
-                    self.arena[idx].compute_self_value(self.agent_color) / (depth as isize) as i64;
+                    self.arena[idx].compute_self_value(self.agent_color) / (depth + 1) as i64;
 
                 let children_value = self.arena[idx]
                     .children
@@ -278,10 +271,10 @@ impl FutureMoves {
 
     fn update_root_idx(&mut self, idx: usize) {
         self.current_root = Some(idx);
-        self.current_depth -= self.depth_of(idx) as isize - 1;
+        self.current_depth -= self.depth_of(idx) - 1;
         self.prune_unrelated();
         self.extend_layers();
-        self.compute_values();
+        self.compute_values(0..self.arena.len());
     }
 
     fn prune_unrelated(&mut self) {
@@ -349,7 +342,7 @@ pub struct ComplexAgent {
 
 impl ComplexAgent {
     pub const fn new(color: Piece) -> Self {
-        const MAX_DEPTH: usize = 17;
+        const MAX_DEPTH: usize = 15;
         Self {
             color,
             future_moves: FutureMoves::new(color, MAX_DEPTH),