eschac 0.3.0

//! Building chess positions.
//!
//! [`Setup`] is a builder for the [`Position`] type.

use crate::bitboard::*;
use crate::board::*;
use crate::lookup;
use crate::position::*;

use alloc::string::String;

/// **A builder type for chess positions.**
///
/// This type is useful to edit a position without having to ensure it stays legal at every step.
/// It must be validated and converted to a [`Position`] using [`Setup::into_position`] before
/// generating moves.
///
/// ## Text description
///
/// Forsyth-Edwards Notation (FEN) is typically used to describe chess positions as text. eschac
/// uses a slightly different notation, which simply removes the last two fields of the FEN record
/// (i.e. the halfmove clock and the fullmove number) as the [`Position`] type does not keep
/// track of those. For example, the starting position is recorded as
/// `rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq -`.
///
/// [`Display`](core::fmt::Display) and [`FromStr`](alloc::str::FromStr) are purposely not
/// implemented on [`Setup`] as there does not exist a truely canonical format for writing chess
/// positions. The format described above is implemented in
/// [`from_text_record`](Setup::from_text_record) and [`to_text_record`](Setup::to_text_record).
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct Setup {
    /// bitboards = [pawns | bishops | queens, knights | bishops | kings, rooks | queens | kings, black]
    pub(crate) bitboards: [Bitboard; 4],
    pub(crate) turn: Color,
    pub(crate) en_passant: Option<Square>,
    pub(crate) castling_rights: CastlingRights,
}

impl Setup {
    /// Creates an empty board, i.e. `8/8/8/8/8/8/8/8 w - -`.
    #[inline]
    pub fn new() -> Self {
        Self {
            bitboards: [Bitboard(0); 4],
            turn: Color::White,
            en_passant: None,
            castling_rights: CastlingRights::new(),
        }
    }

    /// Reads a position from a text record
    /// (see [`from_ascii_record`](Setup::from_ascii_record)).
    #[inline]
    pub fn from_text_record(record: &str) -> Result<Self, ParseRecordError> {
        Self::from_ascii_record(record.as_bytes())
    }

    /// Reads a position from a text record.
    pub fn from_ascii_record(record: &[u8]) -> Result<Self, ParseRecordError> {
        let mut s = record.iter().copied().peekable();
        let mut setup = Setup::new();
        (|| {
            let mut accept_empty_square = true;
            let mut rank: u8 = 7;
            let mut file: u8 = 0;
            for c in s.by_ref() {
                if c == b'/' {
                    (file == 8).then_some(())?;
                    rank = rank.checked_sub(1)?;
                    file = 0;
                    accept_empty_square = true;
                } else if (b'1'..=b'8').contains(&c) && accept_empty_square {
                    file = file + c - b'0';
                    (file <= 8).then_some(())?;
                    accept_empty_square = false;
                } else if c == b' ' {
                    break;
                } else {
                    let role = Role::from_ascii(c)?;
                    let color = match c.is_ascii_uppercase() {
                        true => Color::White,
                        false => Color::Black,
                    };
                    (file < 8).then_some(())?;
                    setup.set(
                        unsafe {
                            Square::from_coords(
                                File::new_unchecked(file),
                                Rank::new_unchecked(rank),
                            )
                        },
                        Some(Piece { role, color }),
                    );
                    file += 1;
                    accept_empty_square = true;
                }
            }
            (rank == 0).then_some(())?;
            (file == 8).then_some(())?;

            match s.next()? {
                b'w' => setup.set_turn(Color::White),
                b'b' => setup.set_turn(Color::Black),
                _ => return None,
            }
            (s.next()? == b' ').then_some(())?;

            if s.next_if_eq(&b'-').is_none() {
                if s.next_if_eq(&b'K').is_some() {
                    setup.set_castling_rights(Color::White, CastlingSide::Short, true);
                }
                if s.next_if_eq(&b'Q').is_some() {
                    setup.set_castling_rights(Color::White, CastlingSide::Long, true);
                }
                if s.next_if_eq(&b'k').is_some() {
                    setup.set_castling_rights(Color::Black, CastlingSide::Short, true);
                }
                if s.next_if_eq(&b'q').is_some() {
                    setup.set_castling_rights(Color::Black, CastlingSide::Long, true);
                }
            }
            (s.next()? == b' ').then_some(())?;

            match s.next()? {
                b'-' => (),
                file => setup.set_en_passant_target_square(Some(Square::from_coords(
                    File::from_ascii(file)?,
                    Rank::from_ascii(s.next()?)?,
                ))),
            }
            s.next().is_none().then_some(())?;
            Some(setup)
        })()
        .ok_or_else(|| ParseRecordError {
            byte: record.len() - s.len(),
        })
    }

