chessfriend/position/src/sight.rs

// Eryn Wells <eryn@erynwells.me>

use crate::position::piece_sets::PieceBitBoards;
use chessfriend_bitboard::BitBoard;
use chessfriend_core::{Color, Direction, PlacedPiece, Shape, Square};

pub(crate) trait SightExt {
    fn sight(&self, pieces: &PieceBitBoards, en_passant_square: Option<Square>) -> BitBoard;

    fn white_pawn_sight(
        &self,
        pieces: &PieceBitBoards,
        en_passant_square: Option<Square>,
    ) -> BitBoard;

    fn black_pawn_sight(
        &self,
        pieces: &PieceBitBoards,
        en_passant_square: Option<Square>,
    ) -> BitBoard;

    fn knight_sight(&self, pieces: &PieceBitBoards) -> BitBoard;
    fn bishop_sight(&self, pieces: &PieceBitBoards) -> BitBoard;
    fn rook_sight(&self, pieces: &PieceBitBoards) -> BitBoard;
    fn queen_sight(&self, pieces: &PieceBitBoards) -> BitBoard;
    fn king_sight(&self, pieces: &PieceBitBoards) -> BitBoard;
}

pub(crate) trait SliderRayToSquareExt {
    fn ray_to_square(&self, square: Square) -> Option<BitBoard>;
}

impl SightExt for PlacedPiece {
    fn sight(&self, pieces: &PieceBitBoards, en_passant_square: Option<Square>) -> BitBoard {
        match self.shape() {
            Shape::Pawn => match self.color() {
                Color::White => self.white_pawn_sight(pieces, en_passant_square),
                Color::Black => self.black_pawn_sight(pieces, en_passant_square),
            },
            Shape::Knight => self.knight_sight(pieces),
            Shape::Bishop => self.bishop_sight(pieces),
            Shape::Rook => self.rook_sight(pieces),
            Shape::Queen => self.queen_sight(pieces),
            Shape::King => self.king_sight(pieces),
        }
    }

    fn white_pawn_sight(
        &self,
        pieces: &PieceBitBoards,
        en_passant_square: Option<Square>,
    ) -> BitBoard {
        let opponent = self.color().other();
        let pawn: BitBoard = self.square().into();
        let pawn = pawn.shift_north_west_one() | pawn.shift_north_east_one();

        let mut possible_squares = pieces.empty_squares() | pieces.all_pieces_of_color(opponent);
        if let Some(en_passant) = en_passant_square {
            let en_passant_bitboard: BitBoard = en_passant.into();
            possible_squares |= en_passant_bitboard;
        }

        pawn & possible_squares
    }

    fn black_pawn_sight(
        &self,
        pieces: &PieceBitBoards,
        en_passant_square: Option<Square>,
    ) -> BitBoard {
        let opponent = self.color().other();

        let pawn: BitBoard = self.square().into();
        let pawn = pawn.shift_south_west_one() | pawn.shift_south_east_one();

        let mut possible_squares = pieces.empty_squares() | pieces.all_pieces_of_color(opponent);
        if let Some(en_passant) = en_passant_square {
            possible_squares |= &en_passant.into();
        }

        pawn & possible_squares
    }

    fn knight_sight(&self, pieces: &PieceBitBoards) -> BitBoard {
        BitBoard::knight_moves(self.square()) & !pieces.all_pieces_of_color(self.color())
    }

    fn bishop_sight(&self, pieces: &PieceBitBoards) -> BitBoard {
        let square = self.square();

        let mut sight = BitBoard::empty();

        let blockers = pieces.all_pieces();

        macro_rules! update_moves_with_ray {
            ($direction:ident, $occupied_squares:tt) => {
                let ray = BitBoard::ray(square, Direction::$direction);
                if let Some(first_occupied_square) =
                    BitBoard::$occupied_squares(&(ray & blockers)).next()
                {
                    let remainder = BitBoard::ray(first_occupied_square, Direction::$direction);
                    let attack_ray = ray & !remainder;
                    sight |= attack_ray;
                } else {
                    sight |= ray;
                }
            };
        }

        update_moves_with_ray!(NorthEast, occupied_squares_trailing);
        update_moves_with_ray!(NorthWest, occupied_squares_trailing);
        update_moves_with_ray!(SouthEast, occupied_squares);
        update_moves_with_ray!(SouthWest, occupied_squares);

        let friendly_pieces = pieces.all_pieces_of_color(self.color());
        sight & !friendly_pieces
    }

