fix: correctly terminate search threads

1 files changed, 37 insertions, 23 deletions

diff --git a/src/grossmeister/UCI.rs b/src/grossmeister/UCI.rs
index 7d23e2d..034b73b 100644
--- a/src/grossmeister/UCI.rs
+++ b/src/grossmeister/UCI.rs
@@ -1,6 +1,6 @@
-use std::{io::stdin, thread::{Builder, JoinHandle, spawn, sleep}, str::Split, time::Duration};
+use std::{io::stdin, thread::{JoinHandle, spawn, sleep}, str::Split, time::Duration, sync::{Arc, atomic::AtomicBool}};
 
-use crate::{board::{Board, io::IO, color::Color}, moves::Move};
+use crate::{board::{Board, io::IO, color::Color}, moves::{Move, MoveKind}};
 
 use super::Grossmeister;
 
@@ -56,11 +56,20 @@ impl Grossmeister {
                                                 let moves = self.board.generate_pseudolegal_moves();
                                                 if let Some(mov) = moves
                                                     .iter()
-                                                    .find(|m| m.source == input_move.source && m.target == input_move.target) {
+                                                    .find(|m| {
+                                                        let promo_matches = match input_move.kind {
+                                                            MoveKind::Promotion(piece) => match m.kind {
+                                                                MoveKind::Promotion(another_piece) => piece == another_piece,
+                                                                _ => false
+                                                            },
+                                                            _ => true,
+                                                        };
+                                                        promo_matches && m.source == input_move.source && m.target == input_move.target
+                                                    }) {
                                                         self.board.make_move(*mov);
-                                                } else {
-                                                    panic!("Illegal move: {}", input_move);
-                                                }
+                                                    } else {
+                                                        panic!("Illegal move: {}", input_move);
+                                                    }
                                             }
                                         } else {
                                             panic!("Unexpected token: {}. CMD: {}", token, cmd);
@@ -133,7 +142,7 @@ impl Grossmeister {
                 "depth" => {
                     if let Some(depth) = tokens.next() {
                         let depth: u8 = depth.parse().unwrap();
-                        return self.create_search_thread(depth);
+                        return self.create_search_thread(depth, Duration::MAX);
                     }
                 },
                 "nodes" => todo!(),
@@ -143,12 +152,11 @@ impl Grossmeister {
                         let time: u64 = time.parse().unwrap();
                         let duration = Duration::from_millis(time);
 
-                        self.create_terminator_thread(duration);
-                        return self.create_search_thread(u8::MAX)
+                        return self.create_search_thread(u8::MAX, duration)
                     }
                 },
                 "infinite" => {
-                    return self.create_search_thread(u8::MAX);
+                    return self.create_search_thread(u8::MAX, Duration::MAX);
                 },
                 _ => {}
             }
@@ -158,26 +166,32 @@ impl Grossmeister {
         let color = self.board.color();
         let duration = (self.board.clock.time[color as usize] + self.board.clock.increment[color as usize]) / 20;
 
-        self.create_terminator_thread(duration);
-        self.create_search_thread(u8::MAX)
+        self.create_search_thread(u8::MAX, duration)
     }
 
-    fn create_terminator_thread(&self, duration: Duration) -> JoinHandle<()> {
+    fn create_search_thread(&self, depth: u8, duration: Duration) -> JoinHandle<Self> {
+        let ordering = std::sync::atomic::Ordering::Relaxed;
         let should_halt = self.should_halt.clone();
-        spawn(move || {
-            println!("info string terminating search in {:?}", duration);
-            sleep(duration);
-            should_halt.store(true, std::sync::atomic::Ordering::Relaxed);
-        })
-    }
+        let should_terminate = Arc::new(AtomicBool::new(true));
+
+        if duration < Duration::MAX {
+            // Create a thread that will terminate search when the duration is expired
+            let should_terminate = should_terminate.clone();
+            spawn(move || {
+                println!("info string terminating search in {:?}", duration);
+                sleep(duration);
+                if should_terminate.load(ordering) {
+                    should_halt.store(true, ordering);
+                }
+            });
+        }
 
-    fn create_search_thread(&self, depth: u8) -> JoinHandle<Self> {
-        // Clone current self and move it
-        // into thread to analyze a position
+        // Clone current self and move it into thread to analyze a position
         let mut gm = self.clone();
         spawn(move || {
             gm.iterative_deepening(depth);
-            gm
+            should_terminate.store(false, ordering); // No need to terminate search anymore
+            gm // Return better version of Self
         })
     }
 }