    /// Returns the text record of the position
    /// (see [`write_text_record`](Setup::write_text_record)).
    pub fn to_text_record(&self) -> String {
        let mut record = String::with_capacity(81);
        self.write_text_record(&mut record).unwrap();
        record
    }

    /// Writes the text record of the position.
    pub fn write_text_record<W>(&self, w: &mut W) -> core::fmt::Result
    where
        W: core::fmt::Write,
    {
        for rank in Rank::all().into_iter().rev() {
            let mut count = 0;
            for file in File::all() {
                match self.get(Square::from_coords(file, rank)) {
                    Some(piece) => {
                        if count > 0 {
                            w.write_char(char::from_u32('0' as u32 + count).unwrap())?;
                        }
                        count = 0;
                        w.write_char(match piece.color {
                            Color::White => piece.role.to_char_uppercase(),
                            Color::Black => piece.role.to_char_lowercase(),
                        })?;
                    }
                    None => {
                        count += 1;
                    }
                }
            }
            if count > 0 {
                w.write_char(char::from_u32('0' as u32 + count).unwrap())?;
            }
            if rank != Rank::First {
                w.write_char('/')?;
            }
        }

        w.write_char(' ')?;

        w.write_char(match self.turn {
            Color::White => 'w',
            Color::Black => 'b',
        })?;

        w.write_char(' ')?;

        let mut no_castle_available = true;
        if self.castling_rights(Color::White, CastlingSide::Short) {
            w.write_char('K')?;
            no_castle_available = false;
        }
        if self.castling_rights(Color::White, CastlingSide::Long) {
            w.write_char('Q')?;
            no_castle_available = false;
        }
        if self.castling_rights(Color::Black, CastlingSide::Short) {
            w.write_char('k')?;
            no_castle_available = false;
        }
        if self.castling_rights(Color::Black, CastlingSide::Long) {
            w.write_char('q')?;
            no_castle_available = false;
        }
        if no_castle_available {
            w.write_char('-')?;
        }

        w.write_char(' ')?;

        match self.en_passant {
            Some(sq) => {
                w.write_char(sq.file().to_char())?;
                w.write_char(sq.rank().to_char())?;
            }
            None => {
                w.write_char('-')?;
            }
        }

        Ok(())
    }

    /// Returns the piece on the square, if any.
    #[inline]
    pub fn get(&self, square: Square) -> Option<Piece> {
        Some(Piece {
            role: self.role(square)?,
            color: self.color(square),
        })
    }

    /// Returns the color to play.
    #[inline]
    pub fn turn(&self) -> Color {
        self.turn
    }

    /// Returns `true` if castling is available for the given color and side.
    #[inline]
    pub fn castling_rights(&self, color: Color, side: CastlingSide) -> bool {
        self.castling_rights.get(color, side)
    }

    /// Returns the optional en passant target square.
    #[inline]
    pub fn en_passant_target_square(&self) -> Option<Square> {
        self.en_passant
    }

    /// Sets the occupancy of a square.
    #[inline]
    pub fn set(&mut self, square: Square, piece: Option<Piece>) {
        let mask = !square.bitboard();
        let [p_b_q, n_b_k, r_q_k, black] = &mut self.bitboards;
        *p_b_q &= mask;
        *n_b_k &= mask;
        *r_q_k &= mask;
        *black &= mask;
        if let Some(piece) = piece {
            let to = square.bitboard();
            match piece.role {
                Role::Pawn => {
                    *p_b_q |= to;
                }
                Role::Knight => {
                    *n_b_k |= to;
                }
                Role::Bishop => {
                    *p_b_q |= to;
                    *n_b_k |= to;
                }
                Role::Rook => {
                    *r_q_k |= to;
                }
                Role::Queen => {
                    *p_b_q |= to;
                    *r_q_k |= to;
                }
                Role::King => {
                    *n_b_k |= to;
                    *r_q_k |= to;
                }
            }
            match piece.color {
                Color::White => (),
                Color::Black => *black |= to,
            }
        }
    }

