chessfriend/moves/src/unmake_move.rs

// Eryn Wells <eryn@erynwells.me>

use crate::MoveRecord;
use chessfriend_board::{Board, BoardProvider, PlacePieceError, PlacePieceStrategy};
use chessfriend_core::{Piece, Square};
use thiserror::Error;

pub type UnmakeMoveResult = Result<(), UnmakeMoveError>;

#[derive(Clone, Debug, Error, Eq, PartialEq)]
pub enum UnmakeMoveError {
    #[error("no move to unmake")]
    NoMove,

    #[error("no piece on {0}")]
    NoPiece(Square),

    #[error("no capture square")]
    NoCaptureSquare,

    #[error("no captured piece to unmake capture move")]
    NoCapturedPiece,

    #[error("{0}")]
    PlacePieceError(#[from] PlacePieceError),
}

pub trait UnmakeMove {
    /// Unmake the given move. Unmaking a move that wasn't the most recent one
    /// made will likely cause strange results.
    ///
    /// ## To-Do
    ///
    /// Some implementations I've seen in other engines take a move to unmake. I
    /// don't understand why they do this because I don't think it makes sense
    /// to unmake any move other than the last move made. I need to do some
    /// research on this to understand if/when passing a move might be useful or
    /// necessary.
    ///
    /// ## Errors
    ///
    /// Returns one of [`UnmakeMoveError`] indicating why the move cannot be
    /// unmade.
    ///
    fn unmake_move(&mut self, record: &MoveRecord) -> UnmakeMoveResult;
}

trait UnmakeMoveInternal {
    fn unmake_quiet_move(&mut self, record: &MoveRecord) -> UnmakeMoveResult;
    fn unmake_capture_move(&mut self, record: &MoveRecord) -> UnmakeMoveResult;
    fn unmake_promotion_move(&mut self, record: &MoveRecord) -> UnmakeMoveResult;
    fn unmake_castle_move(&mut self, record: &MoveRecord) -> UnmakeMoveResult;
}

impl<T: BoardProvider> UnmakeMove for T {
    fn unmake_move(&mut self, record: &MoveRecord) -> UnmakeMoveResult {
        let ply = record.ply;

        if ply.is_quiet() || ply.is_double_push() {
            self.unmake_quiet_move(record)?;
        } else if ply.is_capture() || ply.is_en_passant() {
            self.unmake_capture_move(record)?;
        } else if ply.is_promotion() {
            self.unmake_promotion_move(record)?;
        } else if ply.is_castle() {
            self.unmake_castle_move(record)?;
        } else {
            unreachable!();
        }

        let board = self.board_mut();
        board.set_active_color(record.color);
        board.set_en_passant_target_option(record.en_passant_target);
        board.set_castling_rights(record.castling_rights);
        board.half_move_clock = record.half_move_clock;

        Ok(())
    }
}

impl<T: BoardProvider> UnmakeMoveInternal for T {
    fn unmake_quiet_move(&mut self, record: &MoveRecord) -> UnmakeMoveResult {
        let board = self.board_mut();

        let ply = record.ply;

        let target = ply.target_square();
        let piece = board
            .get_piece(target)
            .ok_or(UnmakeMoveError::NoPiece(target))?;

        let origin = ply.origin_square();
        board.place_piece(piece, origin, PlacePieceStrategy::PreserveExisting)?;

        board.remove_piece(target);

        Ok(())
    }

    fn unmake_capture_move(&mut self, record: &MoveRecord) -> UnmakeMoveResult {
        let board = self.board_mut();

        let ply = record.ply;

        let target = ply.target_square();

        let mut piece = board
            .get_piece(target)
            .ok_or(UnmakeMoveError::NoPiece(target))?;
        if ply.is_promotion() {
            piece = Piece::pawn(piece.color);
        }

        let origin = ply.origin_square();
        board.place_piece(piece, origin, PlacePieceStrategy::PreserveExisting)?;

        let capture_square = ply
            .capture_square()
            .ok_or(UnmakeMoveError::NoCaptureSquare)?;
        let captured_piece = record
            .captured_piece
            .ok_or(UnmakeMoveError::NoCapturedPiece)?;

        board.remove_piece(target);

        board.place_piece(
            captured_piece,
            capture_square,
            PlacePieceStrategy::PreserveExisting,
        )?;

        Ok(())
    }

