blunders_engine/boardrepr/
piece_sets.rs

1//! Piece-Centric representation of a chess board.
2
3use std::fmt::{self, Display};
4use std::ops::{Index, IndexMut, Range};
5
6use crate::bitboard::Bitboard;
7use crate::boardrepr::Mailbox;
8use crate::coretypes::{Color, Piece, PieceKind, Square};
9use crate::coretypes::{Color::*, PieceKind::*};
10
11// These offset impls are used to index their corresponding place in PieceSets.
12// PieceSets contains an array with one index for each kind of piece.
13// Optionally, Color and PieceKind discriminants could directly be these values
14// however they will stay decoupled for now to decrease dependency on enum ordering
15// until order is stabilized.
16
17impl Color {
18    /// Get the position of the start of the block for a color.
19    /// There are 6 piece_kinds per color, so one should start at 0, and the other at 6.
20    #[inline(always)]
21    const fn offset_block(&self) -> usize {
22        match self {
23            White => 0,
24            Black => 6,
25        }
26    }
27}
28impl PieceKind {
29    /// Get the offset of a piece_kind within a block.
30    /// Values must cover all numbers of [0, 1, 2, 3, 4, 5].
31    #[inline(always)]
32    const fn offset_pk(&self) -> usize {
33        match self {
34            King => 0,
35            Pawn => 1,
36            Knight => 2,
37            Queen => 3,
38            Rook => 4,
39            Bishop => 5,
40        }
41    }
42}
43impl Piece {
44    /// Get the completely qualified index for a piece.
45    #[inline(always)]
46    const fn offset(&self) -> usize {
47        self.color.offset_block() + self.piece_kind.offset_pk()
48    }
49}
50
51/// A Piece-Centric representation of pieces on a chessboard.
52/// A Bitboard is used to encode the squares of each chess piece.
53/// PieceSets indexes by piece to get squares, as opposed to Mailbox which
54/// indexes by square to get a piece.
55#[derive(Debug, Copy, Clone, Eq, PartialEq)]
56pub struct PieceSets {
57    pieces: [Bitboard; Self::SIZE],
58}
59
60impl PieceSets {
61    const SIZE: usize = 12; // 1 White, 1 Black BB for each piece type.
62
63    /// Returns PieceSets with all Bitboards set to empty.
64    pub fn new() -> Self {
65        PieceSets {
66            pieces: [Bitboard::EMPTY; Self::SIZE],
67        }
68    }
69
70    /// Returns PieceSets arranged in starting chess position.
71    pub fn start_position() -> Self {
72        let mb: Mailbox = Mailbox::start_position();
73        Self::from(&mb)
74    }
75
76    /// Return a bitboard representing the set of squares occupied by any piece.
77    /// Note: Compiler can auto-vectorize, however looking at assembly on godbolt
78    /// may be limited to avx128. Does not seem to use avx512 on supported cpus.
79    pub fn occupied(&self) -> Bitboard {
80        self.pieces.iter().fold(Bitboard::EMPTY, |acc, bb| acc | bb)
81    }
82
83    /// Return a bitboard representing the set of squares occupied by piece of color.
84    pub fn color_occupied(&self, color: Color) -> Bitboard {
85        self[color].iter().fold(Bitboard::EMPTY, |acc, bb| acc | bb)
86    }
87
88    /// Finds and returns the first piece found on target square, or None.
89    pub fn on_square(&self, sq: Square) -> Option<Piece> {
90        for player in Color::iter() {
91            for pk in PieceKind::iter() {
92                if self[(player, pk)].has_square(sq) {
93                    return Some(Piece::new(player, pk));
94                }
95            }
96        }
97        None
98    }
99
100    /// Finds and returns the first PieceKind found on target square of player's pieces, or None.
101    pub fn on_player_square(&self, player: Color, sq: Square) -> Option<PieceKind> {
102        for pk in PieceKind::iter() {
103            if self[(player, pk)].has_square(sq) {
104                return Some(pk);
105            }
106        }
107        None
108    }
109
110    /// Returns pretty-printed chess board representation of Self.
111    /// Uses Mailbox pretty.
112    pub fn pretty(&self) -> String {
113        Mailbox::from(self).pretty()
114    }
115
116    /// Returns true if all sets in self are disjoint (mutually exclusive).
117    /// In other words, there is no more than 1 piece per square. If a square is in one set, it is in no other.
118    /// PieceSets should be disjoint at all times.
119    pub fn is_disjoint(&self) -> bool {
120        let occupied_sum = self.occupied().count_squares();
121        let individual_sum = self
122            .pieces
123            .iter()
124            .fold(0, |acc, bb| acc + bb.count_squares());
125
126        occupied_sum == individual_sum
127    }
128
129    /// Returns true if Self is valid.
130    /// A valid PieceSets has the following properties:
131    /// * Has a single king per side.
132    /// * Each bitboard is disjoint (mutually exclusive) meaning a square cannot have more than one piece.
133    pub fn is_valid(&self) -> bool {
134        // Illegal if no White King.
135        if self[(White, King)].count_squares() != 1 {
136            return false;
137        }
138        // Illegal if no Black King.
139        if self[(Black, King)].count_squares() != 1 {
140            return false;
141        }
142        // Illegal if more than one piece per any square.
143        if !self.is_disjoint() {
144            return false;
145        }
146        // Illegal if any white pawn on first rank or black pawn on eighth rank.
147        let w_pawns_first_rank = self[(White, Pawn)] & Bitboard::RANK_1;
148        let b_pawns_eighth_rank = self[(Black, Pawn)] & Bitboard::RANK_8;
149        if !w_pawns_first_rank.is_empty() || !b_pawns_eighth_rank.is_empty() {
150            return false;
151        }
152
153        true
154    }
155}
156
157impl Index<Piece> for PieceSets {
158    type Output = Bitboard;
159    fn index(&self, piece: Piece) -> &Self::Output {
160        &self.pieces[piece.offset()]
161    }
162}
163
164impl IndexMut<Piece> for PieceSets {
165    fn index_mut(&mut self, piece: Piece) -> &mut Self::Output {
166        &mut self.pieces[piece.offset()]
167    }
168}
169
170impl Index<&Piece> for PieceSets {
171    type Output = Bitboard;
172    fn index(&self, piece: &Piece) -> &Self::Output {
173        &self.pieces[piece.offset()]
174    }
175}
176
177impl IndexMut<&Piece> for PieceSets {
178    fn index_mut(&mut self, piece: &Piece) -> &mut Self::Output {
179        &mut self.pieces[piece.offset()]
180    }
181}
182
183impl Index<(Color, PieceKind)> for PieceSets {
184    type Output = Bitboard;
185    fn index(&self, (color, piece_kind): (Color, PieceKind)) -> &Self::Output {
186        &self.pieces[color.offset_block() + piece_kind.offset_pk()]
187    }
188}
189
190impl IndexMut<(Color, PieceKind)> for PieceSets {
191    fn index_mut(&mut self, (color, piece_kind): (Color, PieceKind)) -> &mut Self::Output {
192        &mut self.pieces[color.offset_block() + piece_kind.offset_pk()]
193    }
194}
195
196impl Index<&(Color, PieceKind)> for PieceSets {
197    type Output = Bitboard;
198    fn index(&self, (color, piece_kind): &(Color, PieceKind)) -> &Self::Output {
199        &self.pieces[color.offset_block() + piece_kind.offset_pk()]
200    }
201}
202
203impl IndexMut<&(Color, PieceKind)> for PieceSets {
204    fn index_mut(&mut self, (color, piece_kind): &(Color, PieceKind)) -> &mut Self::Output {
205        &mut self.pieces[color.offset_block() + piece_kind.offset_pk()]
206    }
207}
208
209/// Get a slice of all pieces of same color.
210/// ```rust
211/// # use blunders_engine::{coretypes::Color, boardrepr::PieceSets, bitboard::Bitboard};
212/// let ps = PieceSets::start_position();
213/// let w_slice = &ps[Color::White];
214/// let b_slice = &ps[Color::Black];
215/// assert_eq!(b_slice.len(), w_slice.len());
216/// assert_eq!(b_slice.len(), 6);
217/// ```
218impl Index<Color> for PieceSets {
219    type Output = [Bitboard];
220    fn index(&self, color: Color) -> &Self::Output {
221        const RANGES: (Range<usize>, Range<usize>) = color_ranges();
222        match color {
223            White => &self.pieces[RANGES.0],
224            Black => &self.pieces[RANGES.1],
225        }
226    }
227}
228
229impl IndexMut<Color> for PieceSets {
230    fn index_mut(&mut self, color: Color) -> &mut Self::Output {
231        const RANGES: (Range<usize>, Range<usize>) = color_ranges();
232        match color {
233            White => &mut self.pieces[RANGES.0],
234            Black => &mut self.pieces[RANGES.1],
235        }
236    }
237}
238
239/// Used in 4 Index<Color> traits above to get correct ranges to represent each color's block.
240const fn color_ranges() -> (Range<usize>, Range<usize>) {
241    const W_RANGE: Range<usize> = match White.offset_block() < Black.offset_block() {
242        true => White.offset_block()..Black.offset_block(),
243        false => White.offset_block()..PieceSets::SIZE,
244    };
245    const B_RANGE: Range<usize> = match White.offset_block() < Black.offset_block() {
246        true => Black.offset_block()..PieceSets::SIZE,
247        false => Black.offset_block()..White.offset_block(),
248    };
249    (W_RANGE, B_RANGE)
250}
251
252impl From<&Mailbox> for PieceSets {
253    fn from(mb: &Mailbox) -> Self {
254        let mut pieces = Self::new();
255
256        for square in Square::iter() {
257            if let Some(ref piece) = mb[square] {
258                pieces[piece].set_square(square);
259            }
260        }
261        pieces
262    }
263}
264
265/// Defaults to standard chess piece starting positions.
266impl Default for PieceSets {
267    fn default() -> Self {
268        Self::start_position()
269    }
270}
271
272impl Display for PieceSets {
273    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
274        write!(f, "{}", self.pretty())
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281    use Square::*;
282    #[test]
283    fn piece_indexing() {
284        let pieces = PieceSets::start_position();
285        let w_king = &pieces[Piece::new(White, King)];
286        assert_eq!(w_king.count_squares(), 1);
287        assert!(w_king.has_square(E1));
288    }
289
290    #[test]
291    fn color_indexing() {
292        let pieces = PieceSets::start_position();
293        let white_pieces = &pieces[White];
294        let w_occupancy = white_pieces
295            .iter()
296            .fold(Bitboard::EMPTY, |acc, piece| acc | piece);
297        assert_eq!(w_occupancy.count_squares(), 16);
298        for square in [A1, B1, C1, D1, E1, F1, G1, H1] {
299            assert!(w_occupancy.has_square(&square));
300        }
301        for square in [A2, B2, C2, D2, E2, F2, G2, H2] {
302            assert!(w_occupancy.has_square(&square));
303        }
304
305        let black_pieces = &pieces[Black];
306        let b_occupancy = black_pieces
307            .iter()
308            .fold(Bitboard::EMPTY, |acc, piece| acc | piece);
309        assert_eq!(b_occupancy.count_squares(), 16);
310        for square in [A7, B7, C7, D7, E7, F7, G7, H7] {
311            assert!(b_occupancy.has_square(&square));
312        }
313        for square in [A8, B8, C8, D8, E8, F8, G8, H8] {
314            assert!(b_occupancy.has_square(&square));
315        }
316    }
317
318    #[test]
319    fn check_is_valid() {
320        let mut set = PieceSets::start_position();
321        assert!(set.is_valid());
322
323        set[(White, Pawn)].set_square(H8);
324        assert!(!set.is_valid());
325    }
326}