feat: compute rook attacks

1 files changed, 60 insertions, 5 deletions

diff --git a/src/attacks.rs b/src/attacks.rs
index 4d826f6..6cf2a66 100644
--- a/src/attacks.rs
+++ b/src/attacks.rs
@@ -5,6 +5,9 @@ static NOT_B_FILE: Bitboard = 0xFDFDFDFDFDFDFDFD;
 static NOT_G_FILE: Bitboard = 0xBFBFBFBFBFBFBFBF;
 static NOT_H_FILE: Bitboard = 0x7F7F7F7F7F7F7F7F;
 static B_FILE:     Bitboard = 0x0202020202020202;
+static H_FILE:     Bitboard = 0x8080808080808080;
+static DIAG_C2_H7: Bitboard = 0x0080402010080400;
+static DIAG_A1_H8: Bitboard = 0x8040201008040201;
 
 /// An array where N-th item is an attack bitboard
 /// of a piece on N-th square
@@ -215,7 +218,7 @@ impl Attacks {
     ///
     /// Given a square and occupancy masked for rank, diagonal or anti-diagonal (note: not a file!)
     /// return an attack bitboard that considers blocking pieces
-    fn kingergarten_attacks_base(&self, occupancy: Bitboard, mask: Bitboard, square: u8) -> Bitboard {
+    fn kindergarten_attacks_base(&self, occupancy: Bitboard, mask: Bitboard, square: u8) -> Bitboard {
         let file = square % 8;
         let masked_occupancy = occupancy & mask;
         let occupancy_rank = ((masked_occupancy as u128 * B_FILE as u128) >> 58 & 0b111111) << 1;
@@ -228,13 +231,39 @@ impl Attacks {
         filled_up_attacks & mask
     }
 
+    /// https://www.chessprogramming.org/Kindergarten_Bitboards
+    fn kindergarten_attacks_file(&self, occupancy: Bitboard, mask: Bitboard, square: u8) -> Bitboard {
+        let file = square % 8;
+        let rank  = square / 8;
+
+        let masked_occupancy = (occupancy & mask) >> file; // Shift occupancy to A file
+        let occupancy_rank = ((masked_occupancy as u128 * DIAG_C2_H7 as u128) >> 58 & 0b111111) << 1;
+        // Use reversed rank as index, since occupancy is reversed
+        let rank_attacks = self.first_rank_attacks[occupancy_rank as usize][7 - rank as usize] as Bitboard;
+
+        ((rank_attacks as u128 * DIAG_A1_H8 as u128) as Bitboard & H_FILE) >> (7 - file)
+    }
+
     fn bishop(&self, occupancy: Bitboard, square: u8) -> Bitboard {
-        let diagonal_mask = self.ray_attacks[Direction::NoEa as usize][square as usize] |
+        let diagonal_mask =
+            self.ray_attacks[Direction::NoEa as usize][square as usize] |
             self.ray_attacks[Direction::SoWe as usize][square as usize];
-        let anti_diagonal_mask = self.ray_attacks[Direction::NoWe as usize][square as usize] |
+        let anti_diagonal_mask =
+            self.ray_attacks[Direction::NoWe as usize][square as usize] |
             self.ray_attacks[Direction::SoEa as usize][square as usize];
 
-        self.kingergarten_attacks_base(occupancy, diagonal_mask, square) | self.kingergarten_attacks_base(occupancy, anti_diagonal_mask, square)
+        self.kindergarten_attacks_base(occupancy, diagonal_mask, square) | self.kindergarten_attacks_base(occupancy, anti_diagonal_mask, square)
+    }
+
+    fn rook(&self, occupancy: Bitboard, square: u8) -> Bitboard {
+        let vertical =
+            self.ray_attacks[Direction::Nort as usize][square as usize] |
+            self.ray_attacks[Direction::Sout as usize][square as usize];
+        let horizontal =
+            self.ray_attacks[Direction::West as usize][square as usize] |
+            self.ray_attacks[Direction::East as usize][square as usize];
+
+        self.kindergarten_attacks_file(occupancy, vertical, square) | self.kindergarten_attacks_base(occupancy, horizontal, square)
     }
 }
 
@@ -318,7 +347,7 @@ mod tests {
     }
 
     #[test]
-    fn test_diagonal_attacks() {
+    fn test_bishop_attacks() {
         let attacks = Attacks::new();
         let square = Square::E4 as u8;
         let occupancy =
@@ -338,4 +367,30 @@ mod tests {
 
         print(bb);
     }
+
+    #[test]
+    fn test_rook_attacks() {
+        let attacks = Attacks::new();
+        let square = Square::E4 as u8;
+        let occupancy =
+            1 << Square::B7 as usize |
+            1 << Square::B1 as usize |
+            1 << Square::C2 as usize |
+            1 << Square::E3 as usize |
+            1 << Square::F3 as usize;
+        let bb = attacks.rook(occupancy, square);
+
+        assert_ne!(bb & 1 << Square::E8 as u8, 0);
+        assert_ne!(bb & 1 << Square::E7 as u8, 0);
+        assert_ne!(bb & 1 << Square::E6 as u8, 0);
+        assert_ne!(bb & 1 << Square::E5 as u8, 0);
+        assert_ne!(bb & 1 << Square::E3 as u8, 0);
+        assert_eq!(bb & 1 << Square::E2 as u8, 0);
+        assert_eq!(bb & 1 << Square::E1 as u8, 0);
+        assert_ne!(bb & 1 << Square::A4 as u8, 0);
+        assert_ne!(bb & 1 << Square::H4 as u8, 0);
+        assert_eq!(bb & 1 << Square::E4 as u8, 0);
+
+        print(bb);
+    }
 }