chessfriend/board/src/move_generator/pawn.rs

// Eryn Wells <eryn@erynwells.me>

use crate::{
    piece::{Color, Piece, Shape},
    BitBoard, Move, Position,
};

enum MoveList {
    Quiet = 0,
    Promotions = 1,
    Captures = 2,
}

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

struct MoveGenerationParameters {
    starting_rank: BitBoard,
    promotion_rank: BitBoard,
    push_shift: fn(BitBoard) -> BitBoard,
    left_capture_shift: fn(BitBoard) -> BitBoard,
    right_capture_shift: fn(BitBoard) -> BitBoard,
}

pub(super) struct PawnMoveGenerator<'pos> {
    color: Color,
    position: &'pos Position,

    did_populate_move_lists: bool,

    pushes: BitBoard,
    attacks: BitBoard,

    move_lists: [Vec<Move>; 3],
    move_iterator: MoveIterator,
}

impl<'pos> PawnMoveGenerator<'pos> {
    pub(super) fn new(position: &Position, color: Color) -> PawnMoveGenerator {
        PawnMoveGenerator {
            position,
            color,
            did_populate_move_lists: false,
            pushes: BitBoard::empty(),
            attacks: BitBoard::empty(),
            move_lists: [Vec::new(), Vec::new(), Vec::new()],
            move_iterator: MoveIterator(0, 0),
        }
    }

    pub(super) fn iter(&self) -> impl Iterator<Item = &Move> + '_ {
        self.move_lists.iter().flatten()
    }

    fn generate_moves(&mut self) {
        self.generate_move_bitboards();
        self.populate_move_lists();
    }

    fn move_generation_parameters(&self) -> MoveGenerationParameters {
        match self.color {
            Color::White => MoveGenerationParameters {
                starting_rank: BitBoard::rank(1),
                promotion_rank: BitBoard::rank(7),
                push_shift: BitBoard::shift_north_one,
                left_capture_shift: BitBoard::shift_north_west_one,
                right_capture_shift: BitBoard::shift_north_east_one,
            },
            Color::Black => MoveGenerationParameters {
                starting_rank: BitBoard::rank(6),
                promotion_rank: BitBoard::rank(0),
                push_shift: BitBoard::shift_south_one,
                left_capture_shift: BitBoard::shift_south_east_one,
                right_capture_shift: BitBoard::shift_south_west_one,
            },
        }
    }

    #[inline]
    fn quiet_move_list(&mut self) -> &mut Vec<Move> {
        &mut self.move_lists[MoveList::Quiet as usize]
    }

    #[inline]
    fn promotion_move_list(&mut self) -> &mut Vec<Move> {
        &mut self.move_lists[MoveList::Promotions as usize]
    }

    #[inline]
    fn capture_move_list(&mut self) -> &mut Vec<Move> {
        &mut self.move_lists[MoveList::Captures as usize]
    }
}

impl<'pos> PawnMoveGenerator<'pos> {
    fn generate_move_bitboards(&mut self) {
        let parameters = self.move_generation_parameters();
        self.generate_pushes_bitboard(&parameters);
        self.generate_attacks_bitboard(&parameters);
    }

    fn generate_pushes_bitboard(&mut self, parameters: &MoveGenerationParameters) {
        let empty_squares = self.position.empty_squares();
        let bb = self.position.bitboard_for_piece(Piece::pawn(self.color));

        let legal_1square_pushes = (parameters.push_shift)(bb) & empty_squares;
        let legal_2square_pushes =
            (parameters.push_shift)(legal_1square_pushes & BitBoard::rank(2)) & empty_squares;

        self.pushes = legal_1square_pushes | legal_2square_pushes;
    }

    fn generate_attacks_bitboard(&mut self, parameters: &MoveGenerationParameters) {
        let opponent_pieces = self.position.bitboard_for_color(self.color.other());
        let bb = self.position.bitboard_for_piece(Piece::pawn(Color::White));

        self.attacks = ((parameters.left_capture_shift)(bb) | (parameters.right_capture_shift)(bb))
            & opponent_pieces;

        #[allow(unused_variables)]
        if let Some(en_passant) = self.en_passant() {
            // TODO: Add en passant move to the attacks bitboard.
        }
    }
}

impl<'pos> PawnMoveGenerator<'pos> {
    fn populate_move_lists(&mut self) {
        let parameters = self.move_generation_parameters();

        self._populate_move_lists(&parameters);
        self.did_populate_move_lists = true;
    }

