chessfriend/moves/src/generators/pawn.rs

// Eryn Wells <eryn@erynwells.me>

//! Implements a move generator for pawns.

use crate::Move;

use super::GeneratedMove;
use chessfriend_bitboard::{bit_scanner::TrailingBitScanner, BitBoard};
use chessfriend_board::Board;
use chessfriend_core::{Color, Direction, Rank, Square};

pub struct PawnMoveGenerator {
    color: Color,
    single_pushes: BitBoard,
    double_pushes: BitBoard,
    left_captures: BitBoard,
    right_captures: BitBoard,
    en_passant: BitBoard,
    iterator: TrailingBitScanner,
    move_type: MoveType,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum MoveType {
    SinglePushes,
    DoublePushes,
    LeftCaptures,
    RightCaptures,
    EnPassant,
}

impl PawnMoveGenerator {
    #[must_use]
    pub fn new(board: &Board, color: Option<Color>) -> Self {
        let color = board.unwrap_color(color);

        let pawns = board.pawns(color);
        let occupied = board.occupancy();
        let empty = !occupied;
        let enemies = board.enemies(color);

        let (single_pushes, double_pushes) = Self::pushes(pawns, color, empty);
        let (left_captures, right_captures) = Self::captures(pawns, color, enemies);
        let en_passant: BitBoard = board.en_passant_target.into();

        Self {
            color,
            single_pushes,
            double_pushes,
            left_captures,
            right_captures,
            en_passant,
            iterator: single_pushes.occupied_squares_trailing(),
            move_type: MoveType::SinglePushes,
        }
    }

    fn pushes(pawns: BitBoard, color: Color, empty: BitBoard) -> (BitBoard, BitBoard) {
        match color {
            Color::White => {
                const THIRD_RANK: BitBoard = BitBoard::rank(Rank::THREE);

                let single_pushes = pawns.shift_north_one() & empty;
                let double_pushes = (single_pushes & THIRD_RANK).shift_north_one() & empty;

                (single_pushes, double_pushes)
            }
            Color::Black => {
                const SIXTH_RANK: BitBoard = BitBoard::rank(Rank::SIX);

                let single_pushes = pawns.shift_south_one() & empty;
                let double_pushes = (single_pushes & SIXTH_RANK).shift_south_one() & empty;

                (single_pushes, double_pushes)
            }
        }
    }

    fn captures(pawns: BitBoard, color: Color, enemies: BitBoard) -> (BitBoard, BitBoard) {
        match color {
            Color::White => {
                let left_captures = pawns.shift_north_west_one() & enemies;
                let right_captures = pawns.shift_north_east_one() & enemies;
                (left_captures, right_captures)
            }
            Color::Black => {
                let left_captures = pawns.shift_south_east_one() & enemies;
                let right_captures = pawns.shift_south_west_one() & enemies;
                (left_captures, right_captures)
            }
        }
    }

    fn calculate_origin_square(&self, target: Square) -> Option<Square> {
        match self.move_type {
            MoveType::SinglePushes => match self.color {
                Color::White => target.neighbor(Direction::South, None),
                Color::Black => target.neighbor(Direction::North, None),
            },
            MoveType::DoublePushes => match self.color {
                Color::White => target.neighbor(Direction::South, Some(2)),
                Color::Black => target.neighbor(Direction::North, Some(2)),
            },
            MoveType::LeftCaptures => match self.color {
                Color::White => target.neighbor(Direction::NorthWest, None),
                Color::Black => target.neighbor(Direction::SouthEast, None),
            },
            MoveType::RightCaptures => match self.color {
                Color::White => target.neighbor(Direction::NorthEast, None),
                Color::Black => target.neighbor(Direction::SouthWest, None),
            },
            MoveType::EnPassant => match self.color {
                Color::White => unimplemented!(),
                Color::Black => unimplemented!(),
            },
        }
    }

    fn next_move_type(&mut self) -> Option<MoveType> {
        let next_move_type = self.move_type.next();

        if let Some(next_move_type) = next_move_type {
            let next_bitboard = match next_move_type {
                MoveType::SinglePushes => self.single_pushes,
                MoveType::DoublePushes => self.double_pushes,
                MoveType::LeftCaptures => self.left_captures,
                MoveType::RightCaptures => self.right_captures,
                MoveType::EnPassant => self.en_passant,
            };

            self.iterator = next_bitboard.occupied_squares_trailing();
            self.move_type = next_move_type;
        }

        next_move_type
    }
}

impl std::iter::Iterator for PawnMoveGenerator {
    type Item = GeneratedMove;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(target) = self.iterator.next() {
            let origin = self
                .calculate_origin_square(target)
                .expect("Bogus origin square");

            match self.move_type {
                MoveType::SinglePushes => Some(GeneratedMove {
                    ply: Move::quiet(origin, target),
                }),
                MoveType::DoublePushes => Some(GeneratedMove {
                    ply: Move::double_push(origin, target),
                }),
                MoveType::LeftCaptures | MoveType::RightCaptures => Some(GeneratedMove {
                    ply: Move::capture(origin, target),
                }),
                MoveType::EnPassant => Some(GeneratedMove {
                    ply: Move::en_passant_capture(origin, target),
                }),
            }
        } else if self.next_move_type().is_some() {
            self.next()
        } else {
            None
        }
    }
}

impl MoveType {
    fn next(self) -> Option<Self> {
        match self {
            MoveType::SinglePushes => Some(MoveType::DoublePushes),
            MoveType::DoublePushes => Some(MoveType::LeftCaptures),
            MoveType::LeftCaptures => Some(MoveType::RightCaptures),
            MoveType::RightCaptures => Some(MoveType::EnPassant),
            MoveType::EnPassant => None,
        }
    }
}