eschac/src/setup.rs

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

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

/// **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 the [`Setup::validate`] method
/// before generating moves.
///
/// This type implements [`FromStr`](std::str::FromStr) and [`Display`](std::fmt::Display) to parse
/// and print positions from text records.
///
/// 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 string
/// (i.e. the halfmove clock and the fullmove number) as the [`Position`] type does not keep
/// track of those.
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct Setup {
    pub(crate) w: Bitboard,

    pub(crate) p_b_q: Bitboard,
    pub(crate) n_b_k: Bitboard,
    pub(crate) r_q_k: Bitboard,

    pub(crate) turn: Color,
    pub(crate) en_passant: OptionSquare,
    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 {
            w: Bitboard(0),
            p_b_q: Bitboard(0),
            n_b_k: Bitboard(0),
            r_q_k: Bitboard(0),
            turn: Color::White,
            en_passant: OptionSquare::None,
            castling_rights: CastlingRights::new(),
        }
    }

    /// Reads a position from an ascii record.
    pub fn from_ascii(s: &[u8]) -> Result<Self, ParseSetupError> {
        let mut s = s.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::new(File::transmute(file), Rank::transmute(rank)) },
                        Some(Piece { role, color }),
                    );
                    file += 1;
                    accept_empty_square = true;
                }
            }
            (rank == 0).then_some(())?;
            (file == 8).then_some(())?;
            Some(())
        })()
        .ok_or(ParseSetupError::InvalidBoard)?;
        (|| {
            match s.next()? {
                b'w' => setup.set_turn(Color::White),
                b'b' => setup.set_turn(Color::Black),
                _ => return None,
            }
            (s.next()? == b' ').then_some(())
        })()
        .ok_or(ParseSetupError::InvalidTurn)?;
        (|| {
            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(())
        })()
        .ok_or(ParseSetupError::InvalidCastlingRights)?;
        (|| {
            match s.next()? {
                b'-' => (),
                file => setup.set_en_passant_target_square(Some(Square::new(
                    File::from_ascii(file)?,
                    Rank::from_ascii(s.next()?)?,
                ))),
            }
            s.next().is_none().then_some(())
        })()
        .ok_or(ParseSetupError::InvalidEnPassantTargetSquare)?;
        Ok(setup)
    }

    /// Returns the occupancy of a square.
    #[inline]
    pub fn get(&self, square: Square) -> Option<Piece> {
        Some(Piece {
            role: self.get_role(square)?,
            color: match (self.w & square.bitboard()).is_empty() {
                false => Color::White,
                true => Color::Black,
            },
        })
    }

    /// 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.try_into_square()
    }

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

    /// 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 {
        Self {
            w: (self.w ^ (self.p_b_q | self.n_b_k | self.r_q_k)).mirror(),
            p_b_q: self.p_b_q.mirror(),
            n_b_k: self.n_b_k.mirror(),
            r_q_k: self.r_q_k.mirror(),
            turn: !self.turn,
            en_passant: self
                .en_passant
                .try_into_square()
                .map(|square| OptionSquare::from_square(square.mirror()))
                .unwrap_or(OptionSquare::None),
            castling_rights: self.castling_rights.mirror(),
        }
    }

    /// Tries to validate the position, i.e. converting it to a [`Position`].
    ///
    /// See [`IllegalPositionReason`] for details.
    pub fn validate(self) -> Result<Position, IllegalPosition> {
        debug_assert!((self.w & !(self.p_b_q | self.n_b_k | self.r_q_k)).is_empty());
        debug_assert!((self.p_b_q & self.n_b_k & self.r_q_k).is_empty());

        let mut reasons = IllegalPositionReasons::new();
        let d = InitialisedLookup::init();

        let blockers = self.p_b_q | self.n_b_k | self.r_q_k;
        let pieces = self.bitboards();

        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);
            !(d.king(enemy_king) & *pieces.get(Role::King)
                | d.bishop(enemy_king, blockers)
                    & (*pieces.get(Role::Queen) | *pieces.get(Role::Bishop))
                | d.rook(enemy_king, blockers)
                    & (*pieces.get(Role::Queen) | *pieces.get(Role::Rook))
                | d.knight(enemy_king) & *pieces.get(Role::Knight)
                | d.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::new(File::E, color.home_rank()))
                        && pieces
                            .get(color)
                            .get(Role::Rook)
                            .contains(Square::new(side.rook_origin_file(), color.home_rank())))
            })
        }) {
            reasons.add(IllegalPositionReason::InvalidCastlingRights);
        }

        if self.en_passant.try_into_square().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::new(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.try_into_square().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| {
                    !(d.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 get_role(&self, square: Square) -> Option<Role> {
        let mask = square.bitboard();
        let bit0 = (self.p_b_q & mask).0 >> square as u8;
        let bit1 = (self.n_b_k & mask).0 >> square as u8;
        let bit2 = (self.r_q_k & mask).0 >> square as u8;
        match bit0 | bit1 << 1 | bit2 << 2 {
            0 => None,
            i => Some(unsafe { Role::transmute(i as u8) }),
        }
    }

    #[inline]
    pub(crate) fn bitboards(&self) -> ByColor<ByRole<Bitboard>> {
        let Self {
            w,
            p_b_q,
            n_b_k,
            r_q_k,
            ..
        } = self.clone();
        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;
        ByColor::new(|color| {
            let mask = match color {
                Color::White => w,
                Color::Black => !w,
            };
            ByRole::new(|kind| {
                mask & match kind {
                    Role::King => k,
                    Role::Queen => q,
                    Role::Bishop => b,
                    Role::Knight => n,
                    Role::Rook => r,
                    Role::Pawn => p,
                }
            })
        })
    }
}

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

impl std::fmt::Display for Setup {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        use std::fmt::Write;

        for rank in Rank::all().into_iter().rev() {
            let mut count = 0;
            for file in File::all() {
                match self.get(Square::new(file, rank)) {
                    Some(piece) => {
                        if count > 0 {
                            f.write_char(char::from_u32('0' as u32 + count).unwrap())?;
                        }
                        count = 0;
                        f.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 {
                f.write_char(char::from_u32('0' as u32 + count).unwrap())?;
            }
            if rank != Rank::First {
                f.write_char('/')?;
            }
        }

        f.write_char(' ')?;

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

        f.write_char(' ')?;

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

        f.write_char(' ')?;

        match self.en_passant.try_into_square() {
            Some(sq) => {
                f.write_str(sq.to_str())?;
            }
            None => {
                write!(f, "-")?;
            }
        }

        Ok(())
    }
}

impl std::str::FromStr for Setup {
    type Err = ParseSetupError;
    #[inline]
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Self::from_ascii(s.as_bytes())
    }
}

/// An error when trying to parse a position record.
///
/// The variant indicates the field that caused the error.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ParseSetupError {
    InvalidBoard,
    InvalidTurn,
    InvalidCastlingRights,
    InvalidEnPassantTargetSquare,
}
impl std::fmt::Display for ParseSetupError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let details = match self {
            Self::InvalidBoard => "board",
            Self::InvalidTurn => "turn",
            Self::InvalidCastlingRights => "castling rights",
            Self::InvalidEnPassantTargetSquare => "en passant target square",
        };
        write!(f, "invalid text record ({details})")
    }
}
impl std::error::Error for ParseSetupError {}

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

impl IllegalPosition {
    /// Returns an iterator over the reasons why the position is rejected.
    pub fn reasons(&self) -> IllegalPositionReasons {
        self.reasons
    }

    /// Returns the [`Setup`] that failed validation.
    pub fn as_setup(&self) -> &Setup {
        &self.setup
    }

    /// Returns the [`Setup`] that failed validation.
    pub fn into_setup(self) -> Setup {
        self.setup
    }
}

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

impl std::fmt::Debug for IllegalPositionReasons {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::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 { std::mem::transmute::<u8, IllegalPositionReason>(reason) })
        }
    }
}

/// Reasons for illegal positions 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`].
    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 std::fmt::Display for IllegalPositionReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_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",
        })
    }
}