more iterator stuff

This commit is contained in:
2025-02-11 00:27:29 -05:00
parent 3e2ab2fd32
commit 657bb967a0
4 changed files with 80 additions and 56 deletions

View File

@@ -56,6 +56,7 @@ struct FutureMoves {
current_root: Option<usize>,
current_depth: usize,
max_depth: usize,
/// Color w.r.t
agent_color: Piece,
}
@@ -80,22 +81,30 @@ impl FutureMoves {
self.extend_layers(root_nodes, color, self.max_depth);
}
fn extend_layers(&mut self, nodes: Vec<usize>, color: Piece, remaining_depth: usize) {
fn extend_layers(
&mut self,
nodes: impl Iterator<Item = 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);
}
let next_nodes: Vec<usize> = nodes
.flat_map(|node_idx| {
self.generate_children(
Some(node_idx),
&self.arena[node_idx].board.clone(),
next_color,
)
})
.collect();
self.current_depth += 1;
self.extend_layers(next_nodes, next_color, remaining_depth - 1);
self.extend_layers(next_nodes.into_iter(), next_color, remaining_depth - 1);
}
fn generate_children(
@@ -103,31 +112,28 @@ impl FutureMoves {
parent: Option<usize>,
board: &Board,
color: Piece,
) -> Vec<usize> {
let parent_idx = parent;
) -> impl Iterator<Item = usize> {
let start_idx = self.arena.len();
self.arena.extend(
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,
}
.map(|(i, j, (new_board, captured))| Move {
i,
j,
captured,
color,
board: new_board,
winner: new_board.game_winner(color),
parent,
children: Vec::new(),
value: 0,
}),
);
let new_indices: Vec<usize> = (start_idx..self.arena.len()).collect();
if let Some(parent_idx) = parent_idx {
let new_indices = start_idx..self.arena.len();
if let Some(parent_idx) = parent {
self.arena[parent_idx].children.extend(new_indices.clone());
}
@@ -144,10 +150,8 @@ impl FutureMoves {
.collect();
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 self_value = self.arena[idx]
.compute_self_value(self.agent_color, self.current_depth - depth + 1);
let children_value = self.arena[idx]
.children
@@ -191,18 +195,18 @@ impl FutureMoves {
}
pub fn update_root(&mut self, i: usize, j: usize) -> bool {
let new_root = self.arena.iter().enumerate().find_map(|(idx, node)| {
(node.parent == self.current_root && node.i == i && node.j == j).then_some(idx)
});
if let Some(root) = new_root {
self.current_root = Some(root);
self.current_depth = self.max_depth - self.depth_of(Some(root));
self.prune_unrelated();
true
} else {
false
}
self.arena
.iter()
.enumerate()
.find_map(|(idx, node)| {
(node.parent == self.current_root && node.i == i && node.j == j).then_some(idx)
})
.inspect(|&root| {
self.current_root = Some(root);
self.current_depth = self.max_depth - self.depth_of(Some(root));
self.prune_unrelated();
})
.is_some()
}
fn prune_unrelated(&mut self) {
@@ -215,19 +219,26 @@ impl FutureMoves {
// traverse children of the current root
while let Some(idx) = stack.pop() {
retain[idx] = true;
stack.extend(self.arena[idx].children.iter().copied());
stack.extend(self.arena[idx].children.iter());
}
let mut new_arena = Vec::with_capacity(self.arena.len());
let mut index_map = vec![None; self.arena.len()];
// PERF: reverse iterator to avoid unneeded `memcpy`s when deconstructing the Arena
for (old_idx, keep) in retain.iter().enumerate().rev() {
let mut node = self.arena.remove(old_idx);
if *keep {
index_map[old_idx] = Some(new_arena.len());
node.parent = node.parent.and_then(|p| index_map[p]);
node.children = node.children.iter().filter_map(|&c| index_map[c]).collect();
node.children = node
.children
.into_iter()
.filter_map(|c| index_map[c])
.collect();
new_arena.push(node);
}
}
@@ -244,7 +255,7 @@ pub struct ComplexAgent {
impl ComplexAgent {
pub fn new(color: Piece) -> Self {
const MAX_DEPTH: usize = 7;
const MAX_DEPTH: usize = 6;
Self {
color,
future_moves: FutureMoves::new(color, MAX_DEPTH),
@@ -259,12 +270,9 @@ impl Agent for ComplexAgent {
println!("# of moves stored: {}", self.future_moves.arena.len());
if let Some((i, j)) = self.future_moves.best_move() {
self.future_moves.best_move().inspect(|&(i, j)| {
self.future_moves.update_root(i, j);
Some((i, j))
} else {
None
}
})
}
fn name(&self) -> &'static str {