crabchess 0.1.15

//! Tools for navigating and relating the squares on a chessboard.

use crate::{
    error::{Err, Result},
    pieces::{
        Color::{self, Black, White},
        LinearType,
        Side::{self, Kingside, Queenside},
    },
    sq,
};
use compact_str::CompactString;
use smallvec::SmallVec;
use std::{
    cmp::Ordering::{self, Equal, Greater, Less},
    fmt::{Debug, Display},
    str::FromStr,
};
use ufmt::{derive::uDebug, uDisplay, uwrite};
use Direction::{Down, DownLeft, DownRight, Left, Right, Up, UpLeft, UpRight};

const KNIGHT_SHIFTS: [(i8, i8); 8] = [
    (1, -2),
    (2, -1),
    (-1, 2),
    (-2, 1),
    (1, 2),
    (2, 1),
    (-1, -2),
    (-2, -1),
];

/// A file (column) on a chessboard.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum File {
    A,
    B,
    C,
    D,
    E,
    F,
    G,
    H,
}

impl File {
    /// Converts a character to its corresponding file.
    ///
    /// # Errors
    ///
    /// File char must be in `"ABCDEFGH"`.
    pub const fn from_char(ch: char) -> Result<Self> {
        match ch.to_ascii_uppercase() {
            'A' | 'a' => Ok(Self::A),
            'B' | 'b' => Ok(Self::B),
            'C' | 'c' => Ok(Self::C),
            'D' | 'd' => Ok(Self::D),
            'E' | 'e' => Ok(Self::E),
            'F' | 'f' => Ok(Self::F),
            'G' | 'g' => Ok(Self::G),
            'H' | 'h' => Ok(Self::H),
            _ => Err(Err::InvalidFileCharError(ch)),
        }
    }

    /// Get a piece's file from its FEN rank string if it is in the rank. If there is more than one
    /// piece, the piece closest to `File::A` will be returned.
    ///
    /// ```
    /// use crabchess::squares::File;
    /// assert_eq!(File::first_in_fen_rank("3R1K2", 'K').unwrap(), File::F);
    /// ```
    ///
    /// # Errors
    /// Returns an error if the piece character is not present in the string or the string is
    /// invalid.
    #[allow(clippy::cast_possible_truncation)]
    pub fn first_in_fen_rank(fen_rank: &str, piece_char: char) -> Result<Self> {
        let mut idx: i8 = 0;

        for c in fen_rank.chars() {
            match c.to_digit(10) {
                Some(n) => idx += n as i8,
                None if c == piece_char => break,
                _ => idx += 1,
            }
        }

        Self::from_idx(idx).ok_or_else(|| Err::RankFromIndexError(idx))
    }

    /// Get a piece's file from its FEN rank string if it is in the rank. If there is more than one
    /// piece, the piece closest to `File::H` will be returned.
    ///
    /// ```
    /// use crabchess::squares::File;
    /// assert_eq!(File::last_in_fen_rank("RNBQK2R", 'R').unwrap(), File::H);
    /// ```
    ///
    /// # Errors
    /// Returns an error if the piece character is not present in the string or the string is
    /// invalid.
    #[allow(clippy::cast_possible_truncation)]
    pub fn last_in_fen_rank(fen_rank: &str, piece_char: char) -> Result<Self> {
        let mut idx: i8 = 7;

        for c in fen_rank.chars().rev() {
            match c.to_digit(10) {
                Some(n) => idx -= n as i8,
                None if c == piece_char => break,
                _ => idx -= 1,
            }
        }

        Self::from_idx(idx).ok_or_else(|| Err::RankFromIndexError(idx))
    }

    /// Get a piece's file from its FEN rank string if it is in the rank. If there is more than one
    /// piece, the piece closest to `side` will be returned.
    ///
    /// ```
    /// use crabchess::{pieces::Side, squares::File};
    /// assert_eq!(File::closest_in_fen_rank("RNBQK2R", 'R', Side::Kingside).unwrap(), File::H);
    /// assert_eq!(File::closest_in_fen_rank("RNBQK2R", 'R', Side::Queenside).unwrap(), File::A);
    /// ```
    ///
    /// # Errors
    /// Returns an error if the piece character is not present in the string or the string is
    /// invalid.
    pub fn closest_in_fen_rank(fen_rank: &str, piece_char: char, side: Side) -> Result<Self> {
        match side {
            Kingside => Self::last_in_fen_rank(fen_rank, piece_char),
            Queenside => Self::first_in_fen_rank(fen_rank, piece_char),
        }
    }

