[position] Calculate push and capture masks on CheckingPieces

position/src/check.rs

@@ -1,13 +1,12 @@
 // Eryn Wells <eryn@erynwells.me>
 
 use chessfriend_bitboard::BitBoard;
+use chessfriend_core::{Color, Direction, Shape, Square};
+
+use crate::sight::SliderRayToSquareExt;
 
 pub struct CheckingPieces {
-    pawn: BitBoard,
+    bitboards: [BitBoard; 5],
-    knight: BitBoard,
-    bishop: BitBoard,
-    rook: BitBoard,
-    queen: BitBoard,
 }
 
 impl CheckingPieces {
@@ -19,11 +18,80 @@ impl CheckingPieces {
         queen: BitBoard,
     ) -> CheckingPieces {
         CheckingPieces {
-            pawn,
+            bitboards: [pawn, knight, bishop, rook, queen],
-            knight,
-            bishop,
-            rook,
-            queen,
         }
     }
+
+    pub fn count(&self) -> u32 {
+        self.bitboards.iter().map(|b| b.population_count()).sum()
+    }
+
+    /// A BitBoard representing the set of pieces that must be captured to
+    /// resolve check.
+    pub fn capture_mask(&self) -> BitBoard {
+        if self.count() == 0 {
+            BitBoard::FULL
+        } else {
+            self.bitboards
+                .iter()
+                .fold(BitBoard::EMPTY, std::ops::BitOr::bitor)
+        }
+    }
+
+    /// A BitBoard representing the set of squares to which a player can move a
+    /// piece to block a checking piece.
+    pub fn push_mask(&self, king: &BitBoard) -> BitBoard {
+        let target = king.first_occupied_square().unwrap();
+
+        macro_rules! push_mask_for_shape {
+            ($push_mask:expr, $shape:ident, $king:expr) => {{
+                let checking_pieces = self.bitboard_for_shape(Shape::$shape);
+                if !checking_pieces.is_empty() {
+                    if let Some(checking_ray) = checking_pieces
+                        .occupied_squares()
+                        .flat_map(|sq| Shape::$shape.ray_to_square(sq, target).into_iter())
+                        .find(|bb| !(bb & $king).is_empty())
+                    {
+                        $push_mask |= checking_ray & !$king
+                    }
+                }
+            }};
+        }
+
+        let mut push_mask = BitBoard::EMPTY;
+
+        push_mask_for_shape!(push_mask, Bishop, king);
+        push_mask_for_shape!(push_mask, Rook, king);
+        push_mask_for_shape!(push_mask, Queen, king);
+
+        push_mask
+    }
+
+    fn bitboard_for_shape(&self, shape: Shape) -> &BitBoard {
+        &self.bitboards[shape as usize]
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use chessfriend_bitboard::{bitboard, BitBoard};
+
+    /// This is a test position from [this execellent blog post][1] about how to
+    /// efficiently generate legal chess moves.
+    ///
+    /// [1]: https://peterellisjones.com/posts/generating-legal-chess-moves-efficiently/
+    #[test]
+    fn rook_push_mask() {
+        let checks = CheckingPieces::new(
+            BitBoard::EMPTY,
+            BitBoard::EMPTY,
+            BitBoard::EMPTY,
+            bitboard![E5],
+            BitBoard::EMPTY,
+        );
+
+        let push_mask = checks.push_mask(&bitboard![E8]);
+        assert_eq!(push_mask, bitboard![E6, E7]);
+    }
 }