1 files changed, 14 insertions, 14 deletions

diff --git a/src/book.cpp b/src/book.cpp
index 2d499ca..900f94a 100644
--- a/src/book.cpp
+++ b/src/book.cpp
@@ -76,21 +76,21 @@ void InitBook(const std::string &path) {
   std::cerr << "info string " << path << ": " << book.size() << " entries\n";
 }
 
-static uint16_t PolyglotToMono(uint16_t polyglotMove) {
+static uint16_t PolyglotToMono(uint16_t polyglotMove, const Game &g) {
   uint8_t to = polyglotMove & 63;
   uint8_t from = (polyglotMove >> 6) & 63;
   uint8_t promo = (polyglotMove >> 12) & 7;
 
-  // Polyglot encodes castling as king moving to the rook's square.
-  // Convert to king-moves-two-squares form that the engine expects.
-  if (from == 4 && to == 7) {
-    to = 6; // white kingside: e1h1 -> e1g1
-  } else if (from == 4 && to == 0) {
-    to = 2; // white queenside: e1a1 -> e1c1
-  } else if (from == 60 && to == 63) {
-    to = 62; // black kingside: e8h8 -> e8g8
-  } else if (from == 60 && to == 56) {
-    to = 58; // black queenside: e8a8 -> e8c8
+  if (g.pieces[from].type == KING) {
+    if (from == 4 && to == 7) {
+      to = 6; // white kingside: e1h1 -> e1g1
+    } else if (from == 4 && to == 0) {
+      to = 2; // white queenside: e1a1 -> e1c1
+    } else if (from == 60 && to == 63) {
+      to = 62; // black kingside: e8h8 -> e8g8
+    } else if (from == 60 && to == 56) {
+      to = 58; // black queenside: e8a8 -> e8c8
+    }
   }
 
   PieceType promotion = NONEPIECE;
@@ -114,7 +114,7 @@ static uint16_t PolyglotToMono(uint16_t polyglotMove) {
   return CreateMove(from, to, promotion);
 }
 
-std::optional<uint16_t> ProbeBook(uint64_t key) {
+std::optional<uint16_t> ProbeBook(uint64_t key, const Game &g) {
   if (book.empty()) {
     return {};
   }
@@ -147,9 +147,9 @@ std::optional<uint16_t> ProbeBook(uint64_t key) {
   for (const auto *e : matches) {
     r -= e->weight;
     if (r < 0) {
-      return PolyglotToMono(e->move);
+      return PolyglotToMono(e->move, g);
     }
   }
 
-  return PolyglotToMono(matches.back()->move);
+  return PolyglotToMono(matches.back()->move, g);
 }