    /// Get adjacent files.
    #[must_use]
    pub const fn adjacent(self) -> [Option<Self>; 2] {
        match self {
            Self::A => [Some(Self::B), None],
            Self::B => [Some(Self::A), Some(Self::C)],
            Self::C => [Some(Self::B), Some(Self::D)],
            Self::D => [Some(Self::C), Some(Self::E)],
            Self::E => [Some(Self::D), Some(Self::F)],
            Self::F => [Some(Self::E), Some(Self::G)],
            Self::G => [Some(Self::F), Some(Self::H)],
            Self::H => [Some(Self::G), None],
        }
    }

    /// Get files as array.
    #[must_use]
    pub const fn all() -> [Self; 8] {
        [
            Self::A,
            Self::B,
            Self::C,
            Self::D,
            Self::E,
            Self::F,
            Self::G,
            Self::H,
        ]
    }

    /// Get file character.
    #[must_use]
    pub const fn char(self) -> char {
        match self {
            Self::A => 'A',
            Self::B => 'B',
            Self::C => 'C',
            Self::D => 'D',
            Self::E => 'E',
            Self::F => 'F',
            Self::G => 'G',
            Self::H => 'H',
        }
    }

    /// Get lowercase file character.
    #[must_use]
    pub const fn lower(self) -> char {
        match self {
            Self::A => 'a',
            Self::B => 'b',
            Self::C => 'c',
            Self::D => 'd',
            Self::E => 'e',
            Self::F => 'f',
            Self::G => 'g',
            Self::H => 'h',
        }
    }

    /// Gets a file by its index.
    #[must_use]
    pub const fn from_idx(idx: i8) -> Option<Self> {
        match idx {
            0 => Some(Self::A),
            1 => Some(Self::B),
            2 => Some(Self::C),
            3 => Some(Self::D),
            4 => Some(Self::E),
            5 => Some(Self::F),
            6 => Some(Self::G),
            7 => Some(Self::H),
            _ => None,
        }
    }

    /// Steps `by` times right (if positive) or left (if negative).
    #[must_use]
    pub const fn step(self, by: i8) -> Option<Self> {
        Self::from_idx(self as i8 + by)
    }
}

/// A rank (row) on a chessboard.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, uDebug)]
pub enum Rank {
    One,
    Two,
    Three,
    Four,
    Five,
    Six,
    Seven,
    Eight,
}

impl Rank {
    /// Get rank from character.
    ///
    /// # Errors
    ///
    /// Returns an error if rank character is not in '12345678'.
    pub const fn from_char(ch: &char) -> Result<Self> {
        match ch {
            '1' => Ok(Self::One),
            '2' => Ok(Self::Two),
            '3' => Ok(Self::Three),
            '4' => Ok(Self::Four),
            '5' => Ok(Self::Five),
            '6' => Ok(Self::Six),
            '7' => Ok(Self::Seven),
            '8' => Ok(Self::Eight),
            _ => Err(Err::RankFromCharError(*ch)),
        }
    }

    /// Get ranks as an array.
    #[must_use]
    pub const fn all() -> [Self; 8] {
        [
            Self::One,
            Self::Two,
            Self::Three,
            Self::Four,
            Self::Five,
            Self::Six,
            Self::Seven,
            Self::Eight,
        ]
    }

    /// Get ranks as an array in reverse order.
    #[must_use]
    pub const fn reversed() -> [Self; 8] {
        [
            Self::Eight,
            Self::Seven,
            Self::Six,
            Self::Five,
            Self::Four,
            Self::Three,
            Self::Two,
            Self::One,
        ]
    }

    /// Get rank as character.
    #[must_use]
    pub const fn char(self) -> char {
        match self {
            Self::One => '1',
            Self::Two => '2',
            Self::Three => '3',
            Self::Four => '4',
            Self::Five => '5',
            Self::Six => '6',
            Self::Seven => '7',
            Self::Eight => '8',
        }
    }

    /// Gets rank by index.
    ///
    /// # Errors
    ///
    /// Rank index must be in range `0..=7`.
    #[must_use]
    pub const fn from_idx(idx: i8) -> Option<Self> {
        match idx {
            0 => Some(Self::One),
            1 => Some(Self::Two),
            2 => Some(Self::Three),
            3 => Some(Self::Four),
            4 => Some(Self::Five),
            5 => Some(Self::Six),
            6 => Some(Self::Seven),
            7 => Some(Self::Eight),
            _ => None,
        }
    }

    /// Steps `by` times up (if positive) or down (if negative).
    ///
    /// # Errors
    ///
    /// Errs if resultant rank is out of bounds.
    #[must_use]
    pub const fn step(self, by: i8) -> Option<Self> {
        Self::from_idx(self as i8 + by)
    }
}

