feat: implement separate mobility counter

1 files changed, 112 insertions, 1 deletions

diff --git a/src/board/mod.rs b/src/board/mod.rs
index d339f1d..2085894 100644
--- a/src/board/mod.rs
+++ b/src/board/mod.rs
@@ -1,4 +1,4 @@
-use crate::{bitboard::{Bitboard, serialize_bitboard, bitscan}, moves::{Move, MoveKind}, attacks::Attacks, square::Square};
+use crate::{bitboard::{Bitboard, serialize_bitboard, bitscan, pop_count}, moves::{Move, MoveKind}, attacks::Attacks, square::Square};
 mod engine;
 
 pub enum CastlingSide {
@@ -323,6 +323,107 @@ impl Board {
         moves
     }
 
+    /// Count pseudo-legal moves without actually generating them
+    pub fn mobility(&self, color: Color) -> f32 {
+        let mut mobility = 0.;
+        let opponent_occupancy = self.color_occupancy(color.flip());
+        let empty = self.empty();
+        let player_pieces = self.pieces_by_color(color);
+
+        for (piece_type, piece) in player_pieces.iter().enumerate() {
+            match PieceType::from(piece_type) {
+                PieceType::Pawn => {
+                    for source in serialize_bitboard(*piece) {
+                        let ep_bitboard = match self.ep_target {
+                            Some(square) => square.to_bitboard(),
+                            None => 0,
+                        };
+                        mobility += pop_count(self.attacks.pawn[color as usize][source as usize] & (opponent_occupancy | ep_bitboard)) as f32;
+                        mobility += pop_count(self.attacks.pawn_pushes[color as usize][source as usize] & empty) as f32;
+                    }
+                    let able_to_double_push_mask = match color {
+                        Color::White => empty >> 8,
+                        Color::Black => empty << 8,
+                    };
+                    for source in serialize_bitboard(*piece & able_to_double_push_mask) {
+                        mobility += pop_count(self.attacks.pawn_double_pushes[color as usize][source as usize] & empty) as f32;
+                    }
+                }
+                PieceType::King => {
+                    for source in serialize_bitboard(*piece) {
+                        mobility += pop_count(self.attacks.king[source as usize] & (empty | opponent_occupancy)) as f32;
+
+                        // Castling
+                        let king_home_position = match color {
+                            Color::White => Square::E1,
+                            Color::Black => Square::E8,
+                        };
+                        if *piece == king_home_position.to_bitboard() {
+                            for rook_square in serialize_bitboard(player_pieces[PieceType::Rook as usize])
+                                .iter()
+                                .filter(|rook_square| rook_square.rank() == king_home_position.rank())
+                                {
+                                    match rook_square.file() {
+                                        0 => {
+                                            let castle_line = [
+                                                king_home_position.west_one(),
+                                                king_home_position.west_one().west_one(),
+                                            ];
+
+                                            let all_empty = castle_line.iter().all(|square| empty & square.to_bitboard() > 0);
+                                            let any_checks = castle_line.iter().any(|square| self.is_square_attacked(*square, color.flip()));
+
+                                            if all_empty && !any_checks && self.castling_rights[color as usize][CastlingSide::Queen as usize] {
+                                                mobility += 1.;
+                                            }
+                                        },
+                                        7 => {
+                                            let castle_line = [
+                                                king_home_position.east_one(),
+                                                king_home_position.east_one().east_one(),
+                                            ];
+
+                                            let all_empty = castle_line.iter().all(|square| empty & square.to_bitboard() > 0);
+                                            let any_checks = castle_line.iter().any(|square| self.is_square_attacked(*square, color.flip()));
+
+                                            if all_empty && !any_checks && self.castling_rights[color as usize][CastlingSide::King as usize] {
+                                                mobility += 1.;
+                                            }
+                                        },
+                                        _ => {},
+                                    }
+                                }
+                        }
+                    }
+                }
+                PieceType::Knight => {
+                    for source in serialize_bitboard(*piece) {
+                        mobility += pop_count(self.attacks.knight[source as usize] & (empty | opponent_occupancy)) as f32;
+                    }
+                }
+                PieceType::Bishop => {
+                    for source in serialize_bitboard(*piece) {
+                        mobility += pop_count(self.attacks.bishop(self.occupancy, source) & (empty | opponent_occupancy)) as f32;
+                    }
+                }
+                PieceType::Rook => {
+                    for source in serialize_bitboard(*piece) {
+                        mobility += pop_count(self.attacks.rook(self.occupancy, source) & (empty | opponent_occupancy)) as f32;
+                    }
+                }
+                PieceType::Queen => {
+                    for source in serialize_bitboard(*piece) {
+                        mobility += pop_count(self.attacks.queen(self.occupancy, source) & (empty | opponent_occupancy)) as f32;
+                    }
+                }
+                incorrect_type => panic!("Incorrect piece type: {:?}", incorrect_type),
+            }
+        }
+
+        mobility
+    }
+
+
     /// *Blindlessly* apply a move without any validation
     /// Move should be validated beforehand
     pub fn make_move(&mut self, mov: Move) -> Option<PieceType> {
@@ -634,6 +735,16 @@ mod tests {
     }
 
     #[test]
+    fn mobility() {
+        let board = Board::new();
+        let white = board.mobility(Color::White);
+        let black = board.mobility(Color::Black);
+
+        assert_eq!(white, 20.);
+        assert_eq!(black, 20.);
+    }
+
+    #[test]
     fn make_move() {
         let fen = String::from("q1b2k2/5p1p/4p1pb/pPPp4/3N4/3nPB2/P2QKnR1/1R6 w - - 0 25");
         let mut board = Board::from_FEN(fen);