    /// Sets the color to play.
    #[inline]
    pub fn set_turn(&mut self, color: Color) {
        self.turn = color;
    }

    /// Sets the castling rights for the given color and side.
    #[inline]
    pub fn set_castling_rights(&mut self, color: Color, side: CastlingSide, value: bool) {
        match value {
            true => self.castling_rights.set(color, side),
            false => self.castling_rights.unset(color, side),
        }
    }

    /// Sets the en passant target square.
    #[inline]
    pub fn set_en_passant_target_square(&mut self, square: Option<Square>) {
        self.en_passant = square;
    }

    /// Returns the quad-bitboard representation of the board.
    ///
    /// This returns, in order, the squares occpied by:
    /// - pawns, bishops and queens
    /// - knights, bishops and kings
    /// - rooks, queens and kings
    /// - black pieces
    #[inline]
    pub fn bitboards(&self) -> [Bitboard; 4] {
        self.bitboards
    }

    /// Returns the mirror image of the position.
    ///
    /// The mirror of a position is the position obtained after reflecting the placement of pieces
    /// horizontally, inverting the color of all the pieces, inverting the turn, and reflecting the
    /// castling rights as well as the en passant target square.
    ///
    /// For example, the mirror image of `rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b Kq e3`
    /// is `rnbqkbnr/pppp1ppp/8/4p3/8/8/PPPPPPPP/RNBQKBNR w Qk e6`.
    #[inline]
    pub fn mirror(&self) -> Self {
        let [p_b_q, n_b_k, r_q_k, black] = self.bitboards;
        Self {
            bitboards: [
                p_b_q.mirror(),
                n_b_k.mirror(),
                r_q_k.mirror(),
                (black ^ (p_b_q | n_b_k | r_q_k)).mirror(),
            ],
            turn: !self.turn,
            en_passant: self.en_passant.map(|square| square.mirror()),
            castling_rights: self.castling_rights.mirror(),
        }
    }

