From adad4118ae73008bccb27b5e8a7e9df0f51b50aa Mon Sep 17 00:00:00 2001
From: Bruno BELANYI <bruno@belanyi.fr>
Date: Wed, 3 Apr 2024 23:24:12 +0100
Subject: [PATCH] Account for captures in 'ChessBoard::{,un}do_move'

This is a silly thing to forget...
---
 src/board/chess_board/mod.rs | 38 ++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/src/board/chess_board/mod.rs b/src/board/chess_board/mod.rs
index 63dbe51..e7d9efd 100644
--- a/src/board/chess_board/mod.rs
+++ b/src/board/chess_board/mod.rs
@@ -37,6 +37,7 @@ pub struct NonReversibleState {
     castle_rights: [CastleRights; Color::NUM_VARIANTS],
     en_passant: Option<Square>,
     half_move_clock: u32, // Should *probably* never go higher than 100.
+    captured_piece: Option<Piece>,
 }
 
 impl ChessBoard {
@@ -134,11 +135,13 @@ impl ChessBoard {
     /// Play the given [Move], returning all non-revertible state (e.g: en-passant, etc...).
     #[inline(always)]
     pub fn do_move(&mut self, chess_move: Move) -> NonReversibleState {
+        let opponent = !self.current_player();
         // Save non-revertible state
         let state = NonReversibleState {
             castle_rights: self.castle_rights,
             en_passant: self.en_passant,
             half_move_clock: self.half_move_clock,
+            captured_piece: chess_move.capture(),
         };
 
         // Non-revertible state modification
@@ -167,6 +170,11 @@ impl ChessBoard {
                 _ => *castle_rights,
             }
         }
+        if let Some(piece) = chess_move.capture() {
+            *self.piece_occupancy_mut(piece) ^= chess_move.destination();
+            *self.color_occupancy_mut(opponent) ^= chess_move.destination();
+            self.combined_occupancy ^= chess_move.destination();
+        }
 
         // Revertible state modification
         self.xor(
@@ -188,6 +196,11 @@ impl ChessBoard {
         self.castle_rights = previous.castle_rights;
         self.en_passant = previous.en_passant;
         self.half_move_clock = previous.half_move_clock;
+        if let Some(piece) = previous.captured_piece {
+            *self.piece_occupancy_mut(piece) ^= chess_move.destination();
+            *self.color_occupancy_mut(self.current_player()) ^= chess_move.destination();
+            self.combined_occupancy ^= chess_move.destination();
+        }
 
         // Restore revertible state
         self.xor(
@@ -837,4 +850,29 @@ mod test {
                 .unwrap()
         );
     }
+
+    #[test]
+    fn do_move_undo_capture() {
+        let mut position = ChessBoard::from_fen("3q3k/8/8/8/8/8/8/K2Q4 w - - 0 1").unwrap();
+        let expected = ChessBoard::from_fen("3Q3k/8/8/8/8/8/8/K7 b - - 0 1").unwrap();
+        let original = position.clone();
+
+        let capture = MoveBuilder {
+            piece: Piece::Queen,
+            start: Square::D1,
+            destination: Square::D8,
+            capture: Some(Piece::Queen),
+            promotion: None,
+            en_passant: false,
+            double_step: false,
+            castling: false,
+        }
+        .into();
+
+        let state = position.do_move(capture);
+        assert_eq!(position, expected);
+
+        position.undo_move(capture, state);
+        assert_eq!(position, original);
+    }
 }