chess/rs/src/board.rs

use std::collections::HashMap;

#[derive(Eq, PartialEq, Clone)]
pub enum Color {
    Black,
    White,
}

impl Color {
    fn get_home_rank(&self) -> Rank {
        match self {
            Color::Black => Rank::_8,
            Color::White => Rank::_1,
        }
    }
    fn get_pawn_rank(&self) -> Rank {
        match self {
            Color::Black => Rank::_7,
            Color::White => Rank::_2,
        }
    }
    fn get_direction(&self) -> i8 {
        match self {
            Color::Black => -1,
            Color::White => 1,
        }
    }
}

#[derive(PartialEq, Eq, Hash, Clone)]
pub struct Position {
    pub rank: Rank,
    pub file: File,
}

impl Position {
    pub fn new(rank: Rank, file: File) -> Position {
        Position {
            rank,
            file,
        }
    }
    fn delta(&self, delta_rank_file: (i8, i8)) -> Result<Position, ()> {
        let (delta_rank, delta_file) = delta_rank_file;
        Ok(Position::new(
            self.rank.delta(delta_rank)?,
            self.file.delta(delta_file)?,
        ))
    }
}

#[derive(PartialEq, Eq, Hash, Clone, Copy)]
pub enum Rank {
    _1,
    _2,
    _3,
    _4,
    _5,
    _6,
    _7,
    _8,
}

#[derive(PartialEq, Eq, Hash, Clone, Copy)]
pub enum File {
    A,
    B,
    C,
    D,
    E,
    F,
    G,
    H,
}

pub trait GridAxis
where
    Self: Sized,
    Self: Clone,
{
    const ALL_VALUES: [Self; 8];
    fn new(index: u8) -> Result<Self, ()> {
        Ok(Self::ALL_VALUES[Self::validate_index(index)? as usize].clone())
    }
    fn validate_index(index: u8) -> Result<u8, ()> {
        if index as usize >= Self::ALL_VALUES.len() {
            Err(())
        } else {
            Ok(index)
        }
    }
    fn get_index(&self) -> u8;
    fn delta(&self, delta: i8) -> Result<Self, ()> {
        let highest = Self::ALL_VALUES.len() as i8 - 1;
        if delta > highest
            || delta < -highest
            || (delta < 0 && self.get_index() < delta.abs() as u8)
            || (delta > 0 && self.get_index() + delta as u8 > highest as u8)
        {
            Err(())
        } else {
            Self::new(((self.get_index() as i8) + delta) as u8)
        }
    }
}

impl GridAxis for Rank {
    const ALL_VALUES: [Rank; 8] = [
        Rank::_1,
        Rank::_2,
        Rank::_3,
        Rank::_4,
        Rank::_5,
        Rank::_6,
        Rank::_7,
        Rank::_8,
    ];
    fn get_index(&self) -> u8 {
        *self as u8
    }
}

impl GridAxis for File {
    const ALL_VALUES: [File; 8] = [
        File::A,
        File::B,
        File::C,
        File::D,
        File::E,
        File::F,
        File::G,
        File::H,
    ];
    fn get_index(&self) -> u8 {
        *self as u8
    }
}

#[derive(Clone)]
pub enum PieceType {
    Pawn,
}

impl PieceType {
    fn initial_setup(&self) -> Vec<(Piece, Position)> {
        match &self {
            PieceType::Pawn => [Color::Black, Color::White]
                .iter()
                .flat_map(|color| {
                    File::ALL_VALUES.iter().map(|&file| {
                        (
                            Piece {
                                piece_type: self.clone(),
                                color: color.clone(),
                            },
                            Position::new(color.get_pawn_rank(), file),
                        )
                    })
                })
                .collect(),
        }
    }
}

pub struct Piece {
    pub color: Color,
    pub piece_type: PieceType,
}

struct Ray {
    unit: (i8, i8),
    coefs: Option<Vec<u8>>,
    source: Position,
    current: u8,
}

impl Ray {
    fn new(unit: (i8, i8), coefs: Option<Vec<u8>>, source: Position) -> Self {
        Ray {
            unit,
            coefs,
            source,
            current: 0,
        }
    }
}