    fn _populate_move_lists(&mut self, parameters: &MoveGenerationParameters) {
        let piece = Piece::pawn(self.color);

        let bb = self.position.bitboard_for_piece(piece);
        if bb.is_empty() {
            return;
        }

        let empty_squares = self.position.empty_squares();
        let black_pieces = self.position.bitboard_for_color(self.color.other());

        for from_sq in bb.occupied_squares() {
            let pawn: BitBoard = from_sq.into();

            let push = (parameters.push_shift)(pawn);
            if !(push & empty_squares).is_empty() {
                let to_sq = push.occupied_squares().next().unwrap();

                let r#move = Move::new(piece, from_sq, to_sq);
                if !(push & parameters.promotion_rank).is_empty() {
                    for shape in Shape::promotable() {
                        self.promotion_move_list()
                            .push(r#move.clone().promoting_to(*shape));
                    }
                } else {
                    self.quiet_move_list().push(r#move);
                }

                if !(pawn & parameters.starting_rank).is_empty() {
                    let push = (parameters.push_shift)(push);
                    if !(push & empty_squares).is_empty() {
                        let to_sq = push.occupied_squares().next().unwrap();
                        self.quiet_move_list()
                            .push(Move::new(piece, from_sq, to_sq));
                    }
                }
            }

            for attack in [
                (parameters.left_capture_shift)(pawn),
                (parameters.right_capture_shift)(pawn),
            ] {
                if !(attack & black_pieces).is_empty() {
                    let to_sq = attack.occupied_squares().next().unwrap();
                    let captured_piece = self.position.piece_on_square(to_sq).unwrap();

                    let r#move = Move::new(piece, from_sq, to_sq).capturing(captured_piece);
                    if !(attack & parameters.promotion_rank).is_empty() {
                        for shape in Shape::promotable() {
                            self.capture_move_list()
                                .push(r#move.clone().promoting_to(*shape));
                        }
                    } else {
                        self.capture_move_list().push(r#move);
                    }
                }
            }

            if let Some(en_passant) = self.en_passant() {
                self.capture_move_list().push(en_passant);
            }
        }
    }

    fn en_passant(&self) -> Option<Move> {
        // TODO: En passant. I think the way to do this is to have the position mark
        // an en passant square when the conditions are correct, i.e. when the
        // opposing player has pushed a pawn two squares off the initial rank, and
        // then check in these routines if a pawn of this color attacks that square.

        None
    }
}

impl<'pos> Iterator for PawnMoveGenerator<'pos> {
    type Item = Move;

    fn next(&mut self) -> Option<Self::Item> {
        if !self.did_populate_move_lists {
            self.generate_moves();
        }

        let iter = &mut self.move_iterator;

        // Find the next non-empty list.
        loop {
            if iter.0 >= self.move_lists.len() || !self.move_lists[iter.0].is_empty() {
                break;
            }

            iter.0 += 1;
        }

        if let Some(move_list) = self.move_lists.get(iter.0) {
            let next_move = move_list[iter.1].clone();

            iter.1 += 1;
            if iter.1 >= move_list.len() {
                // Increment the list index here. On the next iteration, find the next non-empty list.
                iter.0 += 1;
                iter.1 = 0;
            }

            Some(next_move)
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::piece::PlacedPiece;
    use crate::position::DiagramFormatter;
    use crate::{Position, Square};
    use std::collections::HashSet;

    #[test]
    fn one_2square_push() {
        let mut pos = Position::empty();

        pos.place_piece(Piece::pawn(Color::White), Square::E2)
            .expect("Failed to place pawn on e2");

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

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

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

        assert_eq!(generated_moves, expected_moves);
    }

    #[test]
    fn one_1square_push() {
        let mut pos = Position::empty();

        pos.place_piece(Piece::pawn(Color::White), Square::E3)
            .expect("Failed to place pawn on e3");

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

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

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

        assert_eq!(generated_moves, expected_moves);
    }

    #[test]
    fn one_obstructed_2square_push() {
        let mut pos = Position::empty();

        pos.place_piece(Piece::pawn(Color::White), Square::E2)
            .expect("Failed to place pawn on e2");
        pos.place_piece(Piece::knight(Color::White), Square::E4)
            .expect("Failed to place knight on e4");

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

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

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

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

        assert_eq!(generated_moves, expected_moves);
    }

    #[test]
    fn one_obstructed_1square_push() {
        let mut pos = Position::empty();

        pos.place_piece(Piece::pawn(Color::White), Square::E2)
            .expect("Failed to place pawn on e2");
        pos.place_piece(Piece::knight(Color::White), Square::E3)
            .expect("Failed to place knight on e4");

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

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

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

        assert_eq!(generated_moves, HashSet::new());
    }

    #[test]
    fn one_attack() {
        let mut pos = Position::empty();
        pos.place_piece(Piece::pawn(Color::White), Square::E4)
            .expect("Failed to place pawn on e4");
        pos.place_piece(Piece::bishop(Color::White), Square::E5)
            .expect("Failed to place pawn on e4");
        pos.place_piece(Piece::knight(Color::Black), Square::D5)
            .expect("Failed to place knight on d5");

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

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

        let expected_moves = HashSet::from_iter(
            [Move::new(Piece::pawn(Color::White), Square::E4, Square::D5)
                .capturing(PlacedPiece::new(Piece::knight(Color::Black), Square::D5))]
            .into_iter(),
        );

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

        assert_eq!(generated_moves, expected_moves);
    }

    #[test]
    fn one_double_attack() {
        let mut pos = Position::empty();
        pos.place_piece(Piece::pawn(Color::White), Square::E4)
            .expect("Failed to place pawn on e4");
        pos.place_piece(Piece::bishop(Color::White), Square::E5)
            .expect("Failed to place pawn on e4");
        pos.place_piece(Piece::knight(Color::Black), Square::D5)
            .expect("Failed to place knight on d5");
        pos.place_piece(Piece::queen(Color::Black), Square::F5)
            .expect("Failed to place knight on f5");

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

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

        let expected_moves = HashSet::from_iter(
            [
                Move::new(Piece::pawn(Color::White), Square::E4, Square::D5)
                    .capturing(PlacedPiece::new(Piece::knight(Color::Black), Square::D5)),
                Move::new(Piece::pawn(Color::White), Square::E4, Square::F5)
                    .capturing(PlacedPiece::new(Piece::queen(Color::Black), Square::F5)),
            ]
            .into_iter(),
        );

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

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