1 files changed, 52 insertions, 6 deletions
diff --git a/src/board.rs b/src/board.rs
index 88b87a9..4b8f320 100644
--- a/src/board.rs
+++ b/src/board.rs
@@ -1,4 +1,4 @@
-use crate::{bitboard::{Bitboard, serialize_bitboard}, moves::Move, attacks::Attacks};
+use crate::{bitboard::{Bitboard, serialize_bitboard}, moves::Move, attacks::Attacks, square::Square};
 
 pub static DEFAULT_FEN: &str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1";
 
@@ -102,12 +102,11 @@ impl Board {
         !self.occupancy
     }
 
-    fn pieces_by_color(&self, color: Color) -> [Bitboard; 6] {
-        let mut pieces = [0; 6];
-        for piece_type in 0..6 {
-            pieces[piece_type] = self.pieces[piece_type + 6 * color as usize];
+    fn pieces_by_color(&self, color: Color) -> &[Bitboard] {
+        match color {
+            Color::White => &self.pieces[0..6],
+            Color::Black => &self.pieces[6..12],
         }
-        pieces
     }
 
     fn color_occupancy(&self, color: Color) -> Bitboard {
@@ -318,6 +317,42 @@ impl Board {
 
         total
     }
+
+    fn is_square_attacked(&self, square: Square, attacker_color: Color) -> bool {
+        let square_bb = square.to_bitboard();
+        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) {
+                        return true
+                    }
+                }
+                PieceType::Knight => {
+                    if (self.attacks.knight[square as usize] & piece > 0) {
+                        return true
+                    }
+                }
+                PieceType::Bishop => {
+                    if (self.attacks.bishop(self.occupancy, square) & piece > 0) {
+                        return true
+                    }
+                }
+                PieceType::Rook => {
+                    if (self.attacks.rook(self.occupancy, square) & piece > 0) {
+                        return true
+                    }
+                }
+                PieceType::Queen => {
+                    if (self.attacks.queen(self.occupancy, square) & piece > 0) {
+                        return true
+                    }
+                }
+                PieceType::King => {}
+                _ => panic!("Unexpected piece type! Pieces by color should be considered white")
+            }
+        }
+        false
+    }
 }
 
 
@@ -454,4 +489,15 @@ mod tests {
         assert_eq!(board.perft(2, false), 400);
         // TODO: assert_eq!(board.perft(3, false), 8902);
     }
+
+    #[test]
+    fn test_is_square_attacked() {
+        let fen = String::from(DEFAULT_FEN);
+        let board = Board::from_FEN(fen);
+
+        assert_eq!(board.is_square_attacked(Square::E2, Color::White), true);
+        assert_eq!(board.is_square_attacked(Square::E2, Color::Black), false);
+        assert_eq!(board.is_square_attacked(Square::E4, Color::White), false);
+        assert_eq!(board.is_square_attacked(Square::B6, Color::Black), true);
+    }
 }