feat: correctly display mate distance

1 files changed, 51 insertions, 14 deletions

diff --git a/src/grossmeister/search.rs b/src/grossmeister/search.rs
index b7a4208..c77eaff 100644
--- a/src/grossmeister/search.rs
+++ b/src/grossmeister/search.rs
@@ -16,7 +16,7 @@ pub struct PerftResult {
 }
 
 impl Grossmeister {
-    pub fn negamax_search(&mut self, mut alpha: f32, beta: f32, depth_left: u8) -> (f32, Vec<Move>) {
+    pub fn negamax_search(&mut self, mut alpha: f32, beta: f32, depth_left: u8, root_distance: u8) -> (f32, Vec<Move>) {
         let mut principal_variation = Vec::new();
         let color = self.board.color();
 
@@ -49,7 +49,7 @@ impl Grossmeister {
         }
 
         if depth_left == 0 {
-            return (self.quiscence(alpha, beta), principal_variation);
+            return (self.quiscence(alpha, beta, root_distance), principal_variation);
         }
 
         let mut should_pv_search = true;
@@ -67,16 +67,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)
+                    self.negamax_search(-beta, -alpha, depth_left - 1, root_distance + 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);
+                    let result = self.negamax_search(-(alpha + 0.001), -alpha, depth_left - 1, root_distance + 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)
+                        self.negamax_search(-beta, -alpha, depth_left - 1, root_distance + 1)
                     } else {
                         result
                     }
@@ -134,7 +134,7 @@ impl Grossmeister {
 
         if !legal_move_found {
             if self.board.is_king_in_check(color) {
-                return (-VALUE_WIN, principal_variation);
+                return (-VALUE_WIN + root_distance as f32, principal_variation);
             } else {
                 return (0.0, principal_variation);
             }
@@ -143,7 +143,7 @@ impl Grossmeister {
         (alpha, principal_variation)
     }
 
-    pub fn quiscence(&mut self, mut alpha: f32, beta: f32) -> f32 {
+    pub fn quiscence(&mut self, mut alpha: f32, beta: f32, root_distance: u8) -> f32 {
         let color = self.board.color();
 
         if self.board.positions.iter().filter(|p| **p == self.board.hash).count() >= 3 {
@@ -177,7 +177,7 @@ impl Grossmeister {
 
             if !self.board.is_king_in_check(color) {
                 legal_move_found = true;
-                let evaluation = -self.quiscence(-beta, -alpha);
+                let evaluation = -self.quiscence(-beta, -alpha, root_distance + 1);
                 self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);
 
                 if evaluation >= beta {
@@ -192,7 +192,7 @@ impl Grossmeister {
         }
 
         if !legal_move_found && self.board.is_king_in_check(color) {
-            return -VALUE_WIN
+            return -VALUE_WIN + root_distance as f32
         }
 
         alpha
@@ -213,12 +213,25 @@ impl Grossmeister {
                 println!("info string window {:?}", (alpha, beta));
             }
 
-            let search_result = self.negamax_search(alpha, beta, depth);
+            let search_result = self.negamax_search(alpha, beta, depth, 0);
+
+            if search_result.0.abs() >= VALUE_WIN - 128f32 { // TODO: VALUE_WIN - MAX_PLY
+                let score = search_result.0;
+                let mate_distance = (VALUE_WIN - score.abs()) * if score > 0.0 {
+                    1.
+                } else {
+                    -1. // -N means we are mated in N plies
+                };
+                println!("info mate {:.0}", (mate_distance / 2.0).ceil());
 
-            if search_result.0.abs() >= VALUE_WIN {
-                println!("info mate {:.0}", (search_result.1.len() as f32 / 2.0).ceil(), );
                 result = Some(search_result);
-                break;
+
+                if score.abs() == VALUE_WIN {
+                    break;
+                } else {
+                    depth += 1;
+                    continue;
+                }
             }
 
             if self.should_halt.load(std::sync::atomic::Ordering::SeqCst) {
@@ -263,7 +276,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))
+            result = Some(self.negamax_search(-INFINITY, INFINITY, 1, 0))
         }
 
         match result {
@@ -429,4 +442,28 @@ mod tests {
             "You should fork this bastard!"
         );
     }
+
+    #[test]
+    fn tricky_mate_them_in_2() {
+        let fen = String::from("7k/2P2Q2/3K1b2/p4N1p/P6P/6P1/8/8 w - - 10 61");
+        let board = Board::from_FEN(fen);
+
+        let mut gm = Grossmeister::new(board);
+        gm.debug = true;
+        let (score, pv) = gm.iterative_deepening(5);
+        dbg!(score, pv);
+        assert_eq!(VALUE_WIN - score, 3.0); // Mate in 3 plies
+    }
+
+    #[test]
+    fn tricky_mate_us_in_2() {
+        let fen = String::from("8/2P2Q1k/3K1b2/p4N1p/P6P/6P1/8/8 b - - 5 58");
+        let board = Board::from_FEN(fen);
+
+        let mut gm = Grossmeister::new(board);
+        gm.debug = true;
+        let (score, pv) = gm.iterative_deepening(5);
+        dbg!(score, pv);
+        assert_eq!(score + VALUE_WIN, 4.0); // Mate in 4 plies
+    }
 }