rewrite progress

src/complexagent.rs

@@ -3,6 +3,7 @@ use std::num::NonZero;
 use crate::{agent::Agent, board::Board, piece::Piece};
 use either::Either;
 use indicatif::{ProgressBar, ProgressIterator, ProgressStyle};
+use num::Integer;
 
 #[derive(Clone)]
 struct Move {
@@ -26,6 +27,8 @@ struct Move {
 
     /// Value of this move
     value: i64,
+
+    color: Piece,
 }
 
 impl Move {
@@ -52,51 +55,29 @@ struct FutureMoves {
     current_depth: usize,
 
     max_depth: usize,
-    max_arena: usize,
 
     /// Color w.r.t
     agent_color: Piece,
 }
 
 impl FutureMoves {
-    pub const fn new(agent_color: Piece, max_depth: usize, max_arena: usize) -> Self {
+    pub const fn new(agent_color: Piece, max_depth: usize) -> Self {
         Self {
             arena: Vec::new(),
             current_root: None,
             current_depth: 0,
             max_depth,
-            max_arena,
             agent_color,
         }
     }
 
-    /// Create all future moves based on struct fields
-    pub fn generate(&mut self, board: &Board, color: Piece) {
-        self.arena.clear();
-        self.current_depth = 0;
-
-        let root_nodes = self.generate_children(None, board, color);
-        self.current_root = None;
-        self.extend_layers(root_nodes, color, self.max_depth);
-    }
-
     /// Generate children for all children of `nodes`
-    fn extend_layers(
-        &mut self,
-        nodes: impl Iterator<Item = usize>,
-        color: Piece,
-        remaining_depth: usize,
-    ) {
-        self.current_depth += remaining_depth;
-        let mut next_nodes: Vec<usize> = nodes.collect();
-        let mut next_color = !color;
-
-        for _ in 0..remaining_depth {
-            // early exit if arena is larger than the max size
-            if self.arena.len() >= self.max_arena {
-                break;
-            }
+    fn extend_layers(&mut self) {
+        let mut next_nodes: Vec<usize> = (0..self.arena.len())
+            .filter(|&idx| self.arena[idx].children.is_empty())
+            .collect();
 
+        for _ in (self.current_depth..self.max_depth).progress() {
             let arena_len = self.arena.len();
             next_nodes = next_nodes
                 .into_iter()
@@ -112,14 +93,14 @@ impl FutureMoves {
                     self.generate_children(
                         Some(node_idx),
                         &self.arena[node_idx].board.clone(),
-                        next_color,
+                        !self.arena[node_idx].color,
                     )
                 })
                 .collect();
-            next_color = !next_color;
         }
     }
 
+    /// Creates children for a parent (`parent`), returns an iterator it's children's indexes
     fn generate_children(
         &mut self,
         parent: Option<usize>,
@@ -140,6 +121,7 @@ impl FutureMoves {
                     parent,
                     children: Vec::new(),
                     value: new_board.count(self.agent_color) as i64 - new_board.count(!self.agent_color) as i64,
+                    color: parent.map(|idx| !self.arena[idx].color).unwrap_or(self.agent_color)
                 }).collect();
 
         // we want to keep only the best move of the agent
@@ -221,6 +203,30 @@ impl FutureMoves {
         .map(|x| (x.i, x.j))
     }
 
+    pub fn update(&mut self, board: &Board) {
+        self.current_root = Some(self.current_root.unwrap_or_else(|| {
+            self.arena
+                .iter()
+                .enumerate()
+                .find_map(|(idx, m)| (&m.board == board).then_some(idx))
+                .unwrap_or_else(|| {
+                    self.arena.push(Move {
+                        i: 0,
+                        j: 0,
+                        board: *board,
+                        winner: None,
+                        parent: None,
+                        children: Vec::new(),
+                        value: 0,
+                        color: !self.agent_color,
+                    });
+                    0
+                })
+        }));
+
+        self.extend_layers();
+    }
+
     pub fn update_root(&mut self, i: usize, j: usize) -> bool {
         self.arena
             .iter()
@@ -232,6 +238,7 @@ impl FutureMoves {
                 self.current_root = Some(root);
                 self.current_depth = self.max_depth - self.depth_of(root).get();
                 self.prune_unrelated();
+                self.extend_layers();
                 true
             })
     }
@@ -292,17 +299,16 @@ pub struct ComplexAgent {
 impl ComplexAgent {
     pub const fn new(color: Piece) -> Self {
         const MAX_DEPTH: usize = 17;
-        const MAX_ARENA: usize = 100_000_000;
         Self {
             color,
-            future_moves: FutureMoves::new(color, MAX_DEPTH, MAX_ARENA),
+            future_moves: FutureMoves::new(color, MAX_DEPTH),
         }
     }
 }
 
 impl Agent for ComplexAgent {
     fn next_move(&mut self, board: &Board) -> Option<(usize, usize)> {
-        self.future_moves.generate(board, self.color);
+        self.future_moves.update(board);
         self.future_moves.compute_values();
 
         println!("# of moves stored: {}", self.future_moves.arena.len());