refactor!: implement staged move generation

 - Skip move generation on ttable hit
 - Perform selection sort *iteratively* when pulling items
 - Fix killers probed from incorrect ply (still not ideal)

4 files changed, 277 insertions, 93 deletions

diff --git a/src/board/move_generation.rs b/src/board/move_generation.rs
index 0735309..5a3650a 100644
--- a/src/board/move_generation.rs
+++ b/src/board/move_generation.rs
@@ -1,20 +1,22 @@
+use smallvec::SmallVec;
 use crate::{moves::{Move, MoveKind}, board::{color::Color, piece::Piece, CastlingSide}, bitboard::BitboardFns, square::Square};
 
 use super::Board;
 
+pub type MoveList = SmallVec<[Move; 128]>;
 
 impl Board {
-    pub fn generate_pseudolegal_moves(&self) -> Vec<Move> {
-        let mut result = Vec::new();
+    pub fn generate_pseudolegal_moves(&self) -> MoveList {
+        let mut result = MoveList::new();
         result.append(&mut self.generate_moves_core(true));
         result.append(&mut self.generate_moves_core(false));
         result
     }
 
-    fn generate_moves_core(&self, tactical_only: bool) -> Vec<Move> {
+    pub fn generate_moves_core(&self, tactical_only: bool) -> MoveList {
         let color = self.color();
         let player_pieces = self.pieces_by_color(color);
-        let mut moves = Vec::with_capacity(256);
+        let mut moves = MoveList::new();
         let empty = self.empty();
 
         let targets = if tactical_only {
diff --git a/src/grossmeister/mod.rs b/src/grossmeister/mod.rs
index 67b20bd..03bb828 100644
--- a/src/grossmeister/mod.rs
+++ b/src/grossmeister/mod.rs
@@ -1,17 +1,20 @@
 use std::sync::{atomic::AtomicBool, Arc};
 
+use smallvec::{SmallVec, smallvec};
+
 use crate::board::Board;
-use self::ttable::{TranspositionTable, TTABLE_SIZE};
+use self::{ttable::{TranspositionTable, TTABLE_SIZE}, move_selector::MoveSelector};
 
 mod ttable;
 mod evaluation;
 mod search;
 mod UCI;
+mod move_selector;
 
 /// Grossmeister is a powerful entity that plays the game of Chess.
 /// This structure represents a player, it stores his knowledge
 /// and experience about the game.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct Grossmeister {
     /// GM's internal board representation
     /// This is usually a copy of a real board
@@ -22,6 +25,8 @@ pub struct Grossmeister {
     /// It's indexex by Zobrist hash of a position mod size
     transposition_table: TranspositionTable,
 
+    move_selectors: SmallVec<[MoveSelector; 0]>,
+
     should_halt: Arc<AtomicBool>,
     debug: bool,
 }
@@ -37,6 +42,7 @@ impl Grossmeister {
         Self {
             board,
             transposition_table: vec![None; TTABLE_SIZE as usize],
+            move_selectors: smallvec![MoveSelector::default(); 256],
             should_halt: Arc::new(AtomicBool::new(false)),
             debug: false,
         }
diff --git a/src/grossmeister/move_selector.rs b/src/grossmeister/move_selector.rs
new file mode 100644
index 0000000..3ca2f1a
--- /dev/null
+++ b/src/grossmeister/move_selector.rs
@@ -0,0 +1,239 @@
+use smallvec::SmallVec;
+
+use crate::{moves::Move, board::{Board, move_generation::MoveList}};
+
+use super::Grossmeister;
+
+pub type ScoredMove = (Move, f32);
+pub type ScoredMoveList = SmallVec<[ScoredMove; 128]>;
+
+#[derive(Debug, Clone, Default)]
+pub struct ScoredMoveIter {
+    pub moves: ScoredMoveList,
+    index: usize,
+}
+
+impl Iterator for ScoredMoveIter {
+    type Item = ScoredMove;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index < self.moves.len() {
+            // Perform a selection sort
+            let mut best_score = self.moves[self.index].1;
+            for i in (self.index + 1)..self.moves.len() {
+                if self.moves[i].1 > best_score {
+                    best_score = self.moves[i].1;
+                    self.moves.swap(i, self.index);
+                }
+            }
+
+            self.index += 1;
+            return Some(self.moves[self.index - 1])
+        }
+        None
+    }
+}
+
+impl FromIterator<ScoredMove> for ScoredMoveIter {
+    fn from_iter<T: IntoIterator<Item = ScoredMove>>(iter: T) -> Self {
+        ScoredMoveIter {
+            moves: iter.into_iter().collect(),
+            index: 0,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum MoveGenStage {
+    Hash,
+    WinningOrEqualTactical,
+    Killer,
+    Quiet,
+    LosingTactical,
+}
+
+impl Default for MoveGenStage {
+    fn default() -> Self {
+        Self::Hash
+    }
+}
+
+impl MoveGenStage {
+    pub fn next_stage(&self) -> Self {
+        match self {
+            MoveGenStage::Hash => MoveGenStage::WinningOrEqualTactical,
+            MoveGenStage::WinningOrEqualTactical => MoveGenStage::Killer,
+            MoveGenStage::Killer => MoveGenStage::Quiet,
+            MoveGenStage::Quiet => MoveGenStage::LosingTactical,
+            MoveGenStage::LosingTactical => todo!()
+        }
+    }
+}
+
+#[derive(Default, Debug, Clone)]
+pub struct MoveSelector {
+    all_moves: ScoredMoveList,
+    stage_moves: ScoredMoveIter,
+    pub killer_moves: MoveList,
+    stage: MoveGenStage,
+}
+
+impl Grossmeister {
+    /// Return move selector for the current ply
+    pub fn move_selector(&mut self) -> &mut MoveSelector {
+        // dbg!(self.board.ply, self.move_selectors.len());
+        &mut self.move_selectors[self.board.ply as usize]
+    }
+
+    pub fn cleanup_selector(&mut self) {
+        // Keep the killers!
+        let killers = self.move_selector().killer_moves.clone();
+        *self.move_selector() = MoveSelector::default();
+        self.move_selector().killer_moves = killers;
+    }
+
+    /// Register killer for ply-before
+    pub fn register_killer(&mut self, killer: Move) {
+        if self.board.ply > 1 {
+            let parent_killers = &mut self.move_selectors[(self.board.ply - 2) as usize].killer_moves;
+            match parent_killers.iter().find(|m| **m == killer) {
+                None => {
+                    // println!("Registering killer {:?} for ply {}", killer, self.board.ply - 1);
+                    parent_killers.push(killer);
+                    // We want to have max 3 killers, so if we exceed the limit remove the oldest killer
+                    if parent_killers.len() > 3 {
+                        parent_killers.remove(0);
+                    }
+                }
+                Some(..) => {},
+            }
+        }
+    }
+
+    /// Evaluate move for move ordering, prioritizing efficient captures
+    /// where low-value pieces capture high-value pieces
+    pub fn eval_move(&self, m: Move) -> f32 {
+        if m.is_tactical() {
+            let [source_eval, target_eval] = [m.source, m.target]
+                .map(|sq| self.board.piece_by_square(sq))
+                .map(|p| {
+                    match p {
+                        Some(p) => p.static_eval(),
+                        None => 0.,
+                    }
+                });
+            return 2. * target_eval - source_eval
+        }
+        0.0
+    }
+
+    pub fn score_moves(&self, movelist: MoveList) -> ScoredMoveList {
+        movelist
+            .iter()
+            .map(|&m| (m, self.eval_move(m)))
+            .collect()
+    }
+
+    pub fn next_capture(&mut self) -> Option<Move> {
+        if self.move_selector().stage_moves.moves.is_empty() {
+            let moves = self.board.generate_pseudolegal_moves();
+            self.move_selector().all_moves = self.score_moves(moves);
+            let captures = self.move_selector().all_moves.iter().filter(|(m, _)| m.is_tactical()).copied().collect();
+            self.init_stage(captures);
+        }
+
+        self.move_selector().stage_moves.next().map(|(mov, _)| mov)
+    }
+
+    fn init_stage(&mut self, moves: ScoredMoveList) {
+        self.move_selector().stage_moves = ScoredMoveIter {
+            moves,
+            index: 0,
+        }
+    }
+
+    /// TODO: next stage
+    fn next_stage(&mut self) {
+        let selector = self.move_selector();
+        selector.stage = selector.stage.next_stage();
+        selector.stage_moves = ScoredMoveIter::default();
+        // println!("NEXT STAGE {:?}", selector.stage);
+    }
+
+    pub fn next_move(&mut self) -> Option<Move> {
+        loop {
+            match self.move_selector().stage {
+                MoveGenStage::Hash => {
+                    self.next_stage();
+                    if let Some(transposition) = self.transposition() {
+                        return Some(transposition.mov);
+                    }
+                }
+                MoveGenStage::WinningOrEqualTactical => {
+                    // Init stage moves
+                    if self.move_selector().stage_moves.moves.is_empty() {
+                        // Generate ALL moves (for future re-use)
+                        let moves = self.board.generate_pseudolegal_moves();
+                        self.move_selector().all_moves = self.score_moves(moves);
+
+                        // But we only care about current stage now
+                        let new_stage  =
+                            self.move_selector().all_moves.iter()
+                            .filter(|(m, score)| m.is_tactical() && *score >= 0.0)
+                            .copied()
+                            .collect();
+
+                        self.init_stage(new_stage);
+                    }
+
+                    match self.move_selector().stage_moves.next() {
+                        Some((mov, _)) => return Some(mov),
+                        None => self.next_stage(),
+                    }
+                }
+                MoveGenStage::Killer => {
+                    if self.move_selector().stage_moves.moves.is_empty() {
+                        let new_stage = self.move_selector().killer_moves.clone()
+                            .iter()
+                            .filter(|&m| self.move_selector().all_moves.iter().any(|(m2, _)| m2 == m)) // Test if killer is in the movelist
+                            .map(|&m| (m, 0.0))
+                            .collect();
+                        self.init_stage(new_stage);
+                    }
+                    match self.move_selector().stage_moves.next() {
+                        Some((mov, _)) => return Some(mov),
+                        None => self.next_stage(),
+                    }
+                }
+                MoveGenStage::Quiet => {
+                    if self.move_selector().stage_moves.moves.is_empty() {
+                        let new_stage = self.move_selector().all_moves
+                            .iter()
+                            .filter(|(m, _)| !m.is_tactical())
+                            .copied()
+                            .collect();
+                        self.init_stage(new_stage);
+                    }
+                    match self.move_selector().stage_moves.next() {
+                        Some((mov, _)) => return Some(mov),
+                        None => self.next_stage(),
+                    }
+                }
+                MoveGenStage::LosingTactical => {
+                    if self.move_selector().stage_moves.moves.is_empty() {
+                        let new_stage  =
+                            self.move_selector().all_moves.iter()
+                            .filter(|(m, score)| m.is_tactical() && *score < 0.0)
+                            .copied()
+                            .collect();
+                        self.init_stage(new_stage);
+                    }
+                    match self.move_selector().stage_moves.next() {
+                        Some((mov, _)) => return Some(mov),
+                        None => return None
+                    }
+                }
+            }
+        }
+    }
+}
diff --git a/src/grossmeister/search.rs b/src/grossmeister/search.rs
index 281352e..68756d8 100644
--- a/src/grossmeister/search.rs
+++ b/src/grossmeister/search.rs
@@ -1,9 +1,11 @@
 use std::cmp::Ordering;
 use std::f32::INFINITY;
 
-use crate::{moves::{Move, MoveKind}, board::io::IO};
+use smallvec::SmallVec;
 
-use super::{Grossmeister, ttable::{NodeType, TTABLE_SIZE, TranspositionTableItem}};
+use crate::{moves::{Move, MoveKind}, board::{io::IO, move_generation::MoveList}};
+
+use super::{Grossmeister, ttable::{NodeType, TTABLE_SIZE, TranspositionTableItem}, move_selector::MoveSelector};
 
 const VALUE_WIN: f32 = 20_000.0;
 
@@ -17,9 +19,8 @@ pub struct PerftResult {
 }
 
 impl Grossmeister {
-    pub fn negamax_search(&mut self, mut alpha: f32, beta: f32, depth_left: u8, parent_killers: &mut Vec<Move>) -> (f32, Vec<Move>) {
+    pub fn negamax_search(&mut self, mut alpha: f32, beta: f32, depth_left: u8) -> (f32, Vec<Move>) {
         let mut principal_variation = Vec::new();
-        let mut killer_moves = Vec::new();
         let color = self.board.color();
 
         if self.board.positions.iter().filter(|p| **p == self.board.hash).count() >= 3 {
@@ -54,12 +55,11 @@ impl Grossmeister {
             return (self.quiscence(alpha, beta), principal_variation);
         }
 
-        let mut moves = self.board.generate_pseudolegal_moves();
-        moves = self.order_moves(moves, parent_killers.to_vec());
-
         let mut should_pv_search = true;
         let mut legal_move_found = false;
-        for mov in moves {
+
+        self.cleanup_selector();
+        while let Some(mov) = self.next_move() {
             let ep_target_before = self.board.ep_target;
             let castling_rights_before = self.board.castling_rights;
             let hash_before = self.board.hash;
@@ -70,16 +70,16 @@ impl Grossmeister {
 
                 let (mut score, mut subtree_pv) = if should_pv_search {
                     // Assume PV-node is high in the list (if move ordering is good)
-                    self.negamax_search(-beta, -alpha, depth_left - 1, &mut killer_moves)
+                    self.negamax_search(-beta, -alpha, depth_left - 1)
                 } else {
                     // After we have PV-node (that raised alpha) all other nodes will be searched
                     // with zero-window first to confirm this assumption
                     // TODO: changing 0.001 -> 0.0001 leads to a weird bug
-                    let result = self.negamax_search(-(alpha + 0.001), -alpha, depth_left - 1, &mut killer_moves);
+                    let result = self.negamax_search(-(alpha + 0.001), -alpha, depth_left - 1);
                     // In case some of the other nodes raises alpha, then it's true PV node now,
                     // let's research with full window to find its exact value
                     if -result.0 > alpha {
-                        self.negamax_search(-beta, -alpha, depth_left - 1, &mut killer_moves)
+                        self.negamax_search(-beta, -alpha, depth_left - 1)
                     } else {
                         result
                     }
@@ -97,10 +97,7 @@ impl Grossmeister {
                     });
 
                     if mov.kind == MoveKind::Quiet {
-                        match parent_killers.iter().find(|m| **m == mov) {
-                            None => parent_killers.push(mov),
-                            Some(..) => {},
-                        }
+                        self.register_killer(mov);
                     }
 
                     return (beta, principal_variation);
@@ -151,14 +148,13 @@ impl Grossmeister {
 
     pub fn quiscence(&mut self, mut alpha: f32, beta: f32) -> f32 {
         let color = self.board.color();
-        let mut moves = self.board.generate_pseudolegal_moves();
-        moves = self.order_moves(moves, Vec::new());
 
         if self.board.positions.iter().filter(|p| **p == self.board.hash).count() >= 3 {
             // Draw by repetition
             return 0.0;
         }
 
+        let mut tactical_only = false;
         if !self.board.is_king_in_check(color) {
             // If we are not in check, we can evaluate stand pat
             let stand_pat = self.evaluate();
@@ -170,12 +166,15 @@ impl Grossmeister {
                 alpha = stand_pat;
             }
 
+            tactical_only = true;
+
             // If we are not in check, we can only search tactical moves
-            moves.retain(|m| m.is_tactical())
+            // moves.retain(|m| m.is_tactical())
         }
 
         let mut legal_move_found = false;
-        for mov in moves {
+        self.cleanup_selector();
+        while let Some(mov) = if tactical_only { self.next_capture() } else { self.next_move() } {
             let ep_target_before = self.board.ep_target;
             let castling_rights_before = self.board.castling_rights;
             let hash_before = self.board.hash;
@@ -211,7 +210,6 @@ impl Grossmeister {
         let mut beta = INFINITY;
         let window_size = 0.25;
         let mut gradual_widening_counter = 0;
-        let mut root_killers: Vec<Move> = Vec::new();
 
         while depth <= max_depth {
             println!("info hashfull {}", 1000 * self.transposition_table.iter().filter(|item| item.is_some()).count() / self.transposition_table.len());
@@ -220,7 +218,7 @@ impl Grossmeister {
                 println!("info string window {:?}", (alpha, beta));
             }
 
-            let search_result = self.negamax_search(alpha, beta, depth, &mut root_killers);
+            let search_result = self.negamax_search(alpha, beta, depth);
 
             if search_result.0.abs() >= VALUE_WIN {
                 println!("info mate {:.0}", (search_result.1.len() as f32 / 2.0).ceil(), );
@@ -270,7 +268,7 @@ impl Grossmeister {
 
         if result.is_none() {
             println!("info string could not find move in time, will re-run depth-1 search to avoid panic");
-            result = Some(self.negamax_search(-INFINITY, INFINITY, 1, &mut root_killers))
+            result = Some(self.negamax_search(-INFINITY, INFINITY, 1))
         }
 
         match result {
@@ -286,65 +284,6 @@ impl Grossmeister {
         }
     }
 
-    /// Evaluate move for move ordering, prioritizing efficient captures
-    /// where low-value pieces capture high-value pieces
-    fn eval_move(&self, m: Move) -> f32 {
-        if m.is_tactical() {
-            let [source_eval, target_eval] = [m.source, m.target]
-                .map(|sq| self.board.piece_by_square(sq))
-                .map(|p| {
-                    match p {
-                        Some(p) => p.static_eval(),
-                        None => 0.,
-                    }
-                });
-            return 2. * target_eval - source_eval
-        }
-        0.0
-    }
-
-    pub fn order_moves(&self, moves: Vec<Move>, killer_moves: Vec<Move>) -> Vec<Move> {
-        let mut moves_with_eval: Vec<(Move, f32)> = moves
-            .iter()
-            .map(|m| (*m, self.eval_move(*m)))
-            .collect();
-
-        moves_with_eval.sort_unstable_by(|(a, a_eval), (b, b_eval)| {
-            if *a_eval == 0.0 && *b_eval == 0.0 {
-                // Prioritize equal captures over non-captures
-                if a.is_tactical() && !b.is_tactical() {
-                    return Ordering::Less
-                }
-                if b.is_tactical() && !a.is_tactical() {
-                    return Ordering::Greater
-                }
-            }
-            a_eval.total_cmp(b_eval).reverse()
-        });
-
-        let mut ordered_moves: Vec<Move> = moves_with_eval.iter().map(|(m, _)| *m).collect();
-
-        // Insert killer moves after winning captures
-        let equal_capture_index = moves_with_eval
-            .iter()
-            .position(|(m, eval)| m.is_tactical() && *eval == 0.0)
-            .unwrap_or(0);
-
-        for killer in killer_moves {
-            if let Some(index) = ordered_moves.iter().position(|m| *m == killer) {
-                    let mov = ordered_moves.remove(index);
-                    ordered_moves.insert(equal_capture_index, mov);
-            }
-        }
-
-        if let Some(transposition) = self.transposition() {
-            ordered_moves.insert(0, transposition.mov);
-        }
-
-        ordered_moves
-    }
-
-
     pub fn perft(&mut self, depth: u8, print: bool) -> PerftResult {
         let mut result = PerftResult::default();
 
@@ -354,14 +293,13 @@ impl Grossmeister {
         }
         let color = self.board.color();
 
-        let moves = self.board.generate_pseudolegal_moves();
-
         if print {
-            println!("Running perft for depth {}. Color to move is {:?}\n{} moves available", depth, color, moves.len());
-            println!("{} moves available", moves.len());
+            // println!("Running perft for depth {}. Color to move is {:?}\n{} moves available", depth, color, moves.len());
+            // println!("{} moves available", moves.len());
         }
 
-        for mov in moves {
+        self.cleanup_selector();
+        while let Some(mov) = self.next_move() {
             let ep_target_before = self.board.ep_target;
             let castling_rights_before = self.board.castling_rights;
             let hash_before = self.board.hash;
@@ -414,7 +352,6 @@ impl Grossmeister {
 
 #[cfg(test)]
 mod tests {
-    use std::time::Duration;
     use crate::{board::{Board, io::IO}, square::Square, moves::{Move, MoveKind}, grossmeister::{Grossmeister, search::PerftResult}};
     use super::VALUE_WIN;