chessfriend/board/src/position/pieces.rs

// Eryn Wells <eryn@erynwells.me>

use super::Position;
use crate::bitboard::BitBoard;
use crate::piece::{Color, Piece, PlacedPiece, Shape};
use crate::Square;

pub struct Pieces<'a> {
    color: Color,
    position: &'a Position,

    current_shape: Option<Shape>,

    shape_iterator: Box<dyn Iterator<Item = &'static Shape>>,
    square_iterator: Option<Box<dyn Iterator<Item = Square>>>,
}

impl<'a> Pieces<'a> {
    pub(crate) fn new(position: &Position, color: Color) -> Pieces {
        Pieces {
            color,
            position,
            current_shape: None,
            shape_iterator: Box::new(Shape::iter()),
            square_iterator: None,
        }
    }
}

impl<'a> Iterator for Pieces<'a> {
    type Item = PlacedPiece;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(square_iterator) = &mut self.square_iterator {
            if let (Some(square), Some(shape)) = (square_iterator.next(), self.current_shape) {
                return Some(PlacedPiece::new(Piece::new(self.color, shape), square));
            }
        }

        let mut current_shape: Option<Shape> = None;
        let mut next_nonempty_bitboard: Option<&BitBoard> = None;

        while let Some(shape) = self.shape_iterator.next() {
            let piece = Piece::new(self.color, *shape);

            let bitboard = self.position.bitboard_for_piece(piece);
            if bitboard.is_empty() {
                continue;
            }

            next_nonempty_bitboard = Some(bitboard);
            current_shape = Some(*shape);

            break;
        }

        if let (Some(bitboard), Some(shape)) = (next_nonempty_bitboard, current_shape) {
            let mut square_iterator = bitboard.occupied_squares();

            let mut next_placed_piece: Option<PlacedPiece> = None;
            if let Some(square) = square_iterator.next() {
                next_placed_piece = Some(PlacedPiece::new(Piece::new(self.color, shape), square));
            }

            self.square_iterator = Some(Box::new(square_iterator));
            self.current_shape = Some(shape);

            return next_placed_piece;
        }

        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::piece::{Color, Piece, Shape};
    use crate::Position;
    use crate::Square;
    use std::collections::HashSet;

    fn square_at(sq: &str) -> Square {
        Square::from_algebraic_str(sq).expect(sq)
    }

    fn place_piece_in_position(pos: &mut Position, sq: &str, piece: Piece) {
        pos.place_piece(&piece, &square_at(sq))
            .expect("Unable to place {piece:?} queen on {sq}");
    }

    #[test]
    fn empty() {
        let pos = Position::empty();
        let mut pieces = pos.pieces(Color::White);
        assert_eq!(pieces.next(), None);
    }

    #[test]
    fn one() {
        let sq = Square::from_algebraic_str("e4").expect("e4");

        let mut pos = Position::empty();
        pos.place_piece(&Piece::new(Color::White, Shape::Queen), &sq)
            .expect("Unable to place white queen on e4");
        println!("{:?}", &pos);

        let mut pieces = pos.pieces(Color::White);
        assert_eq!(
            pieces.next(),
            Some(PlacedPiece::new(Piece::new(Color::White, Shape::Queen), sq))
        );
        assert_eq!(pieces.next(), None);
    }

    #[test]
    fn multiple_pieces() {
        let mut pos = Position::empty();

        place_piece_in_position(&mut pos, "e4", Piece::new(Color::White, Shape::Queen));
        place_piece_in_position(&mut pos, "e1", Piece::new(Color::White, Shape::King));
        place_piece_in_position(&mut pos, "b2", Piece::new(Color::White, Shape::Pawn));
        place_piece_in_position(&mut pos, "c2", Piece::new(Color::White, Shape::Pawn));

        println!("{:?}", &pos);

        let expected_placed_pieces = HashSet::from([
            PlacedPiece::new(Piece::new(Color::White, Shape::Queen), square_at("e4")),
            PlacedPiece::new(Piece::new(Color::White, Shape::King), square_at("e1")),
            PlacedPiece::new(Piece::new(Color::White, Shape::Pawn), square_at("b2")),
            PlacedPiece::new(Piece::new(Color::White, Shape::Pawn), square_at("c2")),
        ]);

        let placed_pieces = HashSet::from_iter(pos.pieces(Color::White));

        assert_eq!(placed_pieces, expected_placed_pieces);
    }
}