    fn unmake_promotion_move(&mut self, record: &MoveRecord) -> UnmakeMoveResult {
        let board = self.board_mut();

        let ply = record.ply;

        let target = ply.target_square();

        let piece = Piece::pawn(
            board
                .get_piece(target)
                .ok_or(UnmakeMoveError::NoPiece(target))?
                .color,
        );

        let origin = ply.origin_square();
        board.place_piece(piece, origin, PlacePieceStrategy::PreserveExisting)?;

        board.remove_piece(target);

        Ok(())
    }

    fn unmake_castle_move(&mut self, record: &MoveRecord) -> UnmakeMoveResult {
        let ply = record.ply;

        let wing = ply.castle_wing().expect("no wing for unmaking castle move");
        let color = record.color;
        let parameters = Board::castling_parameters(wing, color);

        let board = self.board_mut();

        let king = board
            .get_piece(parameters.target.king)
            .ok_or(UnmakeMoveError::NoPiece(parameters.target.king))?;
        let rook = board
            .get_piece(parameters.target.rook)
            .ok_or(UnmakeMoveError::NoPiece(parameters.target.rook))?;

        board.place_piece(
            king,
            parameters.origin.king,
            PlacePieceStrategy::PreserveExisting,
        )?;
        board.place_piece(
            rook,
            parameters.origin.rook,
            PlacePieceStrategy::PreserveExisting,
        )?;

        board.remove_piece(parameters.target.king);
        board.remove_piece(parameters.target.rook);

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{MakeMove, Move, PromotionShape, ValidateMove};
    use chessfriend_board::test_board;
    use chessfriend_core::{Color, Square, Wing, piece};

    type TestResult = Result<(), Box<dyn std::error::Error>>;

    /// Helper function to test make/unmake idempotency
    fn test_make_unmake_idempotent(
        initial_board: &mut impl BoardProvider,
        ply: Move,
    ) -> TestResult {
        // Capture initial state
        let initial_state = initial_board.board().clone();

        // Make the move
        let record = initial_board.make_move(ply, ValidateMove::Yes)?;

        // Verify the move changed the board
        assert_ne!(
            *initial_board.board(),
            initial_state,
            "Move should change board state"
        );

        // Unmake the move
        initial_board.unmake_move(&record)?;

        // Verify we're back to the initial state
        assert_eq!(
            *initial_board.board(),
            initial_state,
            "Board should return to initial state after unmake"
        );

        Ok(())
    }

    #[test]
    fn unmake_quiet_move_ai_claude() -> TestResult {
        let mut board = test_board!(White Pawn on C2);

        let ply = Move::quiet(Square::C2, Square::C3);
        let record = board.make_move(ply, ValidateMove::Yes)?;

        // Verify move was made
        assert_eq!(board.get_piece(Square::C2), None);
        assert_eq!(board.get_piece(Square::C3), Some(piece!(White Pawn)));
        assert_eq!(board.active_color(), Color::Black);

        board.unmake_move(&record)?;

        // Verify original state restored
        assert_eq!(board.get_piece(Square::C2), Some(piece!(White Pawn)));
        assert_eq!(board.get_piece(Square::C3), None);
        assert_eq!(board.active_color(), Color::White);

        Ok(())
    }

    #[test]
    fn unmake_double_push_move_ai_claude() -> TestResult {
        let mut board = test_board!(White Pawn on E2);

        let ply = Move::double_push(Square::E2, Square::E4);
        let record = board.make_move(ply, ValidateMove::Yes)?;

        // Verify move was made
        assert_eq!(board.get_piece(Square::E2), None);
        assert_eq!(board.get_piece(Square::E4), Some(piece!(White Pawn)));
        assert_eq!(board.en_passant_target(), Some(Square::E3));
        assert_eq!(board.active_color(), Color::Black);

        board.unmake_move(&record)?;

        // Verify original state restored
        assert_eq!(board.get_piece(Square::E2), Some(piece!(White Pawn)));
        assert_eq!(board.get_piece(Square::E4), None);
        assert_eq!(board.en_passant_target(), None);
        assert_eq!(board.active_color(), Color::White);

        Ok(())
    }

    #[test]
    fn unmake_capture_move_ai_claude() -> TestResult {
        let mut board = test_board![
            White Bishop on C2,
            Black Rook on F5,
        ];

        let ply = Move::capture(Square::C2, Square::F5);
        let record = board.make_move(ply, ValidateMove::Yes)?;

        // Verify move was made
        assert_eq!(board.get_piece(Square::C2), None);
        assert_eq!(board.get_piece(Square::F5), Some(piece!(White Bishop)));
        assert_eq!(record.captured_piece, Some(piece!(Black Rook)));
        assert_eq!(board.active_color(), Color::Black);

        board.unmake_move(&record)?;

        // Verify original state restored
        assert_eq!(board.get_piece(Square::C2), Some(piece!(White Bishop)));
        assert_eq!(board.get_piece(Square::F5), Some(piece!(Black Rook)));
        assert_eq!(board.active_color(), Color::White);

        Ok(())
    }

    #[test]
    fn unmake_en_passant_capture_ai_claude() -> TestResult {
        let mut board = test_board![
            Black Pawn on F4,
            White Pawn on E2
        ];

        // Set up en passant situation
        let double_push = Move::double_push(Square::E2, Square::E4);
        board.make_move(double_push, ValidateMove::Yes)?;

        // Make en passant capture
        let en_passant = Move::en_passant_capture(Square::F4, Square::E3);
        let record = board.make_move(en_passant, ValidateMove::Yes)?;

        // Verify en passant was made
        assert_eq!(board.get_piece(Square::F4), None);
        assert_eq!(board.get_piece(Square::E3), Some(piece!(Black Pawn)));
        assert_eq!(
            board.get_piece(Square::E4),
            None,
            "captured pawn was not removed"
        );
        assert_eq!(record.captured_piece, Some(piece!(White Pawn)));

        board.unmake_move(&record)?;

        // Verify state before en passant is restored
        assert_eq!(board.get_piece(Square::F4), Some(piece!(Black Pawn)));
        assert_eq!(board.get_piece(Square::E3), None);
        assert_eq!(
            board.get_piece(Square::E4),
            Some(piece!(White Pawn)),
            "captured pawn was not restored"
        );
        assert_eq!(board.active_color(), Color::Black);

        Ok(())
    }

    #[test]
    fn unmake_promotion_move_ai_claude() -> TestResult {
        let mut board = test_board![
            White Pawn on F7,
        ];

        let ply = Move::promotion(Square::F7, Square::F8, PromotionShape::Queen);
        let record = board.make_move(ply, ValidateMove::Yes)?;

        // Verify promotion was made
        assert_eq!(board.get_piece(Square::F7), None);
        assert_eq!(board.get_piece(Square::F8), Some(piece!(White Queen)));
        assert_eq!(board.active_color(), Color::Black);

        board.unmake_move(&record)?;

        // Verify original pawn is restored
        assert_eq!(board.get_piece(Square::F7), Some(piece!(White Pawn)));
        assert_eq!(board.get_piece(Square::F8), None);
        assert_eq!(board.active_color(), Color::White);

        Ok(())
    }

    #[test]
    fn unmake_capture_promotion_ai_claude() -> TestResult {
        let mut board = test_board![
            White Pawn on F7,
            Black Rook on G8,
        ];

        let ply = Move::capture_promotion(Square::F7, Square::G8, PromotionShape::Queen);
        let record = board.make_move(ply, ValidateMove::Yes)?;

        // Verify promotion capture was made
        assert_eq!(board.get_piece(Square::F7), None);
        assert_eq!(board.get_piece(Square::G8), Some(piece!(White Queen)));
        assert_eq!(record.captured_piece, Some(piece!(Black Rook)));

        board.unmake_move(&record)?;

        // Verify original state restored
        assert_eq!(board.get_piece(Square::F7), Some(piece!(White Pawn)));
        assert_eq!(board.get_piece(Square::G8), Some(piece!(Black Rook)));
        assert_eq!(board.active_color(), Color::White);

        Ok(())
    }

    #[test]
    fn unmake_white_kingside_castle_ai_claude() -> TestResult {
        let mut board = test_board![
            White King on E1,
            White Rook on H1,
        ];

        let original_castling_rights = *board.castling_rights();

        let ply = Move::castle(Color::White, Wing::KingSide);
        let record = board.make_move(ply, ValidateMove::Yes)?;

        // Verify castle was made
        assert_eq!(board.get_piece(Square::E1), None);
        assert_eq!(board.get_piece(Square::H1), None);
        assert_eq!(board.get_piece(Square::G1), Some(piece!(White King)));
        assert_eq!(board.get_piece(Square::F1), Some(piece!(White Rook)));
        assert!(!board.has_castling_right_unwrapped(Color::White, Wing::KingSide));

        board.unmake_move(&record)?;

        // Verify original state restored
        assert_eq!(board.get_piece(Square::E1), Some(piece!(White King)));
        assert_eq!(board.get_piece(Square::H1), Some(piece!(White Rook)));
        assert_eq!(board.get_piece(Square::G1), None);
        assert_eq!(board.get_piece(Square::F1), None);
        assert_eq!(*board.castling_rights(), original_castling_rights);
        assert_eq!(board.active_color(), Color::White);

        Ok(())
    }

    #[test]
    fn unmake_white_queenside_castle_ai_claude() -> TestResult {
        let mut board = test_board![
            White King on E1,
            White Rook on A1,
        ];

        let original_castling_rights = *board.castling_rights();

        let ply = Move::castle(Color::White, Wing::QueenSide);
        let record = board.make_move(ply, ValidateMove::Yes)?;

        // Verify castle was made
        assert_eq!(board.get_piece(Square::E1), None);
        assert_eq!(board.get_piece(Square::A1), None);
        assert_eq!(board.get_piece(Square::C1), Some(piece!(White King)));
        assert_eq!(board.get_piece(Square::D1), Some(piece!(White Rook)));
        assert!(!board.has_castling_right_unwrapped(Color::White, Wing::QueenSide));

        board.unmake_move(&record)?;

        // Verify original state restored
        assert_eq!(board.get_piece(Square::E1), Some(piece!(White King)));
        assert_eq!(board.get_piece(Square::A1), Some(piece!(White Rook)));
        assert_eq!(board.get_piece(Square::C1), None);
        assert_eq!(board.get_piece(Square::D1), None);
        assert_eq!(*board.castling_rights(), original_castling_rights);
        assert_eq!(board.active_color(), Color::White);

        Ok(())
    }

    #[test]
    fn unmake_black_kingside_castle() -> TestResult {
        let mut board = test_board!(Black, [
            Black King on E8,
            Black Rook on H8,
        ]);

        let original_castling_rights = *board.castling_rights();

        let ply = Move::castle(Color::Black, Wing::KingSide);
        let record = board.make_move(ply, ValidateMove::Yes)?;

        // Verify castle was made
        assert_eq!(board.get_piece(Square::E8), None);
        assert_eq!(board.get_piece(Square::H8), None);
        assert_eq!(board.get_piece(Square::G8), Some(piece!(Black King)));
        assert_eq!(board.get_piece(Square::F8), Some(piece!(Black Rook)));

        board.unmake_move(&record)?;

        // Verify original state restored
        assert_eq!(board.get_piece(Square::E8), Some(piece!(Black King)));
        assert_eq!(board.get_piece(Square::H8), Some(piece!(Black Rook)));
        assert_eq!(board.get_piece(Square::G8), None);
        assert_eq!(board.get_piece(Square::F8), None);
        assert_eq!(*board.castling_rights(), original_castling_rights);
        assert_eq!(board.active_color(), Color::Black);

        Ok(())
    }
}