chessfriend/position/src/fen.rs

// Eryn Wells <eryn@erynwells.me>

use crate::{r#move::Castle, Position};
use chessfriend_core::{Color, File, Piece, PlacedPiece, Rank, Square};
use std::fmt::Write;

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

pub trait ToFen {
    type Error;
    fn to_fen(&self) -> Result<String, Self::Error>;
}

impl ToFen for Position {
    type Error = ToFenError;

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

        let mut empty_squares: u8 = 0;
        for rank in Rank::ALL.iter().rev() {
            for file in File::ALL.iter() {
                let square = Square::from_file_rank(*file, *rank);
                match self.piece_on_square(square) {
                    Some(piece) => {
                        if empty_squares > 0 {
                            write!(fen_string, "{}", empty_squares)
                                .map_err(|err| ToFenError::FmtError(err))?;
                            empty_squares = 0;
                        }
                        write!(fen_string, "{}", piece.to_fen()?)
                            .map_err(|err| ToFenError::FmtError(err))?;
                    }
                    None => empty_squares += 1,
                }
            }

            if empty_squares > 0 {
                write!(fen_string, "{}", empty_squares).map_err(|err| ToFenError::FmtError(err))?;
                empty_squares = 0;
            }
            if rank != &Rank::ONE {
                write!(fen_string, "/").map_err(|err| ToFenError::FmtError(err))?;
            }
        }

        write!(fen_string, " {}", self.player_to_move().to_fen()?)
            .map_err(|err| ToFenError::FmtError(err))?;

        let castling = [
            (Color::White, Castle::KingSide),
            (Color::White, Castle::QueenSide),
            (Color::Black, Castle::KingSide),
            (Color::Black, Castle::QueenSide),
        ]
        .map(|(color, castle)| {
            let can_castle = self.player_has_right_to_castle(color, castle);
            if !can_castle {
                ""
            } else {
                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.len() > 0 { &castling } else { "-" }
        )
        .map_err(|err| ToFenError::FmtError(err))?;

        write!(
            fen_string,
            " {}",
            if let Some(en_passant_square) = self.en_passant_square() {
                en_passant_square.to_string()
            } else {
                "-".to_string()
            }
        )
        .map_err(|err| ToFenError::FmtError(err))?;

        write!(fen_string, " {}", self.ply_counter()).map_err(|err| ToFenError::FmtError(err))?;
        write!(fen_string, " {}", self.move_number()).map_err(|err| ToFenError::FmtError(err))?;

        Ok(fen_string)
    }
}

impl ToFen for Color {
    type Error = ToFenError;

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

impl ToFen for Piece {
    type Error = ToFenError;

    fn to_fen(&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 ToFen for PlacedPiece {
    type Error = ToFenError;

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

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

    #[test]
    fn starting_position() {
        let pos = Position::starting();
        println!("{pos:#?}");
        assert_eq!(
            pos.to_fen(),
            Ok(String::from(
                "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
            ))
        );
    }
}