/// The cardinal and ordinal directions.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, uDebug)]
pub enum Direction {
    Up,
    Down,
    Left,
    Right,
    UpLeft,
    UpRight,
    DownLeft,
    DownRight,
}

impl uDisplay for Direction {
    fn fmt<W>(&self, f: &mut ufmt::Formatter<'_, W>) -> std::prelude::v1::Result<(), W::Error>
    where
        W: ufmt::uWrite + ?Sized,
    {
        uwrite!(f, "{:?}", self)
    }
}

impl std::fmt::Display for Direction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{self:?}")
    }
}

impl Direction {
    /// Returns opposite `Direction`.
    #[must_use]
    pub const fn opposite(self) -> Self {
        match self {
            Up => Down,
            Down => Up,
            Left => Right,
            Right => Left,
            UpLeft => DownRight,
            UpRight => DownLeft,
            DownLeft => UpRight,
            DownRight => UpLeft,
        }
    }

    /// Returns `(f_shift, r_shift)` to step one square in direction.
    #[must_use]
    pub const fn index_modifier(self) -> (i8, i8) {
        match self {
            Up => (0, 1),
            Down => (0, -1),
            Left => (-1, 0),
            Right => (1, 0),
            UpLeft => (-1, 1),
            UpRight => (1, 1),
            DownLeft => (-1, -1),
            DownRight => (1, -1),
        }
    }

    /// Get direction a pawn can move by its color.
    #[must_use]
    pub const fn pawn_forward(color: Color) -> Self {
        match color {
            White => Up,
            Black => Down,
        }
    }

    /// Get directions a pawn can capture towards by its color.
    #[must_use]
    pub const fn pawn_diagonals(color: Color) -> [Self; 2] {
        match color {
            White => [UpLeft, UpRight],
            Black => [DownLeft, DownRight],
        }
    }

    /// Get directions as array.
    #[must_use]
    pub const fn all() -> [Self; 8] {
        [Up, Down, Left, Right, UpLeft, UpRight, DownLeft, DownRight]
    }

    /// Get directions a rook can travel in.
    #[must_use]
    pub const fn rooks() -> [Self; 4] {
        [Up, Down, Left, Right]
    }

    /// Get directions a bishop can travel in.
    #[must_use]
    pub const fn bishops() -> [Self; 4] {
        [UpLeft, UpRight, DownLeft, DownRight]
    }

    /// Get directions a queen can travel in.
    #[must_use]
    pub const fn queens() -> [Self; 8] {
        Self::all()
    }

    /// Get navigable directions by piece type.
    #[must_use]
    pub const fn from_piece_type(pt: LinearType) -> [Option<Self>; 8] {
        match pt {
            LinearType::Bishop => [
                Some(UpLeft),
                Some(UpRight),
                Some(DownLeft),
                Some(DownRight),
                None,
                None,
                None,
                None,
            ],
            LinearType::Rook => [
                Some(Up),
                Some(Down),
                Some(Left),
                Some(Right),
                None,
                None,
                None,
                None,
            ],
            LinearType::Queen => [
                Some(Up),
                Some(Down),
                Some(Left),
                Some(Right),
                Some(UpLeft),
                Some(UpRight),
                Some(DownLeft),
                Some(DownRight),
            ],
        }
    }

    /// Determine the direction between two squares.
    ///
    /// # Errors
    ///
    /// Returns an error if `square` and `other` are equal.
    pub fn between(square: Square, other: Square) -> Result<Self> {
        let file_diff: Ordering = (other.file() as i8).cmp(&(square.file() as i8));
        let rank_diff: Ordering = (other.rank() as i8).cmp(&(square.rank() as i8));

        match (file_diff, rank_diff) {
            (Equal, Greater) => Ok(Up),
            (Equal, Less) => Ok(Down),
            (Less, Equal) => Ok(Left),
            (Greater, Equal) => Ok(Right),
            (Less, Greater) => Ok(UpLeft),
            (Greater, Greater) => Ok(UpRight),
            (Less, Less) => Ok(DownLeft),
            (Greater, Less) => Ok(DownRight),
            _ => Err(Err::EqualSquaresError(square)),
        }
    }
}

/// A square on a chessboard, acting as an index for the `Board` struct.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct Square(pub File, pub Rank);

impl Debug for Square {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Display::fmt(self, f)
    }
}

impl Display for Square {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}{}", self.0.char(), self.1.char())
    }
}

