feat: implement castling rights

1 files changed, 44 insertions, 6 deletions

diff --git a/src/board.rs b/src/board.rs
index 1c06c9b..4a4e2f4 100644
--- a/src/board.rs
+++ b/src/board.rs
@@ -1,5 +1,10 @@
 use crate::{bitboard::{Bitboard, serialize_bitboard, bitscan}, moves::{Move, MoveKind}, attacks::Attacks, square::Square};
 
+pub enum CastlingSide {
+    King,
+    Queen,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Board {
     pub pieces: [Bitboard; 12],
@@ -10,6 +15,12 @@ pub struct Board {
     /// En passsant target square
     pub ep_target: Option<Square>,
 
+    /// Castling rights indexed by Color and CastlingSide
+    /// ```
+    /// let can_castle = castling_rights[Color::White as usize][CastlingSide::Queen as usize];
+    /// ```
+    pub castling_rights: [[bool; 2]; 2],
+
     attacks: Attacks,
 }
 
@@ -91,6 +102,7 @@ impl Board {
             occupancy: 0,
             ply: 0,
             attacks: Attacks::new(),
+            castling_rights: [[true; 2]; 2], // TODO: actualy parse from FEN
             ep_target: None,
         };
         board.update_occupancy();
@@ -223,7 +235,7 @@ impl Board {
                                             let all_empty = castle_line.iter().all(|square| empty & square.to_bitboard() > 0);
                                             let any_checks = castle_line.iter().any(|square| self.is_square_attacked(*square, color.flip()));
 
-                                            if all_empty && !any_checks {
+                                            if all_empty && !any_checks && self.castling_rights[color as usize][CastlingSide::Queen as usize] {
                                                 moves.push(Move {
                                                     source: king_home_position,
                                                     target: king_home_position.west_one().west_one(),
@@ -240,7 +252,7 @@ impl Board {
                                             let all_empty = castle_line.iter().all(|square| empty & square.to_bitboard() > 0);
                                             let any_checks = castle_line.iter().any(|square| self.is_square_attacked(*square, color.flip()));
 
-                                            if all_empty && !any_checks {
+                                            if all_empty && !any_checks && self.castling_rights[color as usize][CastlingSide::King as usize] {
                                                 moves.push(Move {
                                                     source: king_home_position,
                                                     target: king_home_position.east_one().east_one(),
@@ -339,7 +351,7 @@ impl Board {
         }
 
         // Move a piece from source square to target
-        match self.pieces
+        let source_piece = match self.pieces
             .iter()
             .enumerate()
             .find(|(piece_type, bitboard)| *bitboard & mov.source.to_bitboard() > 0)
@@ -350,6 +362,7 @@ impl Board {
 
                     self.pieces[source_piece] |= move_target_bb;
                     self.occupancy            |= move_target_bb;
+                    PieceType::from(source_piece)
                 },
                 None => panic!("Move is malformed: source piece not found"),
             };
@@ -390,13 +403,29 @@ impl Board {
             }
         } else { None };
 
+        // Withdraw castling rights when moving rooks or king
+        match source_piece {
+            PieceType::King => {
+                self.castling_rights[Color::from_piece(source_piece) as usize][CastlingSide::King as usize] = false;
+                self.castling_rights[Color::from_piece(source_piece) as usize][CastlingSide::Queen as usize] = false;
+            },
+            PieceType::Rook => {
+                match mov.source.file() {
+                    0 => self.castling_rights[Color::from_piece(source_piece) as usize][CastlingSide::Queen as usize] = false,
+                    7 => self.castling_rights[Color::from_piece(source_piece) as usize][CastlingSide::King as usize] = false,
+                    _ => {},
+                }
+            },
+            _ => {},
+        }
+
         self.ply += 1;
 
         captured_piece
     }
 
     /// Completely reverse make_move as if it never happened
-    pub fn unmake_move(&mut self, mov: Move, captured_piece: Option<PieceType>, previous_ep_target: Option<Square>) {
+    pub fn unmake_move(&mut self, mov: Move, captured_piece: Option<PieceType>, previous_ep_target: Option<Square>, previous_castling_rights: [[bool; 2]; 2]) {
         let move_source_bb = mov.source.to_bitboard();
         let move_target_bb = mov.target.to_bitboard();
 
@@ -463,6 +492,7 @@ impl Board {
 
 
         self.ep_target = previous_ep_target;
+        self.castling_rights = previous_castling_rights;
         self.ply -= 1;
     }
 
@@ -487,6 +517,7 @@ impl Board {
 
         for mov in moves {
             let ep_target_before = self.ep_target.clone();
+            let castling_rights_before = self.castling_rights.clone();
             let captured_piece = self.make_move(mov);
             // King can not be in check after our own move
             if !self.is_king_in_check(color) {
@@ -521,7 +552,7 @@ impl Board {
                 en_passants += children_ep;
 
             }
-            self.unmake_move(mov, captured_piece, ep_target_before);
+            self.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before);
         }
 
         if print {
@@ -592,6 +623,13 @@ impl Color {
             Self::Black => Self::White,
         }
     }
+    pub fn from_piece(piece: PieceType) -> Self {
+        if (piece as u8) < 6 {
+            Self::White
+        } else {
+            Self::Black
+        }
+    }
 }
 
 #[cfg(test)]
@@ -701,7 +739,7 @@ mod tests {
         let captured_piece = board.make_move(mov);
         board.print();
 
-        board.unmake_move(mov, captured_piece, None);
+        board.unmake_move(mov, captured_piece, None, board.castling_rights);
         board.print();
 
         assert_eq!(board, initial_board, "Board state after unmake_move should be the same as before make_move");