feat: implement is_king_in_check

1 files changed, 24 insertions, 4 deletions

diff --git a/src/board.rs b/src/board.rs
index 4b8f320..bbcd5c7 100644
--- a/src/board.rs
+++ b/src/board.rs
@@ -1,4 +1,4 @@
-use crate::{bitboard::{Bitboard, serialize_bitboard}, moves::Move, attacks::Attacks, square::Square};
+use crate::{bitboard::{Bitboard, serialize_bitboard, bitscan}, moves::Move, attacks::Attacks, square::Square};
 
 pub static DEFAULT_FEN: &str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1";
 
@@ -146,7 +146,7 @@ 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(1 - color as u8));
+        let opponent_occupancy = self.color_occupancy(Color::from(color as u8).flip());
         let empty = self.empty();
         let available_targets = opponent_occupancy | empty;
         for (piece_type, piece) in self.pieces_by_color(color).iter().enumerate() {
@@ -307,7 +307,10 @@ impl Board {
                 println!("{:?}", mov);
                 self.print();
             }
-            total += self.perft(depth - 1, print);
+            // King can not be in check after our own move
+            if !self.is_king_in_check(color) {
+                total += self.perft(depth - 1, print);
+            }
             self.unmake_move(mov, captured_piece);
         }
 
@@ -323,7 +326,7 @@ impl Board {
         for (piece_type, piece) in self.pieces_by_color(attacker_color).iter().enumerate() {
             match PieceType::from(piece_type) {
                 PieceType::Pawn => {
-                    if (self.attacks.pawn[1 - attacker_color as usize][square as usize] & piece > 0) {
+                    if (self.attacks.pawn[attacker_color.flip() as usize][square as usize] & piece > 0) {
                         return true
                     }
                 }
@@ -353,6 +356,15 @@ impl Board {
         }
         false
     }
+
+    fn is_king_in_check(&self, color: Color) -> bool {
+        let king_bb = match color {
+            Color::White => self.pieces[PieceType::King as usize],
+            Color::Black => self.pieces[PieceType::KingBlack as usize],
+        };
+        let square = bitscan(king_bb);
+        self.is_square_attacked(square, color.flip())
+    }
 }
 
 
@@ -363,6 +375,14 @@ pub enum Color {
     White,
     Black,
 }
+impl Color {
+    pub fn flip(&self) -> Self {
+        match self {
+            Self::White => Self::Black,
+            Self::Black => Self::White,
+        }
+    }
+}
 
 #[cfg(test)]
 mod tests {