Add move generation for castling

src/movegen.rs

@@ -3,6 +3,7 @@ use crate::attack::{
     get_rook_attacks,
 };
 use crate::board::Color;
+use crate::game::CastlingRights;
 use u64 as Bitboard;
 
 const NOT_A_FILE: Bitboard = 0xfefefefefefefefe;
@@ -241,9 +242,15 @@ pub fn queen_pseudo_moves(
     moves
 }
 
-pub fn king_pseudo_moves(king: Bitboard, occupancies: Bitboard) -> Vec<Move> {
+pub fn king_pseudo_moves(
+    king: Bitboard,
+    all_occupancies: Bitboard,
+    own_occupancies: Bitboard,
+    castling_rights: CastlingRights,
+    color: Color,
+) -> Vec<Move> {
     let mut moves: Vec<Move> = vec![];
-    let empty: u64 = !occupancies;
+    let empty: u64 = !own_occupancies;
     let king_square = king.trailing_zeros() as usize;
     let from = king_square as u32;
     let mut attacks = get_king_attacks(king_square) & empty;
@@ -252,6 +259,50 @@ pub fn king_pseudo_moves(king: Bitboard, occupancies: Bitboard) -> Vec<Move> {
         moves.push(Move::new(from, attacks.trailing_zeros()));
         attacks &= attacks - 1;
     }
+
+    match color {
+        Color::White => match castling_rights {
+            CastlingRights::Both => {
+                if all_occupancies & 0x60 == 0 {
+                    moves.push(Move::new(4, 6))
+                } else if all_occupancies & 0xe == 0 {
+                    moves.push(Move::new(4, 2))
+                }
+            }
+            CastlingRights::Short => {
+                if all_occupancies & 0x60 == 0 {
+                    moves.push(Move::new(4, 6))
+                }
+            }
+            CastlingRights::Long => {
+                if all_occupancies & 0xe == 0 {
+                    moves.push(Move::new(4, 2))
+                }
+            }
+            CastlingRights::None => (),
+        },
+        Color::Black => match castling_rights {
+            CastlingRights::Both => {
+                if all_occupancies & 0x6000000000000000 == 0 {
+                    moves.push(Move::new(60, 62))
+                } else if all_occupancies & 0xe00000000000000 == 0 {
+                    moves.push(Move::new(60, 58))
+                }
+            }
+            CastlingRights::Short => {
+                if all_occupancies & 0x6000000000000000 == 0 {
+                    moves.push(Move::new(60, 62))
+                }
+            }
+            CastlingRights::Long => {
+                if all_occupancies & 0xe00000000000000 == 0 {
+                    moves.push(Move::new(60, 58))
+                }
+            }
+            CastlingRights::None => (),
+        },
+    }
+
     moves
 }
 
@@ -402,12 +453,16 @@ mod tests {
         Ok(())
     }
 
-    const FEN_KING_MOVES: &str = "R7/6k1/8/8/P6P/6K1/8/4r3 w - - 0 1";
+    const FEN_KING_MOVES: [&str; 3] = [
+        "R7/6k1/8/8/P6P/6K1/8/4r3 w - - 0 1",
+        "rnbqkb1r/ppp1pp1p/5np1/3p4/2PP4/1QN5/PP1BPPPP/R3KBNR w KQkq - 0 1",
+        "rnq1k2r/pp2p2p/2pbb1p1/5pB1/3P2n1/2N3P1/PPQ1PPBP/R3K1NR b KQkq - 0 1",
+    ];
 
     #[test]
     fn test_king_pseudo_moves() -> Result<(), String> {
         init_attacks();
-        let new_game = from_fen(FEN_KING_MOVES)?;
+        let new_game = from_fen(FEN_KING_MOVES[0])?;
         let expected = vec![
             Move::new(22, 13),
             Move::new(22, 14),
@@ -421,6 +476,24 @@ mod tests {
         actual.sort();
         assert_eq!(expected, actual);
 
+        let new_game_2 = from_fen(FEN_KING_MOVES[1])?;
+        let expected = vec![Move::new(4, 2), Move::new(4, 3)];
+        let mut actual = new_game_2.board.pseudo_moves(Color::White, Kind::King);
+        actual.sort();
+        assert_eq!(expected, actual);
+
+        let new_game_3 = from_fen(FEN_KING_MOVES[2])?;
+        let expected = vec![
+            Move::new(60, 51),
+            Move::new(60, 53),
+            Move::new(60, 59),
+            Move::new(60, 61),
+            Move::new(60, 62),
+        ];
+        let mut actual = new_game_3.board.pseudo_moves(Color::Black, Kind::King);
+        actual.sort();
+        assert_eq!(expected, actual);
+
         Ok(())
     }
 }