batched mcts

engine/src/mcts.rs

@@ -59,16 +59,13 @@ impl Clone for Node {
 #[derive(Clone, Debug)]
 pub struct MctsResults {
     pub board_state: BoardState,
-    pub move_dist: HashMap<ChessMove, f32>,
+    // Compact encoded move distribution: (encoded_move_index, probability)
+    pub move_dist: Vec<(usize, f32)>,
     pub value: f32,
 }
 
 impl MctsResults {
-    pub fn new(
-        board_state: BoardState,
-        move_dist: HashMap<ChessMove, f32>,
-        value: f32,
-    ) -> MctsResults {
+    pub fn new(board_state: BoardState, move_dist: Vec<(usize, f32)>, value: f32) -> MctsResults {
         MctsResults {
             board_state,
             move_dist,
@@ -230,12 +227,15 @@ impl<B: Backend> Mcts<B> {
             }
         }
 
-        let mut move_dist: HashMap<ChessMove, f32> = HashMap::new(); // TODO: make vec<(Chessmove, f32)>
-        for idx in nodes[root].children.iter() {
-            move_dist.insert(
-                nodes[*idx].last_move.expect("move didnt exist"),
-                nodes[*idx].visit_count as f32 / self.config.num_simulations as f32,
-            );
+        // Build compact move distribution: encoded move index -> probability (visits / num_simulations)
+        let mut move_dist: Vec<(usize, f32)> = Vec::with_capacity(nodes[root].children.len());
+        let stm = board_state.board.side_to_move();
+        let denom = self.config.num_simulations as f32;
+        for child_idx in nodes[root].children.iter() {
+            let mv = nodes[*child_idx].last_move.expect("move didnt exist");
+            let enc = encode_move(mv, stm);
+            let prob = nodes[*child_idx].visit_count as f32 / denom;
+            move_dist.push((enc, prob));
         }
 
         MctsResults::new(board_state.clone(), move_dist, nodes[root].value())