impl Iterator for Ray {
    type Item = Position;
    fn next(&mut self) -> Option<Position> {
        let (delta_unit_rank, delta_unit_file) = self.unit;
        let coef = i8::try_from(match &self.coefs {
            None => self.current + 1,
            Some(coefs) => {
                let crt_usize = self.current as usize;
                if crt_usize >= coefs.len() {
                    return None;
                }
                coefs[crt_usize]
            }
        })
        .unwrap();
        let delta = (
            delta_unit_rank.checked_mul(coef).unwrap(),
            delta_unit_file.checked_mul(coef).unwrap(),
        );
        let rval = self.source.delta(delta).ok();
        self.current = self.current.wrapping_add(1);
        rval
    }
}

impl Piece {
    fn get_move_rays(&self, position: &Position) -> Vec<Ray> {
        match self.piece_type {
            PieceType::Pawn => {
                let direction = self.color.get_direction();
                let coefs = if position.rank == self.color.get_pawn_rank() {
                    vec![1, 2]
                } else {
                    vec![1]
                };
                vec![Ray::new((direction, 0), Some(coefs), position.clone())]
            }
        }
    }
    fn get_capture_rays(&self, position: &Position) -> Vec<Ray> {
        match self.piece_type {
            PieceType::Pawn => {
                let direction = self.color.get_direction();
                vec![
                    Ray::new((direction, 1), Some(vec![1]), position.clone()),
                    Ray::new((direction, -1), Some(vec![1]), position.clone()),
                ]
            }
        }
    }
}

pub struct Board {
    state: HashMap<Position, Piece>,
}

impl Board {
    pub fn new() -> Board {
        Board {
            state: Default::default(),
        }
    }
    pub fn get_at(&self, position: &Position) -> Option<&Piece> {
        self.state.get(position)
    }
    pub fn set_at(
        &mut self,
        position: Position,
        piece_or_empty: Option<Piece>,
    ) {
        match piece_or_empty {
            Some(piece) => self.state.insert(position, piece),
            None => self.state.remove(&position),
        };
    }
    pub fn relocate(
        &mut self,
        source: &Position,
        target: Position,
    ) -> Result<(), ()> {
        let (_, piece) = self.state.remove_entry(source).ok_or(())?;
        self.set_at(target, Some(piece));
        Ok(())
    }
    pub fn reset(&mut self) {
        self.state.clear();
        for piece_type in &[PieceType::Pawn] {
            for (piece, position) in piece_type.initial_setup() {
                self.set_at(position, Some(piece));
            }
        }
    }
    pub fn iter(&self) -> impl Iterator<Item = (&Position, &Piece)> {
        self.state.iter()
    }
    fn cut_move_ray(&self, ray: Ray) -> impl Iterator<Item = Position> + '_ {
        ray.into_iter().take_while(|position| self.get_at(position).is_none())
    }
    fn find_capture_along_ray(&self, ray: Ray) -> Option<Position> {
        // I'm not expecting to call this if there's no piece in the position
        let crt_piece = self.get_at(&ray.source).unwrap();
        ray.into_iter().find(|p| match self.get_at(p) {
            None => false,
            Some(piece) => piece.color != crt_piece.color,
        })
    }
    pub fn find_moves(
        &self,
        position: &Position,
    ) -> Result<impl Iterator<Item = Position> + '_, ()> {
        if let Some(piece) = self.get_at(position) {
            Ok(piece
                .get_move_rays(position)
                .into_iter()
                .flat_map(|ray| self.cut_move_ray(ray)))
        } else {
            Err(())
        }
    }
    pub fn find_captures(
        &self,
        position: &Position,
    ) -> Result<impl Iterator<Item = Position> + '_, ()> {
        if let Some(piece) = self.get_at(position) {
            Ok(piece
                .get_capture_rays(position)
                .into_iter()
                .filter_map(|ray| self.find_capture_along_ray(ray)))
        } else {
            Err(())
        }
    }
}