chessfriend/board/src/position/position.rs

// Eryn Wells <eryn@erynwells.me>

use super::{flags::Flags, piece_sets::PieceBitBoards, Pieces};
use crate::{
    move_generator::{MoveSet, Moves},
    position::DiagramFormatter,
    r#move::Castle,
    sight::SightExt,
    Move,
};
use chessfriend_bitboard::BitBoard;
use chessfriend_core::{Color, Piece, PlacedPiece, Shape, Square};
use std::{cell::OnceCell, fmt};

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Position {
    color_to_move: Color,
    flags: Flags,
    pieces: PieceBitBoards,
    en_passant_square: Option<Square>,
    sight: [OnceCell<BitBoard>; 2],
    moves: OnceCell<Moves>,

    half_move_counter: u16,
    full_move_number: u16,
}

impl Position {
    pub fn empty() -> Position {
        Position {
            color_to_move: Color::White,
            flags: Default::default(),
            pieces: PieceBitBoards::default(),
            en_passant_square: None,
            sight: [OnceCell::new(), OnceCell::new()],
            moves: OnceCell::new(),
            half_move_counter: 0,
            full_move_number: 1,
        }
    }

    /// Return a starting position.
    pub fn starting() -> Self {
        let black_pieces = [
            BitBoard::new(0b0000000011111111 << 48),
            BitBoard::new(0b0100001000000000 << 48),
            BitBoard::new(0b0010010000000000 << 48),
            BitBoard::new(0b1000000100000000 << 48),
            BitBoard::new(0b0000100000000000 << 48),
            BitBoard::new(0b0001000000000000 << 48),
        ];

        let white_pieces = [
            BitBoard::new(0b1111111100000000),
            BitBoard::new(0b0000000001000010),
            BitBoard::new(0b0000000000100100),
            BitBoard::new(0b0000000010000001),
            BitBoard::new(0b0000000000001000),
            BitBoard::new(0b0000000000010000),
        ];

        Self {
            color_to_move: Color::White,
            flags: Flags::default(),
            pieces: PieceBitBoards::new([white_pieces, black_pieces]),
            en_passant_square: None,
            sight: [OnceCell::new(), OnceCell::new()],
            moves: OnceCell::new(),
            half_move_counter: 0,
            full_move_number: 1,
        }
    }

    pub fn player_to_move(&self) -> Color {
        self.color_to_move
    }

    pub fn move_number(&self) -> u16 {
        self.full_move_number
    }

    pub fn ply_counter(&self) -> u16 {
        self.half_move_counter
    }

    /// Returns true if the player has the right to castle on the given side of
    /// the board.
    ///
    /// The right to castle on a particular side of the board is retained as
    /// long as the player has not moved their king, or the rook on that side of
    /// the board.
    pub(crate) fn player_has_right_to_castle(&self, color: Color, castle: Castle) -> bool {
        self.flags.player_has_right_to_castle(color, castle)
    }

    /// Returns `true` if the player is able to castle on the given side of the board.
    ///
    /// The following requirements must be met:
    ///
    /// 1. The player must still have the right to castle on that side of the
    ///    board. The king and rook involved in the castle must not have moved.
    /// 1. The spaces between the king and rook must be clear
    /// 2. The king must not be in check
    /// 3. In the course of castling on that side, the king must not pass
    ///    through a square that an enemy piece can see
    pub(crate) fn player_can_castle(&self, player: Color, castle: Castle) -> bool {
        if !self.player_has_right_to_castle(player, castle.into()) {
            return false;
        }

        // TODO: Perform a real check that the player can castle.

        true
    }

    /// Return a PlacedPiece representing the rook to use for a castling move.
    pub(crate) fn rook_for_castle(&self, player: Color, castle: Castle) -> Option<PlacedPiece> {
        let square = match (player, castle) {
            (Color::White, Castle::KingSide) => Square::H1,
            (Color::White, Castle::QueenSide) => Square::A1,
            (Color::Black, Castle::KingSide) => Square::H8,
            (Color::Black, Castle::QueenSide) => Square::A8,
        };

        self.piece_on_square(square)
    }

    pub fn moves(&self) -> &Moves {
        self.moves
            .get_or_init(|| Moves::new(self, self.color_to_move))
    }

    /// Return a BitBoard representing the set of squares containing a piece.
    #[inline]
    pub(crate) fn occupied_squares(&self) -> &BitBoard {
        &self.pieces.all_pieces()
    }

    #[inline]
    pub(crate) fn friendly_pieces(&self) -> &BitBoard {
        self.pieces.all_pieces_of_color(self.color_to_move)
    }

    #[inline]
    pub(crate) fn opposing_pieces(&self) -> &BitBoard {
        self.pieces.all_pieces_of_color(self.color_to_move.other())
    }

