diff --git a/position/src/check.rs b/position/src/check.rs index d9d5172..df30749 100644 --- a/position/src/check.rs +++ b/position/src/check.rs @@ -1,13 +1,12 @@ // Eryn Wells use chessfriend_bitboard::BitBoard; +use chessfriend_core::{Color, Direction, Shape, Square}; + +use crate::sight::SliderRayToSquareExt; pub struct CheckingPieces { - pawn: BitBoard, - knight: BitBoard, - bishop: BitBoard, - rook: BitBoard, - queen: BitBoard, + bitboards: [BitBoard; 5], } impl CheckingPieces { @@ -19,11 +18,80 @@ impl CheckingPieces { queen: BitBoard, ) -> CheckingPieces { CheckingPieces { - pawn, - knight, - bishop, - rook, - queen, + bitboards: [pawn, 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]); + } }