crabchess 0.1.15

//! This module contains types representing a chessboard and its squares.

use smallvec::SmallVec;

use crate::{
    error::{Err, Result},
    pieces::{Color, Piece, Type},
    squares::{Direction, File, Rank, Square},
};
use std::{fmt::Display, io::IsTerminal, iter::zip, sync::OnceLock};

/// The standard starting position for a game of chess, in FEN notation.
pub const STAUNTON_PATTERN: &str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1";

/// Simple struct to track kings within a `ChessPosition`.
#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
pub struct Kings {
    pub white: Square,
    pub black: Square,
}

impl Kings {
    /// Get a king's location by its color.
    #[must_use]
    pub const fn get(&self, color: Color) -> Square {
        match color {
            Color::White => self.white,
            Color::Black => self.black,
        }
    }

    /// Set a king's location by its color.
    pub fn set(&mut self, color: Color, value: Square) {
        match color {
            Color::White => self.white = value,
            Color::Black => self.black = value,
        }
    }
}

impl TryFrom<&Board> for Kings {
    type Error = Err;

    fn try_from(value: &Board) -> std::prelude::v1::Result<Self, Self::Error> {
        let mut white_king: SmallVec<[Square; 1]> = SmallVec::new();
        let mut black_king: SmallVec<[Square; 1]> = SmallVec::new();

        for rank in Rank::all() {
            for file in File::all() {
                let square = Square(file, rank);
                if let Some(piece) = value.get(square) {
                    if piece.piece_type == Type::King {
                        match piece.color {
                            Color::White => white_king.push(square),
                            Color::Black => black_king.push(square),
                        }
                    }
                }
            }
        }

        for (color, king_vec) in [(Color::White, &white_king), (Color::Black, &black_king)] {
            let len = king_vec.len();
            if len != 1 {
                return Err(Err::KingCountError(color, len));
            }
        }

        Ok(Self {
            white: white_king[0],
            black: black_king[0],
        })
    }
}

/// A chessboard.
///
/// This struct does not contain all the information needed to understand a chess position, like
/// repetition history, turn, and castling rights. It merely holds the pieces.
#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)]
pub struct Board([[Option<Piece>; 8]; 8]);

impl Board {
    /// Create an empty board.
    #[must_use]
    pub const fn new() -> Self {
        Self([[None; 8]; 8])
    }

    /// Get a square's value by reference.
    #[must_use]
    pub const fn get(&self, sq: Square) -> &Option<Piece> {
        &self.0[sq.0 as usize][sq.1 as usize]
    }

    /// Set a square's value.
    pub fn set(&mut self, square: Square, value: Option<Piece>) {
        self.0[square.0 as usize][square.1 as usize] = value;
    }

    /// Create a board from a FEN string.
    ///
    /// # Errors
    ///
    /// Returns an error if FEN is invalid.
    #[allow(clippy::cast_possible_truncation)]
    pub fn from_fen_slice(fen: &str) -> Result<Self> {
        let mut board = Self::new();
        let piece_grid: &str = match fen.split(' ').next() {
            Some(result) => result,
            None => return Err(Err::InvalidFenError("piece placement data", fen.into())),
        };
        let mut cursor: Square;

        for (rank, segment) in zip(Rank::reversed(), piece_grid.split('/')) {
            cursor = Square(File::A, rank);

            for ch in segment
                .trim_end_matches(|c: char| c.is_ascii_digit())
                .chars()
            {
                if let Some(digit) = ch.to_digit(10) {
                    match cursor.step(digit as i8, 0) {
                        Some(cur) => cursor = cur,
                        None => {
                            return Err(Err::InvalidFenError("piece placement data", fen.into()))
                        }
                    }
                } else {
                    let Ok(piece) = Piece::from_ascii(ch) else {
                        return Err(Err::InvalidFenError("piece placement data", fen.into()));
                    };
                    board.set(cursor, Some(piece));
                    if let Some(cur) = cursor.step(1, 0) {
                        cursor = cur;
                    }
                }
            }
        }

        Ok(board)
    }

