From 598c38efd6082489250be7788f2e02da36b07804 Mon Sep 17 00:00:00 2001
From: Simon Gardling <titaniumtown@proton.me>
Date: Sun, 9 Feb 2025 22:29:42 -0500
Subject: [PATCH] rewrite again using an Arena

---
 src/complexagent.rs | 314 +++++++++++++++++++++++++++++---------------
 1 file changed, 208 insertions(+), 106 deletions(-)

diff --git a/src/complexagent.rs b/src/complexagent.rs
index 227cb46..33cc296 100644
--- a/src/complexagent.rs
+++ b/src/complexagent.rs
@@ -3,21 +3,31 @@ use rayon::prelude::*;
 
 #[derive(Clone)]
 struct Move {
-    /// `i` value of move
+    /// `i` position of move
     i: usize,
-    /// `j` value of move
+
+    /// `j` position of move
     j: usize,
-    /// how many pieces were captured
+
+    /// how many pieces are captured by this move
     captured: usize,
-    /// Turn
+
+    /// Turn this move was made on
     color: Piece,
-    /// [`Board`] after move is made
+
+    /// [`Board`] state after move is made
     board: Board,
-    /// Winner of the match
+
+    /// Current winner of the match
     winner: Option<Piece>,
-    /// Move's parent index in [`FutureMoves`]
-    parent_index: Option<usize>,
-    /// Value determined in [`FutureMoves::compute_values`]
+
+    /// Index of this move's parent
+    parent: Option<usize>,
+
+    /// Indices of this Move's Children
+    children: Vec<usize>,
+
+    /// Value of this move
     value: i64,
 }
 
@@ -42,104 +52,226 @@ impl Move {
 }
 
 struct FutureMoves {
-    /// Contains a Vector of Vectors, each top-level vector of index `i`
-    /// represents moves at depth `i` the contents of that vector are the
-    /// possible moves
-    inner: Vec<Vec<Move>>,
+    arena: Vec<Move>,
+    current_root: Option<usize>,
+    current_depth: usize,
+    max_depth: usize,
     /// Color w.r.t
-    color: Piece,
+    agent_color: Piece,
+    gc: usize,
 }
 
 impl FutureMoves {
-    pub fn generate(color: Piece, board: &Board, depth: usize) -> Self {
-        let initial_layer: Vec<Move> = prob_space_idx(board, color, None).collect();
+    pub const fn new(agent_color: Piece, max_depth: usize) -> Self {
+        Self {
+            arena: Vec::new(),
+            current_root: None,
+            current_depth: 0,
+            max_depth,
+            agent_color,
+            gc: 0,
+        }
+    }
 
-        // 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;
-            }
-            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();
+    pub fn generate(&mut self, board: &Board, color: Piece) {
+        self.arena.clear();
+        self.current_depth = 0;
 
-        let mut tmp = Self {
-            inner: layers,
-            color,
-        };
-        tmp.compute_values();
-        tmp
+        let root_nodes = self.generate_children(None, board, color);
+        self.current_root = None;
+        self.extend_layers(root_nodes, color, self.max_depth);
+    }
+
+    fn extend_layers(&mut self, nodes: Vec<usize>, color: Piece, remaining_depth: usize) {
+        if remaining_depth == 0 {
+            return;
+        }
+
+        let next_color = !color;
+        let mut next_nodes = Vec::new();
+
+        for node_idx in nodes {
+            let board = &self.arena[node_idx].board.clone();
+            let children = self.generate_children(Some(node_idx), board, next_color);
+            next_nodes.extend(children);
+        }
+
+        self.current_depth += 1;
+        self.extend_layers(next_nodes, next_color, remaining_depth - 1);
+    }
+
+    fn generate_children(
+        &mut self,
+        parent: Option<usize>,
+        board: &Board,
+        color: Piece,
+    ) -> Vec<usize> {
+        let children: Vec<_> = board
+            .possible_moves(color)
+            .flat_map(|(i, j)| board.what_if(i, j, color).map(|x| (i, j, x)))
+            .map(|(i, j, (new_board, captured))| {
+                let winner = new_board.game_winner(color);
+                Move {
+                    i,
+                    j,
+                    captured,
+                    color,
+                    board: new_board,
+                    winner,
+                    parent,
+                    children: Vec::new(),
+                    value: 0,
+                }
+            })
+            .collect();
+
+        let parent_idx = parent;
+        let start_idx = self.arena.len();
+        self.arena.extend(children);
+        let new_indices: Vec<usize> = (start_idx..self.arena.len()).collect();
+
+        if let Some(parent_idx) = parent_idx {
+            self.arena[parent_idx].children.extend(new_indices.clone());
+        }
+
+        new_indices
     }
 
     fn compute_values(&mut self) {
-        // could be overhauled via this: https://github.com/rust-lang/rust/issues/75027
-        (0..self.inner.len()).rev().for_each(|depth| {
-            let (parents, children) = self.inner.split_at_mut(depth + 1);
+        for depth in (0..=self.current_depth).rev() {
+            let nodes_at_depth: Vec<usize> = self
+                .arena
+                .iter()
+                .enumerate()
+                .filter(|(_, node)| self.depth_of(node.parent) == depth)
+                .map(|(idx, _)| idx)
+                .collect();
 
-            // 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() };
+            for idx in nodes_at_depth {
+                let self_value = {
+                    let node = &self.arena[idx];
+                    node.compute_self_value(self.agent_color, self.current_depth - depth + 1)
+                };
 
-            let children = children.first().map(Vec::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 = children
+                let children_value = self.arena[idx]
+                    .children
                     .iter()
-                    .filter(|child| {
-                        child.parent_index.is_some() && child.parent_index == mv.parent_index
-                    })
-                    .map(|child| child.value)
+                    .map(|&child| self.arena[child].value)
                     .sum::<i64>()
-                    .checked_div(children.len() as i64)
+                    .checked_div(self.arena[idx].children.len() as i64)
                     .unwrap_or(0);
-                mv.value = self_value + children_value;
-            });
-        });
+
+                self.arena[idx].value = self_value + children_value;
+            }
+        }
     }
 
