chessfriend/board/src/fen.rs

// Eryn Wells <eryn@erynwells.me>

use crate::{piece_sets::PlacePieceStrategy, Board, Castle};
use chessfriend_core::{
    coordinates::ParseSquareError, piece, Color, File, Piece, PlacedPiece, Rank, Square,
};
use std::fmt::Write;

#[macro_export]
macro_rules! fen {
    ($fen_string:literal) => {
        Board::from_fen_str($fen_string)
    };
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ToFenStrError {
    FmtError(std::fmt::Error),
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum FromFenStrError {
    MissingField(Field),
    MissingPlacement,
    InvalidValue,
    ParseIntError(std::num::ParseIntError),
    ParseSquareError(ParseSquareError),
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Field {
    Placements,
    PlayerToMove,
    CastlingRights,
    EnPassantSquare,
    HalfMoveClock,
    FullMoveCounter,
}

pub trait ToFenStr {
    type Error;

    /// Create a FEN string from `Self`.
    ///
    /// # Errors
    ///
    ///
    fn to_fen_str(&self) -> Result<String, Self::Error>;
}

pub trait FromFenStr: Sized {
    type Error;

    /// Create a `Self` from a FEN string.
    ///
    /// # Errors
    ///
    ///
    fn from_fen_str(string: &str) -> Result<Self, Self::Error>;
}

impl ToFenStr for Board {
    type Error = ToFenStrError;

    fn to_fen_str(&self) -> Result<String, Self::Error> {
        let mut fen_string = String::new();

        let mut empty_squares: u8 = 0;
        for rank in Rank::ALL.into_iter().rev() {
            for file in File::ALL {
                let square = Square::from_file_rank(file, rank);
                match self.get_piece(square) {
                    Some(piece) => {
                        if empty_squares > 0 {
                            write!(fen_string, "{empty_squares}")
                                .map_err(ToFenStrError::FmtError)?;
                            empty_squares = 0;
                        }
                        write!(fen_string, "{}", piece.to_fen_str()?)
                            .map_err(ToFenStrError::FmtError)?;
                    }
                    None => empty_squares += 1,
                }
            }

            if empty_squares > 0 {
                write!(fen_string, "{empty_squares}").map_err(ToFenStrError::FmtError)?;
                empty_squares = 0;
            }

            if rank != Rank::ONE {
                write!(fen_string, "/").map_err(ToFenStrError::FmtError)?;
            }
        }

        write!(fen_string, " {}", self.active_color.to_fen_str()?)
            .map_err(ToFenStrError::FmtError)?;

        let castling = [
            (Color::White, Castle::KingSide),
            (Color::White, Castle::QueenSide),
            (Color::Black, Castle::KingSide),
            (Color::Black, Castle::QueenSide),
        ]
        .map(|(color, castle)| {
            if !self.castling_rights.color_has_right(color, castle) {
                return "";
            }

            match (color, castle) {
                (Color::White, Castle::KingSide) => "K",
                (Color::White, Castle::QueenSide) => "Q",
                (Color::Black, Castle::KingSide) => "k",
                (Color::Black, Castle::QueenSide) => "q",
            }
        })
        .concat();

        write!(
            fen_string,
            " {}",
            if castling.is_empty() { "-" } else { &castling }
        )
        .map_err(ToFenStrError::FmtError)?;

        write!(
            fen_string,
            " {}",
            self.en_passant_target
                .map_or("-".to_string(), |square| square.to_string())
        )
        .map_err(ToFenStrError::FmtError)?;

        write!(fen_string, " {}", self.half_move_clock).map_err(ToFenStrError::FmtError)?;
        write!(fen_string, " {}", self.full_move_number).map_err(ToFenStrError::FmtError)?;

        Ok(fen_string)
    }
}

impl ToFenStr for Color {
    type Error = ToFenStrError;

    fn to_fen_str(&self) -> Result<String, Self::Error> {
        match self {
            Color::White => Ok("w".to_string()),
            Color::Black => Ok("b".to_string()),
        }
    }
}

impl ToFenStr for Piece {
    type Error = ToFenStrError;

    fn to_fen_str(&self) -> Result<String, Self::Error> {
        let ascii: char = self.to_ascii();
        Ok(String::from(match self.color {
            Color::White => ascii.to_ascii_uppercase(),
            Color::Black => ascii.to_ascii_lowercase(),
        }))
    }
}

impl ToFenStr for PlacedPiece {
    type Error = ToFenStrError;

    fn to_fen_str(&self) -> Result<String, Self::Error> {
        self.piece().to_fen_str()
    }
}

impl FromFenStr for Board {
    type Error = FromFenStrError;

    fn from_fen_str(string: &str) -> Result<Self, Self::Error> {
        let mut board = Board::empty();

        let mut fields = string.split(' ');

        let placements = fields
            .next()
            .ok_or(FromFenStrError::MissingField(Field::Placements))?;
        let ranks = placements.split('/');

        for (rank, pieces) in Rank::ALL.iter().rev().zip(ranks) {
            let mut files = File::ALL.iter();
            for ch in pieces.chars() {
                if let Some(skip) = ch.to_digit(10) {
                    // TODO: Use advance_by() when it's available.
                    for _ in 0..skip {
                        files.next();
                    }

                    continue;
                }

                let file = files.next().ok_or(FromFenStrError::MissingPlacement)?;
                let piece = Piece::from_fen_str(&ch.to_string())?;

                let _ = board.pieces.place(
                    piece,
                    Square::from_file_rank(*file, *rank),
                    PlacePieceStrategy::default(),
                );
            }

            debug_assert_eq!(files.next(), None);
        }

        let active_color = Color::from_fen_str(
            fields
                .next()
                .ok_or(FromFenStrError::MissingField(Field::PlayerToMove))?,
        )?;
        board.active_color = active_color;

        let castling_rights = fields
            .next()
            .ok_or(FromFenStrError::MissingField(Field::CastlingRights))?;
        if castling_rights == "-" {
            board.castling_rights.revoke_all();
        } else {
            for ch in castling_rights.chars() {
                match ch {
                    'K' => board.castling_rights.grant(Color::White, Castle::KingSide),
                    'Q' => board.castling_rights.grant(Color::White, Castle::QueenSide),
                    'k' => board.castling_rights.grant(Color::Black, Castle::KingSide),
                    'q' => board.castling_rights.grant(Color::Black, Castle::QueenSide),
                    _ => return Err(FromFenStrError::InvalidValue),
                };
            }
        }

        let en_passant_square = fields
            .next()
            .ok_or(FromFenStrError::MissingField(Field::EnPassantSquare))?;
        if en_passant_square != "-" {
            let square = Square::from_algebraic_str(en_passant_square)
                .map_err(FromFenStrError::ParseSquareError)?;
            board.en_passant_target = Some(square);
        }

        let half_move_clock = fields
            .next()
            .ok_or(FromFenStrError::MissingField(Field::HalfMoveClock))?;
        let half_move_clock: u32 = half_move_clock
            .parse()
            .map_err(FromFenStrError::ParseIntError)?;
        board.half_move_clock = half_move_clock;

        let full_move_counter = fields
            .next()
            .ok_or(FromFenStrError::MissingField(Field::FullMoveCounter))?;
        let full_move_counter: u32 = full_move_counter
            .parse()
            .map_err(FromFenStrError::ParseIntError)?;
        board.full_move_number = full_move_counter;

        debug_assert_eq!(fields.next(), None);

        Ok(board)
    }
}

impl FromFenStr for Color {
    type Error = FromFenStrError;

    fn from_fen_str(string: &str) -> Result<Self, Self::Error> {
        if string.len() != 1 {
            return Err(FromFenStrError::InvalidValue);
        }

        match string.chars().take(1).next().unwrap() {
            'w' => Ok(Color::White),
            'b' => Ok(Color::Black),
            _ => Err(FromFenStrError::InvalidValue),
        }
    }
}

impl FromFenStr for Piece {
    type Error = FromFenStrError;

    fn from_fen_str(string: &str) -> Result<Self, Self::Error> {
        if string.len() != 1 {
            return Err(FromFenStrError::InvalidValue);
        }

        match string.chars().take(1).next().unwrap() {
            'P' => Ok(piece!(White Pawn)),
            'N' => Ok(piece!(White Knight)),
            'B' => Ok(piece!(White Bishop)),
            'R' => Ok(piece!(White Rook)),
            'Q' => Ok(piece!(White Queen)),
            'K' => Ok(piece!(White King)),
            'p' => Ok(piece!(Black Pawn)),
            'n' => Ok(piece!(Black Knight)),
            'b' => Ok(piece!(Black Bishop)),
            'r' => Ok(piece!(Black Rook)),
            'q' => Ok(piece!(Black Queen)),
            'k' => Ok(piece!(Black King)),
            _ => Err(FromFenStrError::InvalidValue),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_board;

    #[test]
    fn starting_position() {
        let pos = test_board!(starting);

        assert_eq!(
            pos.to_fen_str().unwrap(),
            "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 0"
        );
    }

    #[test]
    fn from_starting_fen() {
        let board = fen!("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 0").unwrap();
        let expected = Board::starting();
        assert_eq!(board, expected, "{board:#?}\n{expected:#?}");
    }
}