chessfriend/position/src/sight.rs

// Eryn Wells <eryn@erynwells.me>

//! Defines routines for computing sight of a piece. Sight is the set of squares
//! that a piece can see. In other words, it's the set of squares attacked or
//! controled by a piece.

use chessfriend_bitboard::BitBoard;
use chessfriend_board::Board;
use chessfriend_core::{Color, Direction, PlacedPiece, Shape, Square};

macro_rules! ray_in_direction {
    ($square:expr, $blockers:expr, $direction:ident, $first_occupied_square:tt) => {{
        let ray = BitBoard::ray($square, Direction::$direction);
        let ray_blockers = ray & $blockers;
        if let Some(first_occupied_square) = ray_blockers.$first_occupied_square() {
            let remainder = BitBoard::ray(first_occupied_square, Direction::$direction);
            let attack_ray = ray & !remainder;
            attack_ray
        } else {
            ray
        }
    }};
}

/// Compute sight of a white pawn.
fn _white_pawn_sight(
    pawn: BitBoard,
    occupancy: BitBoard,
    blockers: BitBoard,
    en_passant_square: BitBoard,
) -> BitBoard {
    let possible_squares = !occupancy | blockers | en_passant_square;
    let pawn = pawn.shift_north_west_one() | pawn.shift_north_east_one();
    pawn & possible_squares
}

fn _black_pawn_sight(
    pawn: BitBoard,
    occupancy: BitBoard,
    blockers: BitBoard,
    en_passant_square: BitBoard,
) -> BitBoard {
    let possible_squares = !occupancy | blockers | en_passant_square;
    let pawn = pawn.shift_south_west_one() | pawn.shift_south_east_one();
    pawn & possible_squares
}

fn _knight_sight(knight_square: Square, blockers: BitBoard) -> BitBoard {
    BitBoard::knight_moves(knight_square) & !blockers
}

fn _bishop_sight(bishop_square: Square, occupancy: BitBoard) -> BitBoard {
    #[rustfmt::skip]
    let sight = ray_in_direction!(bishop_square, occupancy, NorthEast, first_occupied_square_trailing)
        | ray_in_direction!(bishop_square, occupancy, NorthWest, first_occupied_square_trailing)
        | ray_in_direction!(bishop_square, occupancy, SouthEast, first_occupied_square_leading)
        | ray_in_direction!(bishop_square, occupancy, SouthWest, first_occupied_square_leading);

    sight
}

fn _rook_sight(rook_square: Square, occupancy: BitBoard) -> BitBoard {
    #[rustfmt::skip]
    let sight = ray_in_direction!(rook_square, occupancy, North, first_occupied_square_trailing)
        | ray_in_direction!(rook_square, occupancy, East, first_occupied_square_trailing)
        | ray_in_direction!(rook_square, occupancy, South, first_occupied_square_leading)
        | ray_in_direction!(rook_square, occupancy, West, first_occupied_square_leading);

    sight
}

fn _queen_sight(queen_square: Square, occupancy: BitBoard) -> BitBoard {
    #[rustfmt::skip]
    let sight = ray_in_direction!(queen_square, occupancy, NorthWest, first_occupied_square_trailing)
        | ray_in_direction!(queen_square, occupancy, North, first_occupied_square_trailing)
        | ray_in_direction!(queen_square, occupancy, NorthEast, first_occupied_square_trailing)
        | ray_in_direction!(queen_square, occupancy, East, first_occupied_square_trailing)
        | ray_in_direction!(queen_square, occupancy, SouthEast, first_occupied_square_leading)
        | ray_in_direction!(queen_square, occupancy, South, first_occupied_square_leading)
        | ray_in_direction!(queen_square, occupancy, SouthWest, first_occupied_square_leading)
        | ray_in_direction!(queen_square, occupancy, West, first_occupied_square_leading);

    sight
}