impl uDisplay for Square {
    fn fmt<W: ufmt::uWrite + ?Sized>(
        &self,
        f: &mut ufmt::Formatter<'_, W>,
    ) -> std::prelude::v1::Result<(), W::Error> {
        uwrite!(f, "{}{}", self.0.char(), self.1.char())
    }
}

impl FromStr for Square {
    type Err = Err;

    /// Create a square from a string representation like `"a3"` or `"B4"`.
    ///
    /// # Errors
    ///
    /// Returns an error if file or rank character is missing or invalid.
    fn from_str(s: &str) -> Result<Self> {
        let mut iter = s.chars();

        let Some(file_char) = iter.next() else {
            return Err(Err::MissingFileCharError(s.into()));
        };
        let file = File::from_char(file_char)?;

        let Some(rank_char) = iter.next() else {
            return Err(Err::MissingRankCharError(s.into()));
        };
        let rank = Rank::from_char(&rank_char)?;

        Ok(Self(file, rank))
    }
}

impl TryFrom<&str> for Square {
    type Error = Err;

    fn try_from(value: &str) -> std::prelude::v1::Result<Self, Self::Error> {
        FromStr::from_str(value)
    }
}

impl Square {
    /// File shifts by `f_shift` and rank shifts by `r_shift`, returning a new `Square`.
    #[must_use]
    pub fn step(self, f_shift: i8, r_shift: i8) -> Option<Self> {
        Some(Self(self.0.step(f_shift)?, self.1.step(r_shift)?))
    }

    /// Get square in lowercase, for use in SAN notation.
    #[must_use]
    pub fn to_lowercase_string(self) -> CompactString {
        let mut s = CompactString::default();
        s.push(self.file().lower());
        s.push(self.rank().char());
        s
    }

    /// Get the color of a square on the board.
    #[must_use]
    pub const fn color(self) -> Color {
        match self.rank() {
            Rank::One | Rank::Three | Rank::Five | Rank::Seven => match self.file() {
                File::A | File::C | File::E | File::G => Black,
                File::B | File::D | File::F | File::H => White,
            },
            Rank::Two | Rank::Four | Rank::Six | Rank::Eight => match self.file() {
                File::A | File::C | File::E | File::G => White,
                File::B | File::D | File::F | File::H => Black,
            },
        }
    }

    /// Get a castling king's final square.
    #[must_use]
    pub const fn castling_king_final(color: Color, side: Side) -> Self {
        match (color, side) {
            (White, Kingside) => sq!(G1),
            (White, Queenside) => sq!(C1),
            (Black, Kingside) => sq!(G8),
            (Black, Queenside) => sq!(C8),
        }
    }

    /// Get a castling rook's final square.
    #[must_use]
    pub const fn castling_rook_final(color: Color, side: Side) -> Self {
        match (color, side) {
            (White, Kingside) => sq!(F1),
            (White, Queenside) => sq!(D1),
            (Black, Kingside) => sq!(F8),
            (Black, Queenside) => sq!(D8),
        }
    }

    /// Step `n` times in direction.
    ///
    /// # Errors
    ///
    /// Returns an error if resultant square is out of bounds.
    #[must_use]
    pub fn step_dir(self, direction: Direction, n: i8) -> Option<Self> {
        let (f_shift, r_shift) = direction.index_modifier();
        self.step(f_shift * n, r_shift * n)
    }

    /// Iterate squares in a direction on the board.
    #[must_use]
    pub const fn iter_dir(self, direction: Direction) -> DirectionIter {
        DirectionIter {
            square: self,
            direction,
        }
    }

    /// Get an exclusive range of squares between `self` and `other`.
    ///
    /// # Errors
    ///
    /// Returns an error if `self` is equal to `other` or if squares are not connected by a
    /// straight line.
    pub fn step_to(self, other: Self) -> Result<SmallVec<[Self; 6]>> {
        let direction = Direction::between(self, other)?;
        let mut touched = false;
        let mut result = SmallVec::new();

        for square in self.iter_dir(direction) {
            if square == other {
                touched = true;
                break;
            }
            result.push(square);
        }

        if touched {
            Ok(result)
        } else {
            Err(Err::SquaresNotInlineError(self, other))
        }
    }

