chessfriend/position/src/move_generator/pawn.rs

// Eryn Wells <eryn@erynwells.me>

use super::{move_generator_declaration, MoveGeneratorInternal, MoveSet};
use crate::Position;
use chessfriend_bitboard::BitBoard;
use chessfriend_core::{Piece, PlacedPiece, Rank, Shape, Square};

#[derive(Debug)]
struct MoveIterator(usize, usize);

move_generator_declaration!(PawnMoveGenerator);

impl MoveGeneratorInternal for PawnMoveGenerator {
    fn shape() -> Shape {
        Shape::Pawn
    }

    fn move_set_for_piece(
        position: &Position,
        placed_piece: PlacedPiece,
        capture_mask: BitBoard,
        push_mask: BitBoard,
    ) -> MoveSet {
        let capture_moves = Self::attacks(position, placed_piece) & capture_mask;
        let quiet_moves = Self::pushes(position, placed_piece) & push_mask;

        MoveSet::new(placed_piece)
            .quiet_moves(quiet_moves)
            .capture_moves(capture_moves)
    }
}

impl PawnMoveGenerator {
    fn pushes(position: &Position, piece: PlacedPiece) -> BitBoard {
        let square = piece.square();
        let bitboard: BitBoard = square.into();

        let starting_rank = Rank::PAWN_STARTING_RANKS[piece.color() as usize];

        let empty_squares = position.empty_squares();
        let mut moves = bitboard.shift_north_one() & empty_squares;
        if !(bitboard & BitBoard::rank(starting_rank.as_index())).is_empty() {
            moves |= moves.shift_north_one() & empty_squares;
        }

        moves
    }

    fn attacks(position: &Position, piece: PlacedPiece) -> BitBoard {
        let color = piece.color();

        let opponent_pieces = position.bitboard_for_color(color.other());
        let en_passant_bitboard = match position.en_passant_square() {
            Some(square) => <Square as Into<BitBoard>>::into(square),
            None => BitBoard::empty(),
        };

        BitBoard::pawn_attacks(piece.square(), color) & (opponent_pieces | en_passant_bitboard)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{assert_move_list, position::DiagramFormatter, test_position, Move, MoveBuilder};
    use chessfriend_core::{piece, Color, Square};
    use std::collections::HashSet;

    #[test]
    fn one_2square_push() {
        let pos = test_position![White Pawn on E2];

        let generator = PawnMoveGenerator::new(&pos, Color::White, BitBoard::FULL, BitBoard::FULL);

        let expected_moves = HashSet::from_iter([
            MoveBuilder::new(piece!(White Pawn), Square::E2, Square::E3).build(),
            MoveBuilder::new(piece!(White Pawn), Square::E2, Square::E4).build(),
        ]);

        let generated_moves: HashSet<Move> = generator.iter().collect();

        assert_eq!(generated_moves, expected_moves);
    }

    #[test]
    fn one_1square_push() {
        let pos = test_position![White Pawn on E3];

        let generator = PawnMoveGenerator::new(&pos, Color::White, BitBoard::FULL, BitBoard::FULL);

        let expected_moves = HashSet::from_iter([MoveBuilder::new(
            Piece::pawn(Color::White),
            Square::E3,
            Square::E4,
        )
        .build()]);

        let generated_moves: HashSet<Move> = generator.iter().collect();

        assert_move_list!(generated_moves, expected_moves, pos);
    }

    #[test]
    fn one_obstructed_2square_push() {
        let pos = test_position![
            White Pawn on E2,
            White Knight on E4,
        ];

        println!("{}", DiagramFormatter::new(&pos));

        let generator = PawnMoveGenerator::new(&pos, Color::White, BitBoard::FULL, BitBoard::FULL);

        let expected_moves = HashSet::from_iter([MoveBuilder::new(
            Piece::pawn(Color::White),
            Square::E2,
            Square::E3,
        )
        .build()]);

        let generated_moves: HashSet<Move> = generator.iter().collect();

        assert_move_list!(generated_moves, expected_moves, pos);
    }

    #[test]
    fn one_obstructed_1square_push() {
        let pos = test_position![
            White Pawn on E2,
            White Knight on E3,
        ];

        let generator = PawnMoveGenerator::new(&pos, Color::White, BitBoard::FULL, BitBoard::FULL);

        let generated_moves: HashSet<Move> = generator.iter().collect();
        let expected_moves: HashSet<Move> = HashSet::new();

        assert_move_list!(generated_moves, expected_moves, pos);
    }

    #[test]
    fn one_attack() {
        let pos = test_position![
            White Pawn on E4,
            White Bishop on E5,
            Black Knight on D5,
        ];

        let generator = PawnMoveGenerator::new(&pos, Color::White, BitBoard::FULL, BitBoard::FULL);

        let expected_moves =
            HashSet::from_iter(
                [MoveBuilder::new(piece!(White Pawn), Square::E4, Square::D5)
                    .capturing(piece!(Black Knight on D5))
                    .build()],
            );

        let generated_moves: HashSet<Move> = generator.iter().collect();

        assert_eq!(generated_moves, expected_moves);
    }

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

        let generator = PawnMoveGenerator::new(&pos, Color::White, BitBoard::FULL, BitBoard::FULL);

        let expected_moves = HashSet::from_iter([
            MoveBuilder::new(piece!(White Pawn), Square::E4, Square::D5)
                .capturing(piece!(Black Knight on D5))
                .build(),
            MoveBuilder::new(piece!(White Pawn), Square::E4, Square::F5)
                .capturing(piece!(Black Queen on F5))
                .build(),
        ]);

        let generated_moves: HashSet<Move> = generator.iter().collect();

        assert_eq!(
            generated_moves, expected_moves,
            "generated: {:#?}\nexpected: {:#?}",
            generated_moves, expected_moves
        );
    }
}