fn _king_sight(king_square: Square, blockers: BitBoard) -> BitBoard {
    BitBoard::king_moves(king_square) & !blockers
}
pub(crate) trait BishopSightExt {
    fn bishop_sight(&self, occupancy: BitBoard) -> BitBoard;
}

pub(crate) trait KingSightExt {
    fn king_sight(&self, board: &Board) -> BitBoard;
}

pub(crate) trait KnightSightExt {
    fn knight_sight(&self, board: &Board) -> BitBoard;
}

pub(crate) trait PawnSightExt {
    fn pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard;
    fn white_pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard;
    fn black_pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard;
}
pub(crate) trait QueenSightExt {
    fn queen_sight(&self, occupancy: BitBoard) -> BitBoard;
}

pub(crate) trait RookSightExt {
    fn rook_sight(&self, occupancy: BitBoard) -> BitBoard;
}

pub(crate) trait SliderSightExt: BishopSightExt + QueenSightExt + RookSightExt {}

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

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

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

impl KingSightExt for PlacedPiece {
    fn king_sight(&self, board: &Board) -> BitBoard {
        _king_sight(self.square, board.pieces.all_pieces_of_color(self.color))
    }
}

impl KnightSightExt for PlacedPiece {
    fn knight_sight(&self, board: &Board) -> BitBoard {
        _knight_sight(self.square, board.pieces.all_pieces_of_color(self.color))
    }
}

impl PawnSightExt for PlacedPiece {
    fn pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard {
        match self.color {
            Color::White => self.white_pawn_sight(board, en_passant_square),
            Color::Black => self.black_pawn_sight(board, en_passant_square),
        }
    }

    fn white_pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard {
        let opponent = self.color.other();
        _white_pawn_sight(
            self.square.into(),
            board.pieces.all_pieces(),
            board.pieces.all_pieces_of_color(opponent),
            en_passant_square.into(),
        )
    }

    fn black_pawn_sight(&self, board: &Board, en_passant_square: Option<Square>) -> BitBoard {
        let opponent = self.piece.color.other();
        _black_pawn_sight(
            self.square.into(),
            board.pieces.all_pieces(),
            board.pieces.all_pieces_of_color(opponent),
            en_passant_square.into(),
        )
    }
}

impl BishopSightExt for PlacedPiece {
    fn bishop_sight(&self, occupancy: BitBoard) -> BitBoard {
        _bishop_sight(self.square, occupancy)
    }
}

impl RookSightExt for PlacedPiece {
    fn rook_sight(&self, occupancy: BitBoard) -> BitBoard {
        _rook_sight(self.square, occupancy)
    }
}

impl QueenSightExt for PlacedPiece {
    fn queen_sight(&self, occupancy: BitBoard) -> BitBoard {
        _queen_sight(self.square, occupancy)
    }
}

impl SliderSightExt for PlacedPiece {}

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

        let target_bitboard: BitBoard = target.into();

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

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

        None
    }
}

impl BishopSightExt for Square {
    fn bishop_sight(&self, occupancy: BitBoard) -> BitBoard {
        _bishop_sight(*self, occupancy)
    }
}

impl QueenSightExt for Square {
    fn queen_sight(&self, occupancy: BitBoard) -> BitBoard {
        _queen_sight(*self, occupancy)
    }
}

impl RookSightExt for Square {
    fn rook_sight(&self, occupancy: BitBoard) -> BitBoard {
        _rook_sight(*self, occupancy)
    }
}

#[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!(Shape::Pawn.ray_to_square(Square::F7, Square::E8), None);
        assert_eq!(Shape::Knight.ray_to_square(Square::F6, Square::E8), None);
    }

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

        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 pos = test_position!(White, [
                White Pawn on E5,
                Black Pawn on D5,
            ], 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 = Shape::Bishop.ray_to_square(Square::C5, 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 E1]
        );

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

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

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

            let generated_ray = Shape::Rook.ray_to_square(Square::A6, 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]);
    }
}