chessfriend/moves/src/generators/king.rs

// Eryn Wells <eryn@erynwells.me>

use crate::{GeneratedMove, Move};
use chessfriend_bitboard::{bit_scanner::TrailingBitScanner, BitBoard};
use chessfriend_board::{Board, CastleParameters};
use chessfriend_core::{Color, Square, Wing};
use std::iter::FusedIterator;

#[must_use]
pub struct KingMoveGenerator {
    kings: Vec<KingIterator>,
    next_kings_index: usize,
    current_king: Option<KingIterator>,
    castle_iterator: CastleIterator,
    friends: BitBoard,
    enemies: BitBoard,
}

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

        let friends = board.friendly_occupancy(color);
        let enemies = board.enemies(color);

        let kings: Vec<KingIterator> = board
            .kings(color)
            .occupied_squares_trailing()
            .map(|king| KingIterator {
                origin: king,
                moves: board
                    .king_sight(king, Some(color))
                    .occupied_squares_trailing(),
            })
            .collect();

        Self {
            kings,
            next_kings_index: 0,
            current_king: None,
            castle_iterator: CastleIterator::new(board, color),
            friends,
            enemies,
        }
    }
}

impl Iterator for KingMoveGenerator {
    type Item = GeneratedMove;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.current_king.is_none() {
                if self.next_kings_index < self.kings.len() {
                    self.current_king = Some(self.kings[self.next_kings_index].clone());
                    self.next_kings_index += 1;
                } else {
                    break;
                }
            }

            if let Some(current_king) = self.current_king.as_mut() {
                if let Some(target) = current_king.next() {
                    let target_bitboard: BitBoard = target.into();

                    let is_targeting_friendly_piece =
                        (target_bitboard & self.friends).is_populated();
                    if is_targeting_friendly_piece {
                        continue;
                    }

                    let is_targeting_enemy_piece = (target_bitboard & self.enemies).is_populated();
                    if is_targeting_enemy_piece {
                        return Some(Move::capture(current_king.origin, target).into());
                    }

                    return Some(Move::quiet(current_king.origin, target).into());
                }

                self.current_king = None;
            }
        }

        self.castle_iterator.next()
    }
}

impl FusedIterator for KingMoveGenerator {}

#[derive(Clone, Debug)]
struct KingIterator {
    origin: Square,
    moves: TrailingBitScanner,
}

impl Iterator for KingIterator {
    type Item = Square;

    fn next(&mut self) -> Option<Self::Item> {
        self.moves.next()
    }
}

impl FusedIterator for KingIterator {}

#[derive(Clone, Debug)]
struct CastleIterator {
    kingside: Option<&'static CastleParameters>,
    queenside: Option<&'static CastleParameters>,
}

impl CastleIterator {
    fn new(board: &Board, color: Color) -> Self {
        let kingside = board.color_can_castle(Wing::KingSide, Some(color)).ok();
        let queenside = board.color_can_castle(Wing::QueenSide, Some(color)).ok();

        Self {
            kingside,
            queenside,
        }
    }
}

impl Iterator for CastleIterator {
    type Item = GeneratedMove;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(_parameters) = self.kingside.take() {
            return Some(Move::castle(Wing::KingSide).into());
        }

        if let Some(_parameters) = self.queenside.take() {
            return Some(Move::castle(Wing::QueenSide).into());
        }

        None
    }
}

