optimizations

src/complexagent.rs

@@ -1,7 +1,8 @@
-use either::Either;
-use indicatif::{ProgressBar, ProgressIterator, ProgressStyle};
+use std::num::NonZero;
 
 use crate::{agent::Agent, board::Board, piece::Piece};
+use either::Either;
+use indicatif::{ProgressBar, ProgressIterator, ProgressStyle};
 
 #[derive(Clone)]
 struct Move {
@@ -86,28 +87,28 @@ impl FutureMoves {
     ) {
         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;
             }
+
             let arena_len = self.arena.len();
-            next_nodes =
-                next_nodes
-                    .into_iter()
-                    .flat_map(|node_idx| {
-                        self.generate_children(
-                            Some(node_idx),
-                            &self.arena[node_idx].board.clone(),
-                            next_color,
-                        )
-                    })
-                    .progress_with(ProgressBar::new(arena_len as u64).with_style(
-                        ProgressStyle::with_template("({pos}/{len}) {per_sec}").unwrap(),
-                    ))
-                    .collect();
+            next_nodes = next_nodes
+                .into_iter()
+                .flat_map(|node_idx| {
+                    self.generate_children(
+                        Some(node_idx),
+                        &self.arena[node_idx].board.clone(),
+                        next_color,
+                    )
+                })
+                .progress_with(ProgressBar::new(arena_len as u64).with_style(
+                    ProgressStyle::with_template("Children: ({pos}/{len}) {per_sec}").unwrap(),
+                ))
+                .collect();
             next_color = !next_color;
         }
     }
@@ -159,7 +160,7 @@ impl FutureMoves {
                         true
                     }
                 })
-                .filter(|&idx| self.depth_of(idx) == depth)
+                .filter(|&idx| self.depth_of(idx).get() == depth)
                 .collect();
 
             for idx in nodes_at_depth {
@@ -179,26 +180,25 @@ impl FutureMoves {
         }
     }
 
-    /// Given an index from `self.arena`, what depth is it at? 0-indexed
-    fn depth_of(&self, node_parent: usize) -> usize {
+    /// Given an index from `self.arena`, what depth is it at? 1-indexed (ROOT IS AT INDEX 1)
+    fn depth_of(&self, node_idx: usize) -> NonZero<usize> {
         let mut depth = 0;
-        let mut current = Some(node_parent);
+        let mut current = Some(node_idx);
         while let Some(parent_idx) = current {
             depth += 1;
             current = self.arena[parent_idx].parent;
         }
-        depth
+
+        // SAFETY! because `node_idx` is of type `usize`, depth will never be 0
+        unsafe { NonZero::new_unchecked(depth) }
     }
 
     pub fn best_move(&self) -> Option<(usize, usize)> {
         (match self.current_root {
             Some(root) => Either::Left(vec![root].into_iter()),
-            None => Either::Right(
-                self.arena
-                    .iter()
-                    .enumerate()
-                    .filter_map(|(idx, node)| node.parent.is_none().then_some(idx)),
-            ),
+            None => {
+                Either::Right((0..self.arena.len()).filter(|&idx| self.depth_of(idx).get() == 1))
+            }
         })
         .map(|idx| &self.arena[idx])
         .max_by_key(|x| x.value)
@@ -214,7 +214,7 @@ impl FutureMoves {
             })
             .is_some_and(|root| {
                 self.current_root = Some(root);
-                self.current_depth = self.max_depth - self.depth_of(root);
+                self.current_depth = self.max_depth - self.depth_of(root).get();
                 self.prune_unrelated();
                 true
             })
@@ -229,7 +229,7 @@ impl FutureMoves {
 
         // stack is going to be AT MAXIMUM, the size of the array,
         // so lets just pre-allocate it
-        let mut stack = Vec::with_capacity(self.arena.len());
+        let mut stack: Vec<usize> = Vec::with_capacity(self.arena.len());
         stack.push(root);
 
         // traverse children of the current root
@@ -240,20 +240,12 @@ impl FutureMoves {
 
         let mut index_map = vec![None; self.arena.len()];
 
-        // PERF: reverse iterator to avoid unneeded `memcpy`s when deconstructing the Arena
         self.arena = retain
             .into_iter()
-            .enumerate()
-            .rev()
-            .enumerate()
-            .flat_map(|(new_idx, (old_idx, keep))| {
-                let node = self.arena.remove(old_idx);
-                if keep {
-                    Some((new_idx, (old_idx, node)))
-                } else {
-                    None
-                }
-            })
+            .enumerate() // old_idx
+            .zip(self.arena.drain(..))
+            .flat_map(|((old_idx, keep), node)| keep.then_some((old_idx, node))) // filter out unrelated nodes
+            .enumerate() // new_idx
             .map(|(new_idx, (old_idx, mut node))| {
                 index_map[old_idx] = Some(new_idx);
 
@@ -283,7 +275,7 @@ pub struct ComplexAgent {
 
 impl ComplexAgent {
     pub const fn new(color: Piece) -> Self {
-        const MAX_DEPTH: usize = 10;
+        const MAX_DEPTH: usize = 8;
         const MAX_ARENA: usize = 100_000_000;
         Self {
             color,