-    pub fn inner(&self) -> &[Vec<Move>] {
-        &self.inner
+    fn depth_of(&self, node_parent: Option<usize>) -> usize {
+        let mut depth = 0;
+        let mut current = node_parent;
+        while let Some(parent_idx) = current {
+            depth += 1;
+            current = self.arena[parent_idx].parent;
+        }
+        depth
+    }
+
+    pub fn best_move(&self) -> Option<(usize, usize)> {
+        let root_nodes = match self.current_root {
+            Some(root) => vec![root],
+            None => self
+                .arena
+                .iter()
+                .enumerate()
+                .filter(|(_, node)| node.parent.is_none())
+                .map(|(idx, _)| idx)
+                .collect(),
+        };
+
+        root_nodes
+            .into_iter()
+            .max_by_key(|&idx| self.arena[idx].value)
+            .map(|idx| (self.arena[idx].i, self.arena[idx].j))
+    }
+
+    pub fn update_root(&mut self, i: usize, j: usize) -> bool {
+        let new_root = self
+            .arena
+            .iter()
+            .enumerate()
+            .find(|(_, node)| node.parent == self.current_root && node.i == i && node.j == j)
+            .map(|(idx, _)| idx);
+
+        if let Some(root) = new_root {
+            self.current_root = Some(root);
+            self.current_depth = self.max_depth - self.depth_of(Some(root));
+            self.gc += 1;
+            if self.gc > 3 {
+                self.prune_unrelated();
+            }
+            true
+        } else {
+            false
+        }
+    }
+
+    fn prune_unrelated(&mut self) {
+        let Some(root) = self.current_root else {
+            return;
+        };
+
+        let mut retain = vec![false; self.arena.len()];
+        let mut stack = vec![root];
+        // traverse children of the current root
+        while let Some(idx) = stack.pop() {
+            retain[idx] = true;
+            stack.extend(self.arena[idx].children.iter().copied());
+        }
+
+        let mut new_arena = Vec::new();
+        let mut index_map = vec![None; self.arena.len()];
+
+        for (old_idx, _) in retain.iter().enumerate().filter(|(_, a)| **a) {
+            index_map[old_idx] = Some(new_arena.len());
+            let mut node = self.arena[old_idx].clone();
+            node.parent = node.parent.and_then(|p| index_map[p]);
+            node.children = node.children.iter().filter_map(|&c| index_map[c]).collect();
+            new_arena.push(node);
+        }
+
+        self.arena = new_arena;
+        self.current_root = index_map[root];
     }
 }
 
 pub struct ComplexAgent {
     color: Piece,
+    future_moves: FutureMoves,
 }
 
 impl ComplexAgent {
-    #[allow(dead_code)]
-    pub const fn new(color: Piece) -> Self {
-        Self { color }
+    pub fn new(color: Piece) -> Self {
+        const MAX_DEPTH: usize = 5;
+        Self {
+            color,
+            future_moves: FutureMoves::new(color, MAX_DEPTH),
+        }
     }
 }
 
 impl Agent for ComplexAgent {
     fn next_move(&mut self, board: &Board) -> Option<(usize, usize)> {
-        const LOOPS: usize = 5;
-        let layers = FutureMoves::generate(self.color, board, LOOPS);
+        self.future_moves.generate(board, self.color);
+        self.future_moves.compute_values();
 
-        println!(
-            "# of moves {} deep: {}",
-            LOOPS,
-            layers.inner().iter().map(Vec::len).sum::<usize>()
-        );
+        println!("# of moves stored: {}", self.future_moves.arena.len());
 
-        layers
-            .inner()
-            .first()?
-            .iter()
-            .max_by_key(|m| m.value)
-            .map(|m| (m.i, m.j))
+        if let Some((i, j)) = self.future_moves.best_move() {
+            self.future_moves.update_root(i, j);
+            Some((i, j))
+        } else {
+            None
+        }
     }
 
     fn name(&self) -> &'static str {
@@ -150,33 +282,3 @@ impl Agent for ComplexAgent {
         self.color
     }
 }
-
-fn problem_space(board: &Board, color: Piece) -> Box<dyn Iterator<Item = Move> + '_> {
-    Box::new(
-        board
-            .possible_moves(color)
-            .flat_map(move |(i, j)| board.what_if(i, j, color).map(|x| (i, j, x)))
-            .map(move |(i, j, (board, captured))| Move {
-                i,
-                j,
-                captured,
-                color,
-                board,
-                winner: board.game_winner(color),
-                parent_index: None,
-                value: 0,
-            }),
-    )
-}
-
-fn prob_space_idx(
-    board: &Board,
-    color: Piece,
-    parent_index: Option<usize>,
-) -> Box<dyn Iterator<Item = Move> + '_> {
-    Box::new(problem_space(board, color).map(move |mut m| {
-        m.parent_index = parent_index;
-        m.value = 0;
-        m
-    }))
-}