Add move generation for castling

src/board.rs

@@ -5,7 +5,7 @@ use crate::attack::{
     get_bishop_attacks, get_king_attacks, get_knight_attacks, get_pawn_attacks, get_queen_attacks,
     get_rook_attacks,
 };
-use crate::game::State;
+use crate::game::{CastlingRights, State};
 
 #[derive(Debug, PartialEq, Eq)]
 pub struct Board {
@@ -128,7 +128,82 @@ impl Board {
             Kind::Bishop => bishop_pseudo_moves(pieces, all_occupancies, own_occupancies),
             Kind::Rook => rook_pseudo_moves(pieces, all_occupancies, own_occupancies),
             Kind::Queen => queen_pseudo_moves(pieces, all_occupancies, own_occupancies),
-            Kind::King => king_pseudo_moves(pieces, own_occupancies),
+            Kind::King => king_pseudo_moves(
+                pieces,
+                all_occupancies,
+                own_occupancies,
+                self.castling_ability(color),
+                color,
+            ),
+        }
+    }
+
+    fn castling_ability(&self, color: Color) -> CastlingRights {
+        match (
+            self.castling_from_fen(color),
+            self.castling_pseudo_safety(color),
+        ) {
+            (CastlingRights::Both, _) => self.castling_pseudo_safety(color),
+            (CastlingRights::Short, CastlingRights::Short) => CastlingRights::Short,
+            (CastlingRights::Long, CastlingRights::Long) => CastlingRights::Long,
+            _ => CastlingRights::None,
+        }
+    }
+
+    const fn castling_from_fen(&self, color: Color) -> CastlingRights {
+        let castling_fen = self.state.get_castling_ability();
+
+        let black_king_and_queen = (castling_fen >> 2) & 0b11 == 0b11;
+        let black_king = (castling_fen >> 3) & 1 == 1;
+        let black_queen = (castling_fen >> 2) & 1 == 1;
+        let white_king_and_queen = castling_fen & 0b11 == 0b11;
+        let white_king = (castling_fen >> 1) & 1 == 1;
+        let white_queen = castling_fen & 1 == 1;
+
+        match color {
+            Color::White => match (white_king_and_queen, white_king, white_queen) {
+                (true, _, _) => CastlingRights::Both,
+                (_, true, _) => CastlingRights::Short,
+                (_, _, true) => CastlingRights::Long,
+                _ => CastlingRights::None,
+            },
+            Color::Black => match (black_king_and_queen, black_king, black_queen) {
+                (true, _, _) => CastlingRights::Both,
+                (_, true, _) => CastlingRights::Short,
+                (_, _, true) => CastlingRights::Long,
+                _ => CastlingRights::None,
+            },
+        }
+    }
+
+    fn castling_pseudo_safety(&self, color: Color) -> CastlingRights {
+        match color {
+            Color::White => {
+                if self.is_attacked(4, Color::Black) {
+                    return CastlingRights::None;
+                }
+                let d1_attacked = self.is_attacked(3, Color::Black);
+                let f1_attacked = self.is_attacked(5, Color::Black);
+                match (d1_attacked, f1_attacked) {
+                    (true, true) => CastlingRights::None,
+                    (true, false) => CastlingRights::Short,
+                    (false, true) => CastlingRights::Long,
+                    (false, false) => CastlingRights::Both,
+                }
+            }
+            Color::Black => {
+                if self.is_attacked(60, Color::White) {
+                    return CastlingRights::None;
+                }
+                let d8_attacked = self.is_attacked(59, Color::White);
+                let f8_attacked = self.is_attacked(61, Color::White);
+                match (d8_attacked, f8_attacked) {
+                    (true, true) => CastlingRights::None,
+                    (true, false) => CastlingRights::Short,
+                    (false, true) => CastlingRights::Long,
+                    (false, false) => CastlingRights::Both,
+                }
+            }
         }
     }
 }

src/game.rs

@@ -70,6 +70,10 @@ impl State {
     pub const fn get_en_passant_target_square(&self) -> Option<u8> {
         self.en_passant_target_square
     }
+
+    pub const fn get_castling_ability(&self) -> u8 {
+        self.castling_ability
+    }
 }
 
 use std::fmt;
@@ -87,3 +91,10 @@ impl Default for State {
         Self::new()
     }
 }
+
+pub enum CastlingRights {
+    Short,
+    Long,
+    Both,
+    None,
+}

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(())
     }
 }