    /// Tries to convert the position into the [`Position`] type.
    ///
    /// Some unreachable positions are rejected, see [`IllegalPositionReason`] for details.
    pub fn into_position(self) -> Result<Position, IllegalPosition> {
        let mut reasons = IllegalPositionReasons::new();

        let [p_b_q, n_b_k, r_q_k, black] = self.bitboards;

        debug_assert!((!(p_b_q | n_b_k | r_q_k) & black).is_empty());
        debug_assert!((p_b_q & n_b_k & r_q_k).is_empty());

        let k = n_b_k & r_q_k;
        let q = p_b_q & r_q_k;
        let b = p_b_q & n_b_k;
        let n = n_b_k ^ b ^ k;
        let r = r_q_k ^ q ^ k;
        let p = p_b_q ^ b ^ q;
        let pieces = ByColor::with(|color| {
            let mask = match color {
                Color::White => !black,
                Color::Black => black,
            };
            ByRole::with(|role| {
                mask & match role {
                    Role::Pawn => p,
                    Role::Knight => n,
                    Role::Bishop => b,
                    Role::Rook => r,
                    Role::Queen => q,
                    Role::King => k,
                }
            })
        });

        let blockers = p_b_q | n_b_k | r_q_k;

        if Color::all()
            .into_iter()
            .any(|color| pieces.get(color).king().is_empty())
        {
            reasons.add(IllegalPositionReason::MissingKing);
        }

        if Color::all()
            .into_iter()
            .any(|color| pieces.get(color).get(Role::King).len() > 1)
        {
            reasons.add(IllegalPositionReason::TooManyKings);
        }

        if pieces.get(!self.turn).king().any(|enemy_king| {
            let pieces = pieces.get(self.turn);
            !(lookup::king(enemy_king) & *pieces.get(Role::King)
                | lookup::bishop(enemy_king, blockers)
                    & (*pieces.get(Role::Queen) | *pieces.get(Role::Bishop))
                | lookup::rook(enemy_king, blockers)
                    & (*pieces.get(Role::Queen) | *pieces.get(Role::Rook))
                | lookup::knight(enemy_king) & *pieces.get(Role::Knight)
                | lookup::pawn_attack(!self.turn, enemy_king) & *pieces.get(Role::Pawn))
            .is_empty()
        }) {
            reasons.add(IllegalPositionReason::HangingKing);
        }

        if Color::all().into_iter().any(|color| {
            !(*pieces.get(color).get(Role::Pawn)
                & (Rank::First.bitboard() | Rank::Eighth.bitboard()))
            .is_empty()
        }) {
            reasons.add(IllegalPositionReason::PawnOnBackRank);
        }

        if Color::all().into_iter().any(|color| {
            let dark_squares = Bitboard(0xAA55AA55AA55AA55);
            let light_squares = Bitboard(0x55AA55AA55AA55AA);
            let pieces = pieces.get(color);
            pieces.get(Role::Pawn).len()
                + pieces.get(Role::Queen).len().saturating_sub(1)
                + (*pieces.get(Role::Bishop) & dark_squares)
                    .len()
                    .saturating_sub(1)
                + (*pieces.get(Role::Bishop) & light_squares)
                    .len()
                    .saturating_sub(1)
                + pieces.get(Role::Knight).len().saturating_sub(2)
                + pieces.get(Role::Rook).len().saturating_sub(2)
                > 8
        }) {
            reasons.add(IllegalPositionReason::TooMuchMaterial);
        }

        if Color::all().into_iter().any(|color| {
            CastlingSide::all().into_iter().any(|side| {
                self.castling_rights.get(color, side)
                    && !(pieces
                        .get(color)
                        .get(Role::King)
                        .contains(Square::from_coords(File::E, color.home_rank()))
                        && pieces
                            .get(color)
                            .get(Role::Rook)
                            .contains(Square::from_coords(
                                side.rook_origin_file(),
                                color.home_rank(),
                            )))
            })
        }) {
            reasons.add(IllegalPositionReason::InvalidCastlingRights);
        }

        if self.en_passant.is_some_and(|en_passant| {
            let (target_rank, pawn_rank) = match self.turn {
                Color::White => (Rank::Sixth, Rank::Fifth),
                Color::Black => (Rank::Third, Rank::Fourth),
            };
            let pawn_square = Square::from_coords(en_passant.file(), pawn_rank);
            en_passant.rank() != target_rank
                || blockers.contains(en_passant)
                || !pieces.get(!self.turn).get(Role::Pawn).contains(pawn_square)
        }) {
            reasons.add(IllegalPositionReason::InvalidEnPassant);
        }

        if self.en_passant.is_some_and(|en_passant| {
            let blockers = blockers
                & !en_passant.bitboard().trans(match self.turn {
                    Color::White => Direction::South,
                    Color::Black => Direction::North,
                });
            pieces
                .get(self.turn)
                .king()
                .first()
                .is_some_and(|king_square| {
                    !(lookup::bishop(king_square, blockers)
                        & (pieces.get(!self.turn).queen() | pieces.get(!self.turn).bishop()))
                    .is_empty()
                })
        }) {
            reasons.add(IllegalPositionReason::ImpossibleEnPassantPin);
        }

        if reasons.0 != 0 {
            Err(IllegalPosition {
                setup: self,
                reasons,
            })
        } else {
            Ok(unsafe { Position::from_setup(self) })
        }
    }

    #[inline]
    pub(crate) fn role(&self, square: Square) -> Option<Role> {
        let mask = square.bitboard();
        let [p_b_q, n_b_k, r_q_k, _] = self.bitboards;
        let bit0 = (p_b_q & mask).0 >> square as u8;
        let bit1 = (n_b_k & mask).0 >> square as u8;
        let bit2 = (r_q_k & mask).0 >> square as u8;
        match bit0 | bit1 << 1 | bit2 << 2 {
            0 => None,
            i => Some(unsafe { Role::transmute(i as u8) }),
        }
    }

    /// Returns the color of the piece on the square if any, and `Black` if the square is not occupied.
    #[inline]
    pub(crate) fn color(&self, square: Square) -> Color {
        let [_, _, _, black] = self.bitboards;
        unsafe { Color::new_unchecked(((black & square.bitboard()).0 >> square as u8) as u8) }
    }
}