    /// Get squares between two others. Exclusive on both sides of range.
    ///
    /// # Errors
    ///
    /// Returns an error if squares are not in-line.
    pub fn between(square: Square, other: Square) -> Result<SmallVec<[Square; 6]>> {
        if square == other {
            return Ok(SmallVec::new());
        }
        let direction: Direction = Direction::between(square, other)?;
        let mut output: SmallVec<[Square; 6]> = SmallVec::new();

        let mut touched = false;
        for sq in square.iter_dir(direction) {
            if sq == other {
                touched = true;
                break;
            }
            output.push(sq);
        }

        if !touched {
            return Err(Err::SquaresNotInlineError(square, other));
        }

        Ok(output)
    }

    /// Get the piece placement data in FEN notation. Output only contains the first token of a FEN.
    #[must_use]
    pub fn fen_repr(&self) -> String {
        let mut result = String::new();
        let mut consecutive_empty_squares: u32 = 0;

        for rank in Rank::reversed() {
            for file in File::all() {
                let square = Square(file, rank);
                if let Some(piece) = self.get(square) {
                    if consecutive_empty_squares > 0 {
                        result.push_str(&consecutive_empty_squares.to_string());
                    }
                    result.push(piece.ascii());
                    consecutive_empty_squares = 0;
                } else {
                    consecutive_empty_squares += 1;
                }
            }

            if consecutive_empty_squares > 0 {
                result.push_str(&consecutive_empty_squares.to_string());
                consecutive_empty_squares = 0;
            }

            if rank != Rank::One {
                result.push('/');
            }
        }

        result
    }

    /// Get all pieces and their location for a color, or if unspecified, both colors.
    #[must_use]
    pub fn pieces(&self, color: Option<Color>) -> SmallVec<[(Square, Piece); 32]> {
        let mut output = SmallVec::new();

        for file in File::all() {
            for rank in Rank::all() {
                let square = Square(file, rank);

                if let Some(piece) = self.get(square) {
                    match color {
                        Some(color) if piece.color == color => output.push((square, *piece)),
                        None => output.push((square, *piece)),
                        _ => (),
                    }
                }
            }
        }

        output
    }
}

impl Display for Board {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        fn symbol_getter() -> &'static fn(Piece) -> char {
            static SYMBOL_GETTER: OnceLock<fn(Piece) -> char> = OnceLock::new();
            SYMBOL_GETTER.get_or_init(|| {
                if std::io::stdin().is_terminal() {
                    Piece::unicode_dark_background
                } else {
                    Piece::unicode_light_background
                }
            })
        }

        let mut result = String::new();
        for rank in Rank::reversed() {
            result.push(rank.char());
            result.push(' ');
            File::all()
                .map(|file| self.get(Square(file, rank)).map_or('.', symbol_getter()))
                .into_iter()
                .for_each(|c| {
                    result.push(c);
                    result.push(' ');
                });
            result.push('\n');
        }
        result.push_str("  a b c d e f g h ");
        write!(f, "{result}")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        pieces::{Color::White, Type::*},
        sq, sqs,
    };

    #[test]
    fn fen_piece_placement() {
        let fen = "r2r4/p2p1p1p/b6R/n1p1kp2/2P2P2/3BP3/PP5P/4K2R b K f3 0 22";
        let board = Board::from_fen_slice(fen).unwrap();
        assert_eq!(
            "r2r4/p2p1p1p/b6R/n1p1kp2/2P2P2/3BP3/PP5P/4K2R",
            board.fen_repr()
        );
    }

    #[test]
    fn set_and_get_piece() {
        let mut board = Board::new();
        let piece = Piece {
            piece_type: Rook,
            color: White,
        };
        board.set(sq!(C5), Some(piece));
        assert_eq!(piece, board.get(sq!(C5)).unwrap());
    }

    #[test]
    fn test_between() {
        assert_eq!(
            Board::between(sq!(A2), sq!(A7)).unwrap().to_vec(),
            sqs![A3, A4, A5, A6]
        );
        assert_eq!(
            Board::between(sq!(G6), sq!(D3)).unwrap().to_vec(),
            sqs![F5, E4]
        );
    }

    #[test]
    fn import_fen() {
        let fen = "r2r4/p2p1p1p/b6R/n1p1kp2/2P2P2/3BP3/PP5P/4K2R b K f3 0 22";
        let board = Board::from_fen_slice(fen).unwrap();
        assert_eq!(board.get(sq!(E1)).unwrap(), Piece::new(King, White));
        assert_eq!(board.get(sq!(A2)).unwrap(), Piece::new(Pawn, White));
        assert!(!board.get(sq!(B4)).is_some());
        assert!(!board.get(sq!(F2)).is_some());
    }
}