From 1c80db695fe1a3aba05010ec77fb19bb7764d04a Mon Sep 17 00:00:00 2001
From: stefiosif <si.iosifidis@gmail.com>
Date: Tue, 16 Jul 2024 21:52:14 +0300
Subject: [PATCH] Implement Index traits for Kind

---
 src/board.rs | 49 ++++++++++++++++++++++++++++++++++---------------
 src/move.rs  | 10 +++++-----
 2 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/src/board.rs b/src/board.rs
index ea38a38..eb32679 100644
--- a/src/board.rs
+++ b/src/board.rs
@@ -84,24 +84,21 @@ impl Board {
         };
 
         have_common_bit(
-            opponent[Kind::Pawn.idx()].bitboard,
+            opponent[Kind::Pawn].bitboard,
             fetch_pawn_attacks(square, own_color),
         ) || have_common_bit(
-            opponent[Kind::Knight.idx()].bitboard,
+            opponent[Kind::Knight].bitboard,
             fetch_knight_attacks(square),
         ) || have_common_bit(
-            opponent[Kind::Bishop.idx()].bitboard,
+            opponent[Kind::Bishop].bitboard,
             fetch_bishop_attacks(all_occupancies, square),
         ) || have_common_bit(
-            opponent[Kind::Rook.idx()].bitboard,
+            opponent[Kind::Rook].bitboard,
             fetch_rook_attacks(all_occupancies, square),
         ) || have_common_bit(
-            opponent[Kind::Queen.idx()].bitboard,
+            opponent[Kind::Queen].bitboard,
             fetch_queen_attacks(all_occupancies, square),
-        ) || have_common_bit(
-            opponent[Kind::King.idx()].bitboard,
-            fetch_king_attacks(square),
-        )
+        ) || have_common_bit(opponent[Kind::King].bitboard, fetch_king_attacks(square))
     }
 
     pub fn is_move_legit(&self, square: usize, opponent_color: Color) -> bool {
@@ -123,12 +120,12 @@ impl Board {
         let all_occupancies = self.all_occupancies();
         let (pieces, opponent_occupancies, own_occupancies) = match color {
             Color::White => (
-                self.white_pieces[kind.idx()].bitboard,
+                self.white_pieces[kind].bitboard,
                 self.black_occupancies(),
                 self.white_occupancies(),
             ),
             Color::Black => (
-                self.black_pieces[kind.idx()].bitboard,
+                self.black_pieces[kind].bitboard,
                 self.white_occupancies(),
                 self.black_occupancies(),
             ),
@@ -152,8 +149,8 @@ impl Board {
 
     pub fn set_piece(&mut self, piece: Piece) {
         match piece.color {
-            Color::Black => self.black_pieces[piece.kind.idx()].bitboard |= piece.bitboard,
-            Color::White => self.white_pieces[piece.kind.idx()].bitboard |= piece.bitboard,
+            Color::Black => self.black_pieces[piece.kind].bitboard |= piece.bitboard,
+            Color::White => self.white_pieces[piece.kind].bitboard |= piece.bitboard,
         };
     }
 
@@ -196,8 +193,30 @@ pub enum Kind {
 }
 
 impl Kind {
-    pub const fn idx(&self) -> usize {
-        *self as usize
+    fn idx(&self) -> usize {
+        match self {
+            Kind::Pawn => 0,
+            Kind::Knight => 1,
+            Kind::Bishop => 2,
+            Kind::Rook => 3,
+            Kind::Queen => 4,
+            Kind::King => 5,
+        }
+    }
+}
+use std::ops::{Index, IndexMut};
+
+impl Index<Kind> for [Piece] {
+    type Output = Piece;
+
+    fn index(&self, kind: Kind) -> &Self::Output {
+        &self[kind.idx()]
+    }
+}
+
+impl IndexMut<Kind> for [Piece] {
+    fn index_mut(&mut self, kind: Kind) -> &mut Self::Output {
+        &mut self[kind.idx()]
     }
 }
 
diff --git a/src/move.rs b/src/move.rs
index 2f8a79c..422bf5a 100644
--- a/src/move.rs
+++ b/src/move.rs
@@ -63,7 +63,7 @@ impl Board {
             .update_castling_state_capture(mv.dst, Color::opponent_color(color));
         self.state.change_side();
 
-        let own_king_square = lsb(pieces[Kind::King.idx()].bitboard);
+        let own_king_square = lsb(pieces[Kind::King].bitboard);
         self.is_move_legit(own_king_square, Color::opponent_color(color))
     }
 
@@ -138,10 +138,10 @@ impl Board {
 
     fn promote_piece(square: usize, pieces: &mut [Piece; 6], promote: Promote) {
         match promote {
-            Promote::Knight => pieces[Kind::Knight.idx()].bitboard |= square_to_bitboard(square),
-            Promote::Bishop => pieces[Kind::Bishop.idx()].bitboard |= square_to_bitboard(square),
-            Promote::Rook => pieces[Kind::Rook.idx()].bitboard |= square_to_bitboard(square),
-            Promote::Queen => pieces[Kind::Queen.idx()].bitboard |= square_to_bitboard(square),
+            Promote::Knight => pieces[Kind::Knight].bitboard |= square_to_bitboard(square),
+            Promote::Bishop => pieces[Kind::Bishop].bitboard |= square_to_bitboard(square),
+            Promote::Rook => pieces[Kind::Rook].bitboard |= square_to_bitboard(square),
+            Promote::Queen => pieces[Kind::Queen].bitboard |= square_to_bitboard(square),
         };
     }