    pub(crate) fn all_pieces(&self) -> (&BitBoard, &BitBoard) {
        (self.friendly_pieces(), self.opposing_pieces())
    }

    /// Return a BitBoard representing the set of squares containing a piece.
    /// This set is the inverse of `occupied_squares`.
    #[inline]
    pub(crate) fn empty_squares(&self) -> BitBoard {
        !self.occupied_squares()
    }

    pub fn piece_on_square(&self, sq: Square) -> Option<PlacedPiece> {
        for color in Color::iter() {
            for shape in Shape::iter() {
                let piece = Piece::new(color, *shape);
                if self.pieces.bitboard_for_piece(&piece).is_set(sq) {
                    return Some(PlacedPiece::new(piece, sq));
                }
            }
        }

        None
    }

    pub fn pieces(&self, color: Color) -> Pieces {
        Pieces::new(&self, color)
    }

    pub fn en_passant_square(&self) -> Option<Square> {
        self.en_passant_square
    }

    pub(crate) fn sight_of_player(&self, color: Color) -> BitBoard {
        *self.sight[color as usize].get_or_init(|| {
            self.pieces(color).fold(BitBoard::empty(), |acc, pp| {
                acc | pp.sight_in_position(&self)
            })
        })
    }

    pub(crate) fn is_king_in_check(&self) -> bool {
        let sight_of_opposing_player = self.sight_of_player(self.color_to_move.other());
        sight_of_opposing_player.is_set(self.king_square())
    }

    fn king_square(&self) -> Square {
        self.pieces
            .bitboard_for_piece(&Piece::king(self.color_to_move))
            .occupied_squares()
            .next()
            .unwrap()
    }

    pub(crate) fn move_is_legal(&self, mv: Move) -> bool {
        true
    }
}

// crate::position methods
impl Position {
    pub(super) fn new(
        player_to_move: Color,
        flags: Flags,
        pieces: PieceBitBoards,
        en_passant_square: Option<Square>,
        half_move_counter: u16,
        full_move_number: u16,
    ) -> Self {
        Self {
            color_to_move: player_to_move,
            flags,
            en_passant_square,
            pieces,
            sight: [OnceCell::new(), OnceCell::new()],
            moves: OnceCell::new(),
            half_move_counter: 0,
            full_move_number: 1,
        }
    }

    pub(super) fn flags(&self) -> Flags {
        self.flags
    }

    pub(super) fn piece_bitboards(&self) -> &PieceBitBoards {
        &self.pieces
    }
}

// crate methods
impl Position {
    pub(crate) fn bitboard_for_color(&self, color: Color) -> &BitBoard {
        self.pieces.bitboard_for_color(color)
    }

    pub(crate) fn bitboard_for_piece(&self, piece: Piece) -> &BitBoard {
        self.pieces.bitboard_for_piece(&piece)
    }
}

#[cfg(test)]
impl Position {
    pub(crate) fn test_set_en_passant_square(&mut self, square: Square) {
        self.en_passant_square = Some(square);
    }
}

impl Default for Position {
    fn default() -> Self {
        Self::empty()
    }
}

impl fmt::Display for Position {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", DiagramFormatter::new(self))
    }
}

#[cfg(test)]
mod tests {
    use crate::{position, test_position, Castle, Position};
    use chessfriend_core::{piece, Color, Square};

    #[test]
    fn piece_on_square() {
        let pos = test_position![
            Black Bishop on F7,
        ];

        let piece = pos.piece_on_square(Square::F7);
        assert_eq!(piece, Some(piece!(Black Bishop on F7)));
    }

    #[test]
    fn piece_in_starting_position() {
        let pos = Position::starting();
        println!("{pos}");

        assert_eq!(
            pos.piece_on_square(Square::H1),
            Some(piece!(White Rook on H1))
        );
        assert_eq!(
            pos.piece_on_square(Square::A8),
            Some(piece!(Black Rook on A8))
        );
    }

    #[test]
    fn king_is_in_check() {
        let pos = position![
            White King on E1,
            Black Rook on E8,
        ];
        assert!(pos.is_king_in_check());
    }

    #[test]
    fn king_is_not_in_check() {
        let pos = position![
            White King on F1,
            Black Rook on E8,
        ];
        assert!(!pos.is_king_in_check());
    }

    #[test]
    fn rook_for_castle() {
        let pos = position![
            White King on E1,
            White Rook on H1,
            White Rook on A1,
        ];

        assert_eq!(
            pos.rook_for_castle(Color::White, Castle::KingSide),
            Some(piece!(White Rook on H1))
        );
        assert_eq!(
            pos.rook_for_castle(Color::White, Castle::QueenSide),
            Some(piece!(White Rook on A1))
        );
    }
}