chessfriend/position/src/move_generator/move_set.rs

// Eryn Wells <eryn@erynwells.me>

use chessfriend_bitboard::BitBoard;
use chessfriend_core::{Color, Piece, PlacedPiece, Shape, Square};
use chessfriend_moves::{Builder as MoveBuilder, Castle, Move};

/// A set of bitboards defining the moves for a single piece on the board.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
struct BitBoardSet {
    quiet: BitBoard,
    captures: BitBoard,
}

/// A set of moves for a single piece on the board.
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct MoveSet {
    piece: PlacedPiece,
    bitboards: BitBoardSet,
    special: u8,
}

impl MoveSet {
    pub(super) fn new(piece: PlacedPiece) -> MoveSet {
        MoveSet {
            piece,
            bitboards: BitBoardSet::default(),
            special: 0,
        }
    }

    pub(crate) fn can_move_to_square(&self, to: Square) -> bool {
        self.bitboard().is_set(to)
    }

    pub(crate) fn can_castle(&self, castle: Castle) -> bool {
        match castle {
            Castle::KingSide => (self.special & 0b1) != 0,
            Castle::QueenSide => (self.special & 0b10) != 0,
        }
    }

    pub(super) fn quiet_moves(mut self, bitboard: BitBoard) -> MoveSet {
        self.bitboards.quiet = bitboard;
        self
    }

    pub(super) fn capture_moves(mut self, bitboard: BitBoard) -> MoveSet {
        self.bitboards.captures = bitboard;
        self
    }

    pub(super) fn kingside_castle(&mut self) -> &mut MoveSet {
        self.special |= 0b1;
        self
    }

    pub(super) fn queenside_castle(&mut self) -> &mut MoveSet {
        self.special |= 0b10;
        self
    }

    /// Return a BitBoard representing all possible moves.
    pub(super) fn bitboard(&self) -> BitBoard {
        self.bitboards.captures | self.bitboards.quiet
    }

    pub(crate) fn moves(&self) -> impl Iterator<Item = Move> + '_ {
        let piece = self.piece.piece();
        let from_square = self.piece.square();

        self.bitboards
            .quiet
            .occupied_squares()
            .map(move |to_square| MoveBuilder::new(*piece, from_square, to_square).build())
            .chain(
                self.bitboards
                    .captures
                    .occupied_squares()
                    .map(move |to_square| {
                        MoveBuilder::new(*piece, from_square, to_square)
                            .capturing(PlacedPiece::new(
                                Piece::new(Color::White, Shape::Pawn),
                                to_square,
                            ))
                            .build()
                    }),
            )
            .chain(
                if (self.special & 0b1) != 0 {
                    Some(())
                } else {
                    None
                }
                .map(|()| {
                    MoveBuilder::new(
                        *piece,
                        from_square,
                        Castle::KingSide.target_squares(piece.color()).king,
                    )
                    .castle(Castle::KingSide)
                    .build()
                }),
            )
            .chain(
                if (self.special & 0b10) != 0 {
                    Some(())
                } else {
                    None
                }
                .map(|()| {
                    MoveBuilder::new(
                        *piece,
                        from_square,
                        Castle::QueenSide.target_squares(piece.color()).king,
                    )
                    .castle(Castle::QueenSide)
                    .build()
                }),
            )
    }
}