1 files changed, 54 insertions, 9 deletions

diff --git a/src/board.rs b/src/board.rs
index 2afbe90..ef8dde7 100644
--- a/src/board.rs
+++ b/src/board.rs
@@ -2,7 +2,7 @@ use crate::{bitboard::{Bitboard, serialize_bitboard}, moves::Move, attacks::Atta
 
 pub static DEFAULT_FEN: &str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1";
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Board {
     pub pieces: [Bitboard; 12],
 
@@ -201,28 +201,28 @@ impl Board {
     /// *Blindlessly* apply a move without any validation
     /// Move should be validated beforehand
     pub fn make_move(&mut self, mov: Move) -> Option<PieceType> {
-        // Target
+        // Remove existing piece (if any) from target square
         let captured_piece = match self.pieces
             .iter()
             .enumerate()
             .find(|(piece_type, bitboard)| *bitboard & mov.target.to_bitboard() > 0)
             {
-                Some((piece_type, _)) => {
-                    self.pieces[piece_type] ^= mov.target.to_bitboard();
-                    Some(PieceType::from(piece_type))
+                Some((target_piece, _)) => {
+                    self.pieces[target_piece] ^= mov.target.to_bitboard();
+                    Some(PieceType::from(target_piece))
                 },
                 None => None,
             };
 
-        // Source
+        // Move a piece from source square to target
         match self.pieces
             .iter()
             .enumerate()
             .find(|(piece_type, bitboard)| *bitboard & mov.source.to_bitboard() > 0)
             {
-                Some((piece_type, _)) => {
-                    self.pieces[piece_type] ^= mov.source.to_bitboard();
-                    self.pieces[piece_type] |= mov.target.to_bitboard();
+                Some((source_piece, _)) => {
+                    self.pieces[source_piece] ^= mov.source.to_bitboard();
+                    self.pieces[source_piece] |= mov.target.to_bitboard();
                 },
                 None => panic!("Move is malformed: source piece not found"),
             };
@@ -232,6 +232,33 @@ impl Board {
 
         captured_piece
     }
+
+    /// Completely reverse make_move as if it never happened
+    pub fn unmake_move(&mut self, mov: Move, captured_piece: Option<PieceType>) {
+        // Remove source piece from target square
+        match self.pieces
+            .iter()
+            .enumerate()
+            .find(|(piece_type, bitboard)| *bitboard & mov.target.to_bitboard() > 0)
+            {
+                Some((source_piece, _)) => {
+                    self.pieces[source_piece] ^= mov.target.to_bitboard();
+                    self.pieces[source_piece] |= mov.source.to_bitboard();
+                },
+                None => panic!("Trying to unmake move which was not made: no piece was found on target square"),
+            };
+
+        // Return captured piece to target square
+        match captured_piece {
+            Some(target_piece) => {
+                self.pieces[target_piece as usize] |= mov.target.to_bitboard();
+            },
+            None => {}
+        }
+
+        self.update_occupancy();
+        self.ply -= 1;
+    }
 }
 
 
@@ -335,6 +362,24 @@ mod tests {
         assert!(board.pieces[PieceType::Queen as usize] & Square::D2.to_bitboard() == 0);
         assert!(board.pieces[PieceType::Queen as usize] & Square::A5.to_bitboard() > 0);
         assert!(board.ply == 2);
+    }
+
+    #[test]
+    fn test_unmake_move() {
+        let fen = String::from("q1b2k2/5p1p/4p1pb/pPPp4/3N4/3nPB2/P2QKnR1/1R6 w - - 0 25");
+        let mut board = Board::from_FEN(fen);
+        let initial_board = board.clone();
+
+        let mov = Move { source: Square::D2, target: Square::A5 };
+
+        board.print();
+
+        let captured_piece = board.make_move(mov);
+        board.print();
+
+        board.unmake_move(mov, captured_piece);
+        board.print();
 
+        assert_eq!(board, initial_board, "Board state after unmake_move should be the same as before make_move");
     }
 }