impl FusedIterator for CastleIterator {}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{assert_move_list, ply};
    use chessfriend_board::test_board;

    #[test]
    fn white_king_center_square_ai_claude() {
        let board = test_board!(White King on E4);
        assert_move_list!(
            KingMoveGenerator::new(&board, None),
            [
                // All 8 adjacent squares
                ply!(E4 - D3), // Southwest
                ply!(E4 - D4), // West
                ply!(E4 - D5), // Northwest
                ply!(E4 - E3), // South
                ply!(E4 - E5), // North
                ply!(E4 - F3), // Southeast
                ply!(E4 - F4), // East
                ply!(E4 - F5), // Northeast
            ]
        );
    }

    #[test]
    fn white_king_corner_square_ai_claude() {
        let board = test_board!(White King on A1);
        assert_move_list!(
            KingMoveGenerator::new(&board, None),
            [
                // Only 3 possible moves from corner
                ply!(A1 - A2), // North
                ply!(A1 - B1), // East
                ply!(A1 - B2), // Northeast
            ]
        );
    }

    #[test]
    fn white_king_edge_square_ai_claude() {
        let board = test_board!(White King on E1);
        assert_move_list!(
            KingMoveGenerator::new(&board, None),
            [
                // 5 possible moves from edge
                ply!(E1 - D1), // West
                ply!(E1 - D2), // Northwest
                ply!(E1 - E2), // North
                ply!(E1 - F1), // East
                ply!(E1 - F2), // Northeast
            ]
        );
    }

    #[test]
    fn white_king_with_captures_ai_claude() {
        let board = test_board!(
            White King on D4,
            Black Pawn on C5,    // Can capture
            Black Knight on E5,  // Can capture
            Black Bishop on D3   // Can capture
        );
        assert_move_list!(
            KingMoveGenerator::new(&board, None),
            [
                // Regular moves
                ply!(D4 - C3),
                ply!(D4 - C4),
                ply!(D4 - D5),
                ply!(D4 - E3),
                ply!(D4 - E4),
                // Captures
                ply!(D4 x C5), // Capture pawn
                ply!(D4 x E5), // Capture knight
                ply!(D4 x D3), // Capture bishop
            ]
        );
    }

    #[test]
    fn white_king_blocked_by_friendly_pieces_ai_claude() {
        let board = test_board!(
            White King on D4,
            White Pawn on C4,    // Blocks west
            White Knight on D5,  // Blocks north
            White Bishop on E3   // Blocks southeast
        );
        assert_move_list!(
            KingMoveGenerator::new(&board, None),
            [
                ply!(D4 - C3),
                ply!(D4 - C5),
                ply!(D4 - D3),
                ply!(D4 - E4),
                ply!(D4 - E5),
                // Cannot move to C4, D5, E3 (friendly pieces)
            ]
        );
    }

    #[test]
    fn white_king_castling_kingside_ai_claude() {
        let board = test_board!(
            White King on E1,
            White Rook on H1
            // Assuming squares F1, G1 are empty and king/rook haven't moved
        );
        assert_move_list!(
            KingMoveGenerator::new(&board, None),
            [
                // Regular king moves
                ply!(E1 - D1),
                ply!(E1 - D2),
                ply!(E1 - E2),
                ply!(E1 - F1),
                ply!(E1 - F2),
                ply!(0 - 0),
            ]
        );
    }

    #[test]
    fn white_king_castling_queenside_ai_claude() {
        let board = test_board!(
            White King on E1,
            White Rook on A1
            // Assuming squares B1, C1, D1 are empty and king/rook haven't moved
        );
        assert_move_list!(
            KingMoveGenerator::new(&board, None),
            [
                // Regular king moves
                ply!(E1 - D1),
                ply!(E1 - D2),
                ply!(E1 - E2),
                ply!(E1 - F1),
                ply!(E1 - F2),
                ply!(0 - 0 - 0),
            ]
        );
    }

    #[test]
    fn white_king_no_castling_through_check_ai_claude() {
        let board = test_board!(
            White King on E1,
            White Rook on H1,
            Black Rook on F8    // Attacks F1, preventing kingside castling
        );
        assert_move_list!(
            KingMoveGenerator::new(&board, None),
            [
                ply!(E1 - D1),
                ply!(E1 - D2),
                ply!(E1 - E2),
                ply!(E1 - F1),
                ply!(E1 - F2),
                // No castling moves - cannot castle through check
            ]
        );
    }

    #[test]
    fn white_king_castling_blocked_by_pieces_ai_claude() {
        let board = test_board!(
            White King on E1,
            White Rook on H1,
            White Knight on G1, // Blocks kingside castling
        );
        assert_move_list!(
            KingMoveGenerator::new(&board, None),
            [
                ply!(E1 - D1),
                ply!(E1 - D2),
                ply!(E1 - E2),
                ply!(E1 - F1),
                ply!(E1 - F2),
                // No castling - path blocked by knight
            ]
        );
    }

    #[test]
    fn black_king_movement_ai_claude() {
        let board = test_board!(Black King on E8);
        assert_move_list!(
            KingMoveGenerator::new(&board, Some(Color::Black)),
            [
                ply!(E8 - D7),
                ply!(E8 - D8),
                ply!(E8 - E7),
                ply!(E8 - F7),
                ply!(E8 - F8),
            ]
        );
    }

    #[test]
    fn white_king_surrounded_by_enemies_ai_claude() {
        let board = test_board!(
            White King on E4,
            Black Pawn on D3,
            Black Pawn on D4,
            Black Pawn on D5,
            Black Pawn on E3,
            Black Pawn on E5,
            Black Pawn on F3,
            Black Pawn on F4,
            Black Pawn on F5
        );
        assert_move_list!(
            KingMoveGenerator::new(&board, None),
            [
                // Can capture all surrounding enemy pieces
                ply!(E4 x D3),
                ply!(E4 x D4),
                ply!(E4 x D5),
                ply!(E4 x E3),
                ply!(E4 x E5),
                ply!(E4 x F3),
                ply!(E4 x F4),
                ply!(E4 x F5),
            ]
        );
    }
}