feat: implement castling

3 files changed, 158 insertions, 41 deletions

diff --git a/src/board.rs b/src/board.rs
index efe0b80..1c06c9b 100644
--- a/src/board.rs
+++ b/src/board.rs
@@ -158,10 +158,12 @@ impl Board {
 
     pub fn generate_moves(&self, color: Color) -> Vec<Move> {
         let mut moves = Vec::with_capacity(1024);
-        let opponent_occupancy = self.color_occupancy(Color::from(color as u8).flip());
+        let opponent_occupancy = self.color_occupancy(color.flip());
         let empty = self.empty();
         let available_targets = opponent_occupancy | empty;
-        for (piece_type, piece) in self.pieces_by_color(color).iter().enumerate() {
+        let player_pieces = self.pieces_by_color(color);
+
+        for (piece_type, piece) in player_pieces.iter().enumerate() {
             match PieceType::from(piece_type) {
                 PieceType::Pawn => {
                     for source in serialize_bitboard(*piece) {
@@ -200,6 +202,56 @@ impl Board {
                         for target in serialize_bitboard(self.attacks.king[source as usize] & opponent_occupancy) {
                             moves.push(Move { source, target, kind: MoveKind::Capture });
                         };
+
+                        // Castling
+                        let king_home_position = match color {
+                            Color::White => Square::E1,
+                            Color::Black => Square::E8,
+                        };
+                        if *piece == king_home_position.to_bitboard() {
+                            for rook_square in serialize_bitboard(player_pieces[PieceType::Rook as usize])
+                                .iter()
+                                .filter(|rook_square| rook_square.rank() == king_home_position.rank())
+                                {
+                                    match rook_square.file() {
+                                        0 => {
+                                            let castle_line = [
+                                                king_home_position.west_one(),
+                                                king_home_position.west_one().west_one(),
+                                            ];
+
+                                            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 {
+                                                moves.push(Move {
+                                                    source: king_home_position,
+                                                    target: king_home_position.west_one().west_one(),
+                                                    kind: MoveKind::Castle,
+                                                })
+                                            }
+                                        },
+                                        7 => {
+                                            let castle_line = [
+                                                king_home_position.east_one(),
+                                                king_home_position.east_one().east_one(),
+                                            ];
+
+                                            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 {
+                                                moves.push(Move {
+                                                    source: king_home_position,
+                                                    target: king_home_position.east_one().east_one(),
+                                                    kind: MoveKind::Castle,
+                                                })
+                                            }
+                                        },
+                                        _ => {},
+                                    }
+                                }
+                        }
                     }
                 }
                 PieceType::Knight => {
@@ -270,14 +322,13 @@ impl Board {
             };
 
         // En Passant captures diffirently
-        captured_piece = match mov.kind {
-            MoveKind::EnPassant => {
-                debug_assert!(captured_piece.is_none(), "No capture should be found at this point");
-                let captured_bb = Square::from(mov.source.rank() * 8 + mov.target.file()).to_bitboard();
-                match self.pieces
-                    .iter()
-                    .enumerate()
-                    .find(|(piece_type, bitboard)| *bitboard & captured_bb > 0)
+        if mov.kind == MoveKind::EnPassant {
+            debug_assert!(captured_piece.is_none(), "No capture should be found at this point");
+            let captured_bb = Square::from_coords(mov.source.rank(), mov.target.file()).to_bitboard();
+            captured_piece = match self.pieces
+                .iter()
+                .enumerate()
+                .find(|(piece_type, bitboard)| *bitboard & captured_bb > 0)
                 {
                     Some((pawn_type, _)) => {
                         self.pieces[pawn_type] ^= captured_bb;
@@ -285,9 +336,7 @@ impl Board {
                     }
                     None => panic!("Pawn captured by En Passant was not found"),
                 }
-            },
-            _ => captured_piece,
-        };
+        }
 
         // Move a piece from source square to target
         match self.pieces
@@ -305,20 +354,42 @@ impl Board {
                 None => panic!("Move is malformed: source piece not found"),
             };
 
-        // Double push should set En Passant target square
-        match mov.kind {
-            MoveKind::DoublePush => {
-                self.ep_target = match mov.source.rank() {
-                    1 => Some(mov.source.nort_one()),
-                    6 => Some(mov.source.sout_one()),
-                    rank => panic!("Double-push was used from invalid rank({}) when trying to make {:?}", rank, mov),
-                };
-            },
-            _ => {
-                self.ep_target = None
-            },
+        // When castling, also move a rook
+        if mov.kind == MoveKind::Castle {
+            debug_assert!(mov.source.file() == 4, "Castle can only be done from E file");
+            let (rook_source_file, rook_target_file) = match mov.target.file() {
+                2 => (0, 3),
+                6 => (7, 5),
+                _ => panic!("Malformed castle, target square invalid: {:?}", mov),
+            };
+            let rook_source_bb = Square::from_coords(mov.target.rank(), rook_source_file).to_bitboard();
+            let rook_target_bb = Square::from_coords(mov.target.rank(), rook_target_file).to_bitboard();
+
+            match self.pieces
+                .iter()
+                .enumerate()
+                .find(|(rook_type, bitboard)| *bitboard & rook_source_bb > 0)
+                {
+                    Some((rook_type, _)) => {
+                        self.pieces[rook_type] ^= rook_source_bb;
+                        self.occupancy         ^= rook_source_bb;
+
+                        self.pieces[rook_type] |= rook_target_bb;
+                        self.occupancy         |= rook_target_bb;
+                    },
+                    None => panic!("Rook was not found when castling"),
+                }
         }
 
+        // Double push should set En Passant target square
+        self.ep_target = if mov.kind == MoveKind::DoublePush {
+            match mov.source.rank() {
+                1 => Some(mov.source.nort_one()),
+                6 => Some(mov.source.sout_one()),
+                rank => panic!("Double-push was used from invalid rank({}) when trying to make {:?}", rank, mov),
+            }
+        } else { None };
+
         self.ply += 1;
 
         captured_piece
@@ -345,13 +416,39 @@ impl Board {
                 None => panic!("Trying to unmake move which was not made: no piece was found on target square"),
             };
 
+        // If unmaking castle, also return rook to its place
+        if mov.kind == MoveKind::Castle {
+            let (rook_source_file, rook_target_file) = match mov.target.file() {
+                2 => (0, 3),
+                6 => (7, 5),
+                _ => panic!("Malformed castle, target square invalid: {:?}", mov),
+            };
+            let rook_source_bb = Square::from_coords(mov.target.rank(), rook_source_file).to_bitboard();
+            let rook_target_bb = Square::from_coords(mov.target.rank(), rook_target_file).to_bitboard();
+
+            match self.pieces
+                .iter()
+                .enumerate()
+                .find(|(rook_type, bitboard)| *bitboard & rook_target_bb > 0)
+                {
+                    Some((rook_type, _)) => {
+                        self.pieces[rook_type] |= rook_source_bb;
+                        self.occupancy         |= rook_source_bb;
+
+                        self.pieces[rook_type] ^= rook_target_bb;
+                        self.occupancy         ^= rook_target_bb;
+                    },
+                    None => panic!("Rook was not found when castling"),
+                }
+        }
+
         // Return captured piece to target square
         match captured_piece {
             Some(target_piece) => {
                 match mov.kind {
                     // Return pawn captured by En Passant pawn if needed
                     MoveKind::EnPassant => {
-                        let original_dead_pawn_bb = Square::from(mov.source.rank() * 8 + mov.target.file()).to_bitboard();
+                        let original_dead_pawn_bb = Square::from_coords(mov.source.rank(), mov.target.file()).to_bitboard();
                         self.pieces[target_piece as usize] |= original_dead_pawn_bb;
                         self.occupancy                     |= original_dead_pawn_bb;
                     },
@@ -369,9 +466,9 @@ impl Board {
         self.ply -= 1;
     }
 
-    fn perft(&mut self, depth: u8, print: bool) -> (u64, u64, u64, u64) {
+    fn perft(&mut self, depth: u8, print: bool) -> (u64, u64, u64, u64, u64) {
         if depth == 0 {
-            return (1, 0, 0, 0) // This a leaf, exactly one node
+            return (1, 0, 0, 0, 0) // This a leaf, exactly one node
         }
         let color = self.color_to_move();
 
@@ -385,6 +482,7 @@ impl Board {
         let mut total = 0;
         let mut captures = 0;
         let mut checks = 0;
+        let mut castles = 0;
         let mut en_passants = 0;
 
         for mov in moves {
@@ -401,6 +499,9 @@ impl Board {
                             en_passants += 1;
                             captures += 1;
                         }
+                        MoveKind::Castle => {
+                            castles += 1;
+                        }
                         _ => {}
                     }
                     if self.is_king_in_check(color.flip()) {
@@ -412,10 +513,11 @@ impl Board {
                     println!("{:?}", mov);
                     self.print();
                 }
-                let (children_total, children_tactical, children_checks, children_ep) = self.perft(depth - 1, print);
+                let (children_total, children_tactical, children_checks, children_castles, children_ep) = self.perft(depth - 1, print);
                 total += children_total;
                 captures += children_tactical;
                 checks += children_checks;
+                castles += children_castles;
                 en_passants += children_ep;
 
             }
@@ -426,7 +528,7 @@ impl Board {
             println!("Found {} nodes in this subtree (depth {})", total, depth);
         }
 
-        (total, captures, checks, en_passants)
+        (total, captures, checks, castles, en_passants)
     }
 
     fn is_square_attacked(&self, square: Square, attacker_color: Color) -> bool {
@@ -609,22 +711,23 @@ mod tests {
     fn test_perft() {
         let mut board = Board::new();
 
-        assert_eq!(board.perft(0, false), (1, 0, 0, 0));
-        assert_eq!(board.perft(1, false), (20, 0, 0, 0));
-        assert_eq!(board.perft(2, false), (400, 0, 0, 0));
-        assert_eq!(board.perft(3, false), (8902, 34, 12, 0));
-        assert_eq!(board.perft(4, false), (197281, 1576, 469, 0));
-        // assert_eq!(board.perft(5, false), (4865609, 82719, 27351, 258));
-        // assert_eq!(board.perft(6, false), (119060324 , 2812008, 809099 , 5248));
+        assert_eq!(board.perft(0, false), (1, 0, 0, 0, 0));
+        assert_eq!(board.perft(1, false), (20, 0, 0, 0, 0));
+        assert_eq!(board.perft(2, false), (400, 0, 0, 0, 0));
+        assert_eq!(board.perft(3, false), (8902, 34, 12, 0, 0));
+        assert_eq!(board.perft(4, false), (197281, 1576, 469, 0, 0));
+        // assert_eq!(board.perft(5, false), (4865609, 82719, 27351, 0, 258));
+        // assert_eq!(board.perft(6, false), (119060324, 2812008, 809099, 0, 5248));
     }
 
     #[test]
     fn test_position_perft() {
         let fen = String::from("r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq - ");
         let mut board = Board::from_FEN(fen);
-        assert_eq!(board.perft(0, false), (1, 0, 0, 0));
-        assert_eq!(board.perft(1, false), (48, 8, 0, 0));
-        assert_eq!(board.perft(2, false), (2039, 351, 3, 1));
+        assert_eq!(board.perft(0, false), (1, 0, 0, 0, 0));
+        assert_eq!(board.perft(1, false), (48, 8, 0, 2, 0));
+        assert_eq!(board.perft(2, false), (2039, 351, 3, 91, 1));
+        assert_eq!(board.perft(3, false), (97862, 17102, 993, 3162, 45));
     }
 
     #[test]
diff --git a/src/moves.rs b/src/moves.rs
index d0e3eb6..f31d864 100644
--- a/src/moves.rs
+++ b/src/moves.rs
@@ -1,6 +1,6 @@
 use crate::{square::Square, bitboard::print};
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, PartialEq, Eq, Copy)]
 pub enum MoveKind {
     Quiet,
     Capture,
diff --git a/src/square.rs b/src/square.rs
index 1adca83..efa682c 100644
--- a/src/square.rs
+++ b/src/square.rs
@@ -16,6 +16,10 @@ pub enum Square {
   A8, B8, C8, D8, E8, F8, G8, H8,
 }
 impl Square {
+    pub fn from_coords(rank: u8, file: u8) -> Self {
+        Self::from(rank * 8 + file)
+    }
+
     pub fn to_bitboard(&self) -> Bitboard {
         1u64 << *self as u8
     }
@@ -35,6 +39,16 @@ impl Square {
     pub fn sout_one(&self) -> Self {
         Self::from(*self as u8 - 8)
     }
+
+    /// Warning: this wraps which is not what you might expect!
+    pub fn west_one(&self) -> Self {
+        Self::from(*self as u8 - 1)
+    }
+
+    /// Warning: this wraps which is not what you might expect!
+    pub fn east_one(&self) -> Self {
+        Self::from(*self as u8 + 1)
+    }
 }