impl {Index,IndexMut} for {ByColor,ByRole,ByCastlingSide}

niklasf · Jan 4, 2025 · a5e70fd · a5e70fd
1 parent fbc087b
commit a5e70fd
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 14 deletions.
diff --git a/src/castling_side.rs b/src/castling_side.rs
@@ -244,6 +244,22 @@ impl<T> IntoIterator for ByCastlingSide<T> {
     }
 }
 
+impl<T> ops::Index<CastlingSide> for ByCastlingSide<T> {
+    type Output = T;
+
+    #[inline]
+    fn index(&self, index: CastlingSide) -> &T {
+        self.get(index)
+    }
+}
+
+impl<T> ops::IndexMut<CastlingSide> for ByCastlingSide<T> {
+    #[inline]
+    fn index_mut(&mut self, index: CastlingSide) -> &mut T {
+        self.get_mut(index)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

diff --git a/src/color.rs b/src/color.rs
@@ -396,3 +396,19 @@ impl<T> IntoIterator for ByColor<T> {
         [self.white, self.black].into_iter()
     }
 }
+
+impl<T> ops::Index<Color> for ByColor<T> {
+    type Output = T;
+
+    #[inline]
+    fn index(&self, index: Color) -> &T {
+        self.get(index)
+    }
+}
+
+impl<T> ops::IndexMut<Color> for ByColor<T> {
+    #[inline]
+    fn index_mut(&mut self, index: Color) -> &mut T {
+        self.get_mut(index)
+    }
+}
diff --git a/src/role.rs b/src/role.rs
@@ -1,4 +1,4 @@
-use core::{array, convert::identity, num};
+use core::{array, convert::identity, num, ops};
 
 use crate::{color::Color, types::Piece, util::out_of_range_error};
 
@@ -339,3 +339,19 @@ impl<T> IntoIterator for ByRole<T> {
         .into_iter()
     }
 }
+
+impl<T> ops::Index<Role> for ByRole<T> {
+    type Output = T;
+
+    #[inline]
+    fn index(&self, role: Role) -> &T {
+        self.get(role)
+    }
+}
+
+impl<T> ops::IndexMut<Role> for ByRole<T> {
+    #[inline]
+    fn index_mut(&mut self, role: Role) -> &mut T {
+        self.get_mut(role)
+    }
+}
diff --git a/src/zobrist.rs b/src/zobrist.rs
@@ -27,8 +27,8 @@ use core::{
 };
 
 use crate::{
-    Board, ByColor, ByRole, CastlingSide, Color, EnPassantMode, File, Piece, Position,
-    RemainingChecks, Role, Square,
+    Board, ByCastlingSide, ByColor, ByRole, CastlingSide, Color, EnPassantMode, File, Piece,
+    Position, RemainingChecks, Role, Square,
 };
 
 /// Integer type that can be returned as a Zobrist hash.
@@ -940,7 +940,7 @@ macro_rules! zobrist_value_impl {
                     },
                 };
 
-                $t(PIECE_MASKS.get(piece.color).get(piece.role)[usize::from(square)])
+                $t(PIECE_MASKS.piece(piece)[usize::from(square)])
             }
 
             #[inline]
@@ -952,18 +952,18 @@ macro_rules! zobrist_value_impl {
             #[inline]
             fn zobrist_for_castling_right(color: Color, side: CastlingSide) -> $t {
                 #[allow(overflowing_literals)]
-                static CASTLING_RIGHT_MASKS: ByColor<[$proxy; 2]> = ByColor {
-                    black: [
-                        0x7756_8e6e_6151_6b92_a57e_6339_dd2c_f3a0,
-                        0xd153_e6cf_8d19_84ea_1ef6_e6db_b196_1ec9,
-                    ],
-                    white: [
-                        0xca3c_7f8d_050c_44ba_31d7_1dce_64b2_c310,
-                        0x8f50_a115_834e_5414_f165_b587_df89_8190,
-                    ],
+                static CASTLING_RIGHT_MASKS: ByColor<ByCastlingSide<$proxy>> = ByColor {
+                    black: ByCastlingSide {
+                        king_side: 0x7756_8e6e_6151_6b92_a57e_6339_dd2c_f3a0,
+                        queen_side: 0xd153_e6cf_8d19_84ea_1ef6_e6db_b196_1ec9,
+                    },
+                    white: ByCastlingSide {
+                        king_side: 0xca3c_7f8d_050c_44ba_31d7_1dce_64b2_c310,
+                        queen_side: 0x8f50_a115_834e_5414_f165_b587_df89_8190,
+                    },
                 };
 
-                $t(CASTLING_RIGHT_MASKS.get(color)[side as usize])
+                $t(CASTLING_RIGHT_MASKS.get(color)[side])
             }
 
             #[inline]