1 files changed, 41 insertions, 34 deletions

diff --git a/src/board/board.cpp b/src/board/board.cpp
index a76ec42..22db36e 100644
--- a/src/board/board.cpp
+++ b/src/board/board.cpp
@@ -2,6 +2,7 @@
 #include "../moves.hpp"
 #include "zobrist/zobrist.hpp"
 #include <cassert>
+#include <cstdlib>
 
 Piece createPiece(PieceType type, bool color) {
   Piece piece;
@@ -69,11 +70,11 @@ UndoMove MakeMove(Move move, Game *g) {
   if (piece.type == PAWN && g->canEnpassant && move.To == g->enPassant) {
     Position capturedPawn = move.To;
     capturedPawn.rank += piece.color ? 1 : -1;
-    g->pieces[PositionToIndex(capturedPawn)] = {false, NONEPIECE};
 
     undo.wasEnPassantCapture = true;
     undo.enPassantCapturedSquare = capturedPawn;
     undo.enPassantCapturedPiece = g->pieces[PositionToIndex(capturedPawn)];
+    g->pieces[PositionToIndex(capturedPawn)] = {false, NONEPIECE};
   }
 
   g->canEnpassant = false;
@@ -97,33 +98,38 @@ UndoMove MakeMove(Move move, Game *g) {
   }
 
   // caslte
-  if (piece.type == KING && move.From.rank == (piece.color ? 0 : 7) &&
-      (move.From.file + 2 == move.To.file ||
-       move.From.file - 2 == move.To.file)) {
+  if (piece.type == KING && move.From.rank == (piece.color ? 7 : 0) &&
+      std::abs(move.To.file - move.From.file) == 2) {
 
     undo.wasCastle = true;
-    undo.CastledSide = (move.To.file > move.From.file);
-    if (move.From.file + 2 == move.To.file) {
-      Position rook_pos = move.To;
-      rook_pos.file += 1;
-      Piece rook = g->pieces[PositionToIndex(rook_pos)];
-      g->pieces[PositionToIndex(rook_pos)] = {false, NONEPIECE};
-      rook_pos = move.From;
-      rook_pos.file -= 1;
-      g->pieces[PositionToIndex(rook_pos)] = rook;
-    }
-    if (move.From.file - 2 == move.To.file) {
-      Position rook_pos = move.To;
-      rook_pos.file -= 1;
-      Piece rook = g->pieces[PositionToIndex(rook_pos)];
-      g->pieces[PositionToIndex(rook_pos)] = {false, NONEPIECE};
-      rook_pos = move.From;
-      rook_pos.file += 1;
-      g->pieces[PositionToIndex(rook_pos)] = rook;
+    undo.CastledSide = move.To.file > move.From.file;
+
+    if (undo.CastledSide) {
+      Position rookFrom = move.From;
+      rookFrom.file = 7;
+
+      Position rookTo = move.From;
+      rookTo.file = 5;
+
+      g->pieces[PositionToIndex(rookTo)] = g->pieces[PositionToIndex(rookFrom)];
+      g->pieces[PositionToIndex(rookFrom)] = {false, NONEPIECE};
+    } else {
+      // Queenside: a -> d
+      Position rookFrom = move.From;
+      rookFrom.file = 0;
+
+      Position rookTo = move.From;
+      rookTo.file = 3;
+
+      g->pieces[PositionToIndex(rookTo)] = g->pieces[PositionToIndex(rookFrom)];
+      g->pieces[PositionToIndex(rookFrom)] = {false, NONEPIECE};
     }
-  };
+  }
 
-  // removing caslte rights
+  /*
+   * remove castling right if king has moved
+   * but only for color of the king
+   */
   if (piece.type == KING) {
     if (piece.color) {
       g->whiteCastleKing = false;
@@ -219,27 +225,28 @@ void UnMakeMove(UndoMove undo, Game *g) {
 
   if (undo.wasCastle) {
     if (undo.CastledSide) {
-      // rook f-file -> h-file
-      Position rookFrom = undo.to;
-      rookFrom.file -= 1; // f1/f8
+      // f -> h
+      Position rookFrom = undo.from;
+      rookFrom.file = 5;
 
-      Position rookTo = undo.to;
-      rookTo.file += 1; // h1/h8
+      Position rookTo = undo.from;
+      rookTo.file = 7;
 
       g->pieces[PositionToIndex(rookTo)] = g->pieces[PositionToIndex(rookFrom)];
       g->pieces[PositionToIndex(rookFrom)] = {false, NONEPIECE};
     } else {
-      // rook d-file -> a-file
-      Position rookFrom = undo.to;
-      rookFrom.file += 1; // d1/d8
+      // d -> a
+      Position rookFrom = undo.from;
+      rookFrom.file = 3;
 
-      Position rookTo = undo.to;
-      rookTo.file -= 2; // a1/a8
+      Position rookTo = undo.from;
+      rookTo.file = 0;
 
       g->pieces[PositionToIndex(rookTo)] = g->pieces[PositionToIndex(rookFrom)];
       g->pieces[PositionToIndex(rookFrom)] = {false, NONEPIECE};
     }
   }
+
   g->ThreeFoldMap[undo.zobristKey] -= 1;
   if (g->ThreeFoldMap[undo.zobristKey] <= 0) {
     g->ThreeFoldMap.erase(undo.zobristKey);