    /// Squares a knight could navigate to.
    pub fn knight_navigable(&self) -> impl Iterator<Item = Self> + '_ {
        KNIGHT_SHIFTS
            .into_iter()
            .filter_map(|(f, r)| self.step(f, r))
    }

    /// Squares a king could navigate to.
    pub fn king_navigable(self) -> impl Iterator<Item = Self> + 'static {
        Direction::all()
            .into_iter()
            .filter_map(move |dir| self.step_dir(dir, 1))
    }

    /// Get file.
    #[must_use]
    pub const fn file(self) -> File {
        self.0
    }

    /// Get rank.
    #[must_use]
    pub const fn rank(self) -> Rank {
        self.1
    }

    /// Get final square from SAN notation.
    ///
    /// # Errors
    ///
    /// Returns an error if SAN is invalid or represents a castling move.
    pub fn final_square_from_san(san: &str) -> Result<Self> {
        match san.rfind(|c| "abcdefgh".contains(c)) {
            Some(idx) => Ok(Self::from_str(&san[idx..=idx + 1])?),
            None => Err(Err::MissingFinalSquareError(san.into())),
        }
    }

    /// Get en passant target/capture square from final square and color of pawn to be moved.
    ///
    /// For example, if a black pawn makes a double forward advance from `C7` to `C5`, the en
    /// passant final square (indicated in FEN notation) would be `C6`, and this method would
    /// return `C5` if given `C6` and `White` as parameters.
    #[must_use]
    pub fn en_passant_target_from_final(&self, color: Color) -> Option<Self> {
        let direction = match color.other() {
            White => Up,
            Black => Down,
        };

        self.step_dir(direction, 1)
    }

    /// Squares a pawn could capture to, assuming the square is occupied by an opponent's piece and
    /// the move is legal.
    #[must_use]
    pub fn pawn_diagonals(&self, color: Color) -> [Option<Self>; 2] {
        let rank = match color {
            White => self.rank().step(1),
            Black => self.rank().step(-1),
        };
        let Some(rank) = rank else {
            return [None, None];
        };

        self.file().adjacent().map(|f| Some(Self(f?, rank)))
    }

    /// Square(s) which a pawn of `color` could capture from.
    #[must_use]
    pub fn vulnerable_to_pawns(&self, color: Color) -> [Option<Self>; 2] {
        self.pawn_diagonals(color.other())
    }

    /// Square(s) navigable by a pawn forward advance (i.e. _without_ capture).
    #[must_use]
    pub fn pawn_forwards(self, color: Color) -> Vec<Self> {
        let direction = Direction::pawn_forward(color);
        let mut result = Vec::new();

        if let Some(square) = self.step_dir(direction, 1) {
            result.push(square);
        } else {
            return result;
        }

        if let (Rank::Two, White) | (Rank::Seven, Black) = (self.rank(), color) {
            if let Some(square) = self.step_dir(direction, 2) {
                result.push(square);
            }
        }

        result
    }

    /// Squares which a pawn of `color` could make a forward advance from.
    #[must_use]
    pub fn navigable_to_pawns(self, color: Color) -> [Option<Self>; 2] {
        let direction = Direction::pawn_forward(color.other());

        let first = if let Some(square) = self.step_dir(direction, 1) {
            Some(square)
        } else {
            return [None, None];
        };

        let second = if let (Rank::Four, White) | (Rank::Five, Black) = (self.rank(), color) {
            self.step_dir(direction, 2)
        } else {
            None
        };

        [first, second]
    }

    /// Get squares adjacent to `self`. For example, `C3` and `E3` are adjacent to `D3`.
    #[must_use]
    pub fn adjacent(self) -> SmallVec<[Self; 2]> {
        let mut output = SmallVec::new();

        let rank = self.rank();

        for file in self.file().adjacent() {
            let Some(file) = file else { break };
            output.push(Self(file, rank));
        }

        output
    }
}

/// Iterates by square in a direction until the end of the board is reached.
pub struct DirectionIter {
    square: Square,
    direction: Direction,
}

impl Iterator for DirectionIter {
    type Item = Square;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(square) = self.square.step_dir(self.direction, 1) {
            self.square = square;
            Some(square)
        } else {
            None
        }
    }
}

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

    #[test]
    fn file_index() {
        let c = File::C;
        let x = c.step(2).unwrap();
        assert_eq!(x, File::E);
        let x = c.step(6);
        assert!(x.is_none(), "Overshifting file did not return Err");
        let x = c.step(-2).unwrap();
        assert_eq!(x, File::A);
        let x = c.step(-3);
        assert!(x.is_none(), "Overshifting file did not return Err");
        let x = File::from_idx(3).unwrap();
        assert_eq!(x, File::D);
        let x = File::from_idx(8);
        assert!(x.is_none(), "Overshifting file did not return Err");
    }

    #[test]
    fn direction_between() {
        assert_eq!(Direction::between(sq!(A5), sq!(A3)).unwrap(), Down);
        assert_eq!(Direction::between(sq!(C1), sq!(E3)).unwrap(), UpRight);
        assert_eq!(Direction::between(sq!(D1), sq!(B1)).unwrap(), Left);
    }
}