pub(crate) struct TextRecord<'a>(pub(crate) &'a Setup);
impl<'a> core::fmt::Display for TextRecord<'a> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        self.0.write_text_record(f)
    }
}
impl<'a> core::fmt::Debug for TextRecord<'a> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        use core::fmt::Write;
        f.write_char('"')?;
        self.0.write_text_record(f)?;
        f.write_char('"')?;
        Ok(())
    }
}

impl core::fmt::Debug for Setup {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
        f.debug_tuple("Setup").field(&TextRecord(&self)).finish()
    }
}

impl Default for Setup {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

/// An invalid position.
///
/// This is an illegal position that can't be represented with the [`Position`] type.
#[derive(Debug)]
pub struct IllegalPosition {
    pub setup: Setup,
    pub reasons: IllegalPositionReasons,
}
impl core::fmt::Display for IllegalPosition {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        use core::fmt::Write;
        let setup = &self.setup;
        write!(f, "`{}` is illegal:", &TextRecord(setup))?;
        let mut first = true;
        for reason in self.reasons {
            if !first {
                f.write_char(',')?;
            }
            first = false;
            f.write_char(' ')?;
            f.write_str(reason.to_str())?;
        }
        Ok(())
    }
}
impl core::error::Error for IllegalPosition {}

/// A set of [`IllegalPositionReason`]s.
#[derive(Clone, Copy)]
pub struct IllegalPositionReasons(u8);

impl core::fmt::Debug for IllegalPositionReasons {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_list().entries(*self).finish()
    }
}

impl IllegalPositionReasons {
    /// Returns `true` if the given reason appears in the set.
    pub fn contains(&self, reason: IllegalPositionReason) -> bool {
        (self.0 & reason as u8) != 0
    }

    fn new() -> Self {
        IllegalPositionReasons(0)
    }

    fn add(&mut self, reason: IllegalPositionReason) {
        self.0 |= reason as u8;
    }
}

impl Iterator for IllegalPositionReasons {
    type Item = IllegalPositionReason;
    fn next(&mut self) -> Option<Self::Item> {
        if self.0 == 0 {
            None
        } else {
            let reason = 1 << self.0.trailing_zeros();
            self.0 &= !reason;
            Some(unsafe { core::mem::transmute::<u8, IllegalPositionReason>(reason) })
        }
    }
}

/// Reasons for an illegal position to be rejected by eschac.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum IllegalPositionReason {
    /// One of the colors misses its king.
    MissingKing = 1,
    /// There is more than one king of the same color.
    TooManyKings = 2,
    /// The opponent's king is in check.
    HangingKing = 4,
    /// There is a pawn on the first or eighth rank.
    PawnOnBackRank = 8,
    /// Some castling rights are invalid regarding the positions of the rooks and kings.
    InvalidCastlingRights = 16,
    /// The en passant target square is invalid, either because:
    /// - it is not on the correct rank
    /// - it is occupied
    /// - it is not behind an opponent's pawn
    InvalidEnPassant = 32,
    /// There is an impossible number of pieces.
    ///
    /// Enforcing this enables to put an upper limit on the number of legal moves on any position,
    /// allowing to reduce the size of [`Moves`](crate::moves::Moves).
    TooMuchMaterial = 64,
    /// The pawn that can be taken en passant is pinned diagonally to the playing king.
    ///
    /// This can't happen on a legal position, as it would imply that the king could have be taken
    /// on that move. Enforcing this makes it unnecessary to test for a discovery check on the
    /// diagonal when taking en passant.
    ImpossibleEnPassantPin = 128,
}

impl IllegalPositionReason {
    fn to_str(self) -> &'static str {
        match self {
            Self::MissingKing => "missing king",
            Self::TooManyKings => "too many kings",
            Self::HangingKing => "hanging king",
            Self::PawnOnBackRank => "pawn on back rank",
            Self::InvalidCastlingRights => "invalid castling rights",
            Self::InvalidEnPassant => "invalid en passant",
            Self::TooMuchMaterial => "too much material",
            Self::ImpossibleEnPassantPin => "illegal en passant",
        }
    }
}

impl core::fmt::Display for IllegalPositionReason {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.write_str(self.to_str())
    }
}

/// An error when trying to parse a position record.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ParseRecordError {
    /// The index where the error was detected.
    pub byte: usize,
}
impl core::fmt::Display for ParseRecordError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "invalid text record (at byte {})", self.byte)
    }
}
impl core::error::Error for ParseRecordError {}