    fn rook_sight(&self, pieces: &PieceBitBoards) -> BitBoard {
        let square = self.square();

        let mut sight = BitBoard::empty();

        let blockers = pieces.all_pieces();

        macro_rules! update_moves_with_ray {
            ($direction:ident, $occupied_squares:tt) => {
                let ray = BitBoard::ray(square, Direction::$direction);
                if let Some(first_occupied_square) =
                    BitBoard::$occupied_squares(&(ray & blockers)).next()
                {
                    let remainder = BitBoard::ray(first_occupied_square, Direction::$direction);
                    let attack_ray = ray & !remainder;
                    sight |= attack_ray;
                } else {
                    sight |= ray;
                }
            };
        }

        update_moves_with_ray!(North, occupied_squares_trailing);
        update_moves_with_ray!(East, occupied_squares_trailing);
        update_moves_with_ray!(South, occupied_squares);
        update_moves_with_ray!(West, occupied_squares);

        let friendly_pieces = pieces.all_pieces_of_color(self.color());
        sight & !friendly_pieces
    }

    fn queen_sight(&self, pieces: &PieceBitBoards) -> BitBoard {
        let square = self.square();

        let mut sight = BitBoard::empty();

        let blockers = pieces.all_pieces();

        macro_rules! update_moves_with_ray {
            ($direction:ident, $occupied_squares:tt) => {
                let ray = BitBoard::ray(square, Direction::$direction);
                if let Some(first_occupied_square) =
                    BitBoard::$occupied_squares(&(ray & blockers)).next()
                {
                    let remainder = BitBoard::ray(first_occupied_square, Direction::$direction);
                    let attack_ray = ray & !remainder;
                    sight |= attack_ray;
                } else {
                    sight |= ray;
                }
            };
        }

        update_moves_with_ray!(NorthWest, occupied_squares_trailing);
        update_moves_with_ray!(North, occupied_squares_trailing);
        update_moves_with_ray!(NorthEast, occupied_squares_trailing);
        update_moves_with_ray!(East, occupied_squares_trailing);
        update_moves_with_ray!(SouthEast, occupied_squares);
        update_moves_with_ray!(South, occupied_squares);
        update_moves_with_ray!(SouthWest, occupied_squares);
        update_moves_with_ray!(West, occupied_squares);

        let friendly_pieces = pieces.all_pieces_of_color(self.color());
        sight & !friendly_pieces
    }

    fn king_sight(&self, pieces: &PieceBitBoards) -> BitBoard {
        BitBoard::king_moves(self.square()) & !pieces.all_pieces_of_color(self.color())
    }
}

