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.
//!
//! These functions use some common terms to describe arguments.
//!
//! occupancy
//! : The set of occupied squares on the board.
//!
//! vacancy
//! : The set of empty squares on the board, `!occpuancy`.
//!
//! blockers
//! : The set of squares occupied by friendly pieces that block moves to that
//! square and beyond.

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

pub trait Sight {
    fn sight(&self, square: Square, board: &Board) -> BitBoard;
}

impl Sight for Piece {
    fn sight(&self, square: Square, board: &Board) -> BitBoard {
        let occupancy = board.occupancy();
        let info = SightInfo {
            square,
            occupancy,
            friendly_occupancy: board.friendly_occupancy(self.color),
        };

        match self.shape {
            Shape::Pawn => {
                let en_passant_square: BitBoard = board.en_passant_target.into();
                match self.color {
                    Color::White => white_pawn_sight(&info, en_passant_square),
                    Color::Black => black_pawn_sight(&info, en_passant_square),
                }
            }
            Shape::Knight => knight_sight(&info),
            Shape::Bishop => bishop_sight(&info),
            Shape::Rook => rook_sight(&info),
            Shape::Queen => queen_sight(&info),
            Shape::King => king_sight(&info),
        }
    }
}

struct SightInfo {
    square: Square,
    occupancy: BitBoard,
    friendly_occupancy: BitBoard,
}

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(info: &SightInfo, en_passant_square: BitBoard) -> BitBoard {
    let possible_squares = !info.friendly_occupancy | en_passant_square;

    let pawn: BitBoard = info.square.into();
    let pawn = pawn.shift_north_west_one() | pawn.shift_north_east_one();

    pawn & possible_squares
}

fn black_pawn_sight(info: &SightInfo, en_passant_square: BitBoard) -> BitBoard {
    let possible_squares = !info.friendly_occupancy | en_passant_square;

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

    pawn & possible_squares
}

fn knight_sight(info: &SightInfo) -> BitBoard {
    BitBoard::knight_moves(info.square)
}

fn bishop_sight(info: &SightInfo) -> BitBoard {
    let bishop = info.square;
    let occupancy = info.occupancy;

    #[rustfmt::skip]
    let sight = ray_in_direction!(bishop, occupancy, NorthEast, first_occupied_square_trailing)
        | ray_in_direction!(bishop, occupancy, SouthEast, first_occupied_square_leading)
        | ray_in_direction!(bishop, occupancy, SouthWest, first_occupied_square_leading)
        | ray_in_direction!(bishop, occupancy, NorthWest, first_occupied_square_trailing);

    sight
}

fn rook_sight(info: &SightInfo) -> BitBoard {
    let rook = info.square;
    let occupancy = info.occupancy;

    #[rustfmt::skip]
    let sight = ray_in_direction!(rook, occupancy, North, first_occupied_square_trailing)
        | ray_in_direction!(rook, occupancy, East, first_occupied_square_trailing)
        | ray_in_direction!(rook, occupancy, South, first_occupied_square_leading)
        | ray_in_direction!(rook, occupancy, West, first_occupied_square_leading);

    sight
}

fn queen_sight(info: &SightInfo) -> BitBoard {
    let queen = info.square;
    let occupancy = info.occupancy;

    #[rustfmt::skip]
    let sight = ray_in_direction!(queen, occupancy, NorthWest, first_occupied_square_trailing)
        | ray_in_direction!(queen, occupancy, North, first_occupied_square_trailing)
        | ray_in_direction!(queen, occupancy, NorthEast, first_occupied_square_trailing)
        | ray_in_direction!(queen, occupancy, East, first_occupied_square_trailing)
        | ray_in_direction!(queen, occupancy, SouthEast, first_occupied_square_leading)
        | ray_in_direction!(queen, occupancy, South, first_occupied_square_leading)
        | ray_in_direction!(queen, occupancy, SouthWest, first_occupied_square_leading)
        | ray_in_direction!(queen, occupancy, West, first_occupied_square_leading);

    sight
}

fn king_sight(info: &SightInfo) -> BitBoard {
    BitBoard::king_moves(info.square)
}

#[cfg(test)]
mod tests {
    use chessfriend_bitboard::bitboard;
    use chessfriend_core::piece;

    macro_rules! sight_test {
        ($test_name:ident, $position:expr, $piece:expr, $square:expr, $bitboard:expr) => {
            #[test]
            fn $test_name() {
                use chessfriend_core::Square;
                use $crate::sight::Sight;

                let pos = $position;
                let piece = $piece;
                let sight = piece.sight($square, &pos.board);

                assert_eq!(sight, $bitboard);
            }
        };
        ($test_name:ident, $piece:expr, $square:expr, $bitboard:expr) => {
            sight_test! {$test_name, $crate::Position::empty(), $piece, $square, $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::{sight::Sight, test_position};
        use chessfriend_bitboard::{bitboard, BitBoard};
        use chessfriend_core::{piece, Square};

        sight_test!(e4_pawn, piece!(White Pawn), Square::E4, bitboard![D5 F5]);

        sight_test!(
            e4_pawn_one_blocker,
            test_position![
                White Bishop on D5,
                White Pawn on E4,
            ],
            piece!(White Pawn),
            Square::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);
            let sight = piece.sight(Square::E4, &pos.board);

            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);
            let sight = piece.sight(Square::E4, &pos.board);

            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);
            let sight = piece.sight(Square::E5, &pos.board);

            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),
            Square::F6,
            bitboard![H7 G8 E8 D7 D5 E4 G4 H5]
        );
    }

    mod bishop {
        use super::*;

        sight_test!(
            c2_bishop,
            piece!(Black Bishop),
            Square::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),
            Square::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 E2,
                Black King on E7,
            ],
            piece!(White Rook),
            Square::E4,
            bitboard![A4 B4 C4 D4 F4 G4 H4 E2 E3 E5 E6 E7]
        );

        // #[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),
            Square::E1,
            bitboard![D1 D2 E2 F2 F1]
        );
    }
}