impl SliderRayToSquareExt for PlacedPiece {
    fn ray_to_square(&self, target: Square) -> Option<BitBoard> {
        macro_rules! ray {
            ($square:expr, $direction:ident) => {
                (
                    BitBoard::ray($square, Direction::$direction),
                    Direction::$direction,
                )
            };
        }

        let square = self.square();
        let target_bitboard: BitBoard = target.into();

        let ray_and_direction = match self.shape() {
            Shape::Bishop => [
                ray!(square, NorthEast),
                ray!(square, SouthEast),
                ray!(square, SouthWest),
                ray!(square, NorthWest),
            ]
            .into_iter()
            .find(|(&ray, _)| !(target_bitboard & ray).is_empty()),
            Shape::Rook => [
                ray!(square, North),
                ray!(square, East),
                ray!(square, South),
                ray!(square, West),
            ]
            .into_iter()
            .find(|(&ray, _)| !(target_bitboard & ray).is_empty()),
            Shape::Queen => [
                ray!(square, North),
                ray!(square, NorthEast),
                ray!(square, East),
                ray!(square, SouthEast),
                ray!(square, South),
                ray!(square, SouthWest),
                ray!(square, West),
                ray!(square, NorthWest),
            ]
            .into_iter()
            .find(|(&ray, _)| !(target_bitboard & ray).is_empty()),
            _ => None,
        };

        if let Some((ray, direction)) = ray_and_direction {
            let first_occupied_square = match direction {
                Direction::East
                | Direction::NorthWest
                | Direction::NorthEast
                | Direction::North => ray.first_occupied_square_trailing(),
                Direction::West
                | Direction::SouthWest
                | Direction::SouthEast
                | Direction::South => ray.first_occupied_square(),
            };

            if let Some(occupied_square) = first_occupied_square {
                let remainder = BitBoard::ray(target, direction);
                return Some(ray & !remainder);
            }
        }

        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chessfriend_bitboard::bitboard;
    use chessfriend_core::{piece, Square};

    macro_rules! sight_test {
        ($test_name:ident, $position:expr, $piece:expr, $bitboard:expr) => {
            #[test]
            fn $test_name() {
                let pos = $position;
                let piece = $piece;
                let sight = pos.sight_of_piece(&piece);

                assert_eq!(sight, $bitboard);
            }
        };
        ($test_name:ident, $piece:expr, $bitboard:expr) => {
            sight_test! {$test_name, $crate::Position::empty(), $piece, $bitboard}
        };
    }

    #[test]
    fn pawns_and_knights_cannot_make_rays() {
        assert_eq!(piece!(White Pawn on F7).ray_to_square(Square::E8), None);
        assert_eq!(piece!(White Knight on F6).ray_to_square(Square::E8), None);
    }

    mod pawn {
        use crate::test_position;
        use chessfriend_bitboard::{bitboard, BitBoard};
        use chessfriend_core::{piece, Square};

        sight_test!(e4_pawn, piece!(White Pawn on E4), bitboard!(D5, F5));

        sight_test!(
            e4_pawn_one_blocker,
            test_position![
                White Bishop on D5,
                White Pawn on E4,
            ],
            piece!(White Pawn on E4),
            bitboard!(F5)
        );

        #[test]
        fn e4_pawn_two_blocker() {
            let pos = test_position!(
                White Bishop on D5,
                White Queen on F5,
                White Pawn on E4,
            );

            let piece = piece!(White Pawn on E4);
            let sight = pos.sight_of_piece(&piece);

            assert_eq!(sight, BitBoard::empty());
        }

        #[test]
        fn e4_pawn_capturable() {
            let pos = test_position!(
                Black Bishop on D5,
                White Queen on F5,
                White Pawn on E4,
            );

            let piece = piece!(White Pawn on E4);
            let sight = pos.sight_of_piece(&piece);

            assert_eq!(sight, bitboard!(D5));
        }

        #[test]
        fn e5_en_passant() {
            let mut pos = test_position!(
                White Pawn on E5,
                Black Pawn on D5,
            );
            pos.test_set_en_passant_square(Square::D6);
            let piece = piece!(White Pawn on E5);
            let sight = pos.sight_of_piece(&piece);

            assert_eq!(sight, bitboard!(D6, F6));
        }
    }

    #[macro_use]
    mod knight {
        use chessfriend_bitboard::bitboard;
        use chessfriend_core::piece;

        sight_test!(
            f6_knight,
            piece!(Black Knight on F6),
            bitboard!(H7, G8, E8, D7, D5, E4, G4, H5)
        );
    }

    mod bishop {
        use super::*;

        sight_test!(
            c2_bishop,
            piece!(Black Bishop on C2),
            bitboard!(D1, B3, A4, B1, D3, E4, F5, G6, H7)
        );

        #[test]
        fn ray_to_square() {
            let generated_ray = piece!(White Bishop on C5).ray_to_square(Square::E7);
            let expected_ray = bitboard![D6, E7];
            assert_eq!(generated_ray, Some(expected_ray));
        }
    }

    mod rook {
        use super::*;
        use crate::test_position;

        sight_test!(
            g3_rook,
            piece!(White Rook on G3),
            bitboard!(G1, G2, G4, G5, G6, G7, G8, A3, B3, C3, D3, E3, F3, H3)
        );

        sight_test!(
            e4_rook_with_e1_white_king_e7_black_king,
            test_position![
                White Rook on E4,
                White King on E1,
                Black King on E7,
            ],
            piece!(White Rook on E4),
            bitboard!(A4, B4, C4, D4, F4, G4, H4, E2, E3, E5, E6, E7)
        );

        #[test]
        fn ray_to_square() {
            let generated_ray = piece!(White Rook on C2).ray_to_square(Square::C6);
            let expected_ray = bitboard![C3, C4, C5, C6];
            assert_eq!(generated_ray, Some(expected_ray));

            let generated_ray = piece!(White Rook on D2).ray_to_square(Square::H2);
            let expected_ray = bitboard![E2, F2, G2, H2];
            assert_eq!(generated_ray, Some(expected_ray));

            let generated_ray = piece!(White Rook on G6).ray_to_square(Square::B6);
            let expected_ray = bitboard![B6, C6, D6, E6, F6];
            assert_eq!(generated_ray, Some(expected_ray));

            let generated_ray = piece!(White Rook on A6).ray_to_square(Square::A3);
            let expected_ray = bitboard![A5, A4, A3];
            assert_eq!(generated_ray, Some(expected_ray));
        }
    }

    mod king {
        use chessfriend_bitboard::bitboard;
        use chessfriend_core::piece;

        sight_test!(
            e1_king,
            piece!(White King on E1),
            bitboard![D1, D2, E2, F2, F1]
        );
    }
}