path: root/src/attacks.rs

use crate::{bitboard::Bitboard, board::Color};

static NOT_A_FILE: Bitboard = 0xFEFEFEFEFEFEFEFE;
static NOT_B_FILE: Bitboard = 0xFDFDFDFDFDFDFDFD;
static NOT_G_FILE: Bitboard = 0xBFBFBFBFBFBFBFBF;
static NOT_H_FILE: Bitboard = 0x7F7F7F7F7F7F7F7F;
static B_FILE:     Bitboard = 0x0202020202020202;
static H_FILE:     Bitboard = 0x8080808080808080;
static DIAG_C2_H7: Bitboard = 0x0080402010080400;
static DIAG_A1_H8: Bitboard = 0x8040201008040201;
static RANK_2:     Bitboard = 0x000000000000FF00;
static RANK_6:     Bitboard = 0x0000FF0000000000;

/// An array where N-th item is an attack bitboard
/// of a piece on N-th square
type AttackTable = [Bitboard; 64];

/// One byte stores exactly 256 states which is every
/// possible state of the occupancy on a rank
pub type Byte = u8;

/// First index is *occupancy* of the first rank (all possible states)
/// and the second index is the file of the piece on the first rank
/// ```
/// rank_attacks = first_rank_attacks[occupancy][piece_file];
/// ```
type FirstRankAttacks = [[Byte; 8]; 256];

#[allow(dead_code)]
enum Direction {
    Nort,
    NoEa,
    East,
    SoEa,
    Sout,
    SoWe,
    West,
    NoWe,
}

#[derive(Debug)]
pub struct Attacks {
    pub knight: AttackTable,
    pub king: AttackTable,
    /// TODO: compute pawn attacks set-wise
    pub pawn: [AttackTable; 2],
    pub pawn_pushes: [AttackTable; 2],

    pub first_rank_attacks: FirstRankAttacks,
    /// Should be indexed by Direction
    pub ray_attacks: [AttackTable; 8],
}

#[allow(unused)]
impl Attacks {
    pub fn new() -> Self {
        let knight = Self::precompute_knight_attacks();
        let king = Self::precompute_king_attacks();
        let first_rank_attacks = Self::precompute_first_rank_attacks();
        let ray_attacks = Self::precompute_ray_attacks();
        let pawn = Self::precompute_pawn_attacks();
        let pawn_pushes = Self::precompute_pawn_pushes();

        Self {
            knight,
            king,
            pawn,
            pawn_pushes,
            first_rank_attacks,
            ray_attacks,
        }
    }

    fn precompute_pawn_attacks() -> [AttackTable; 2] {
        let mut attacks = [[0; 64]; 2];
        for index in 0..64 {
            let square = 1u64 << index;
            attacks[Color::White as usize][index] = ((square & NOT_A_FILE) << 7) | ((square & NOT_H_FILE) << 9);
            attacks[Color::Black as usize][index] = ((square & NOT_A_FILE) >> 9) | ((square & NOT_H_FILE) >> 7);
        }
        attacks
    }

    fn precompute_pawn_pushes() -> [AttackTable; 2] {
        let mut pushes = [[0; 64]; 2];
        for index in 0..64 {
            let square = 1u64 << index;
            pushes[Color::White as usize][index] = (square << 8) | ((square & RANK_2) << 16);
            pushes[Color::Black as usize][index] = (square >> 8) | ((square & RANK_6) >> 16);
        }
        pushes
    }

    fn precompute_knight_attacks() -> AttackTable {
        let mut attacks = [0; 64];
        for index in 0..64 {
            let square = 1u64 << index;
            attacks[index] =
                ((square & NOT_A_FILE & NOT_B_FILE) << 6) |
                ((square & NOT_G_FILE & NOT_H_FILE) << 10) |
                ((square & NOT_A_FILE) << 15) |
                ((square & NOT_H_FILE) << 17) |
                ((square & NOT_G_FILE & NOT_H_FILE) >> 6) |
                ((square & NOT_A_FILE & NOT_B_FILE) >> 10) |
                ((square & NOT_H_FILE) >> 15) |
                ((square & NOT_A_FILE) >> 17);
        }
        attacks
    }

    fn precompute_king_attacks() -> AttackTable {
        let mut attacks = [0; 64];
        for index in 0..64 {
            let square = 1u64 << index;
            attacks[index] =
                ((square & NOT_A_FILE) >> 1) |
                ((square & NOT_A_FILE) << 7) |
                ((square & NOT_A_FILE) >> 9) |
                ((square & NOT_H_FILE) << 1) |
                ((square & NOT_H_FILE) >> 7) |
                ((square & NOT_H_FILE) << 9) |
                (square << 8) |
                (square >> 8);
        }
        attacks
    }


    // Compute rook-like attacks on the first rank based on
    // occupancy and rook file.
    fn precompute_first_rank_attacks() -> FirstRankAttacks {
        let mut attacks = [[0; 8]; 256];

        /// Really slow implementation of Most Significant One Bit which is fine to use when pre-computing
        fn log_base_2(bb: Bitboard) -> u64 {
            (bb as f64).log2() as u64
        }

        for occupancy in 0..256 {
            for file in 0..8 {
                let mut left_bound = 0;
                let mut right_bound = 7;

                for bit in 0..8 {
                    if (occupancy & (1 << bit)) > 0 {
                        if (bit > left_bound) && (bit < file) {
                            left_bound = bit
                        }
                        if (bit < right_bound) && (bit > file) {
                            right_bound = bit
                        }
                    }
                }

                for index in left_bound..right_bound+1 {
                    if index != file {
                        attacks[occupancy as usize][file as usize] |= 1 << index;
                    }
                }
            }
        }
        attacks
    }

    fn precompute_ray_attacks() -> [AttackTable; 8] {
        let mut nort = [0; 64];
        let mut noea = [0; 64];
        let mut east = [0; 64];
        let mut soea = [0; 64];
        let mut sout = [0; 64];
        let mut sowe = [0; 64];
        let mut west = [0; 64];
        let mut nowe = [0; 64];

        for rank in 0..8 {
            for file in 0..8 {
                let index = rank * 8 + file;
                {
                    let mut rank = rank;
                    while rank != 7 {
                        rank += 1;
                        nort[index] |= 1 << (rank * 8 + file);
                    }
                }
                {
                    let mut rank = rank;
                    let mut file = file;
                    while (rank < 7) && (file < 7) {
                        rank += 1;
                        file += 1;
                        noea[index] |= 1 << (rank * 8 + file);
                    }
                }
                {
                    let mut file = file;
                    while file < 7 {
                        file += 1;
                        east[index] |= 1 << (rank * 8 + file);
                    }
                }
                {
                    let mut rank = rank;
                    let mut file = file;
                    while (rank > 0) && (file < 7) {
                        rank -= 1;
                        file += 1;
                        soea[index] |= 1 << (rank * 8 + file);
                    }
                }
                {
                    let mut rank = rank;
                    while rank > 0 {
                        rank -= 1;
                        sout[index] |= 1 << (rank * 8 + file);
                    }
                }
                {
                    let mut rank = rank;
                    let mut file = file;
                    while (rank > 0) && (file > 0) {
                        rank -= 1;
                        file -= 1;
                        sowe[index] |= 1 << (rank * 8 + file);
                    }
                }
                {
                    let mut file = file;
                    while file > 0 {
                        file -= 1;
                        west[index] |= 1 << (rank * 8 + file);
                    }
                }
                {
                    let mut rank = rank;
                    let mut file = file;
                    while (file > 0) && (rank < 7) {
                        file -= 1;
                        rank += 1;
                        nowe[index] |= 1 << (rank * 8 + file);
                    }
                }
            }
        }
        [nort, noea, east, soea, sout, sowe, west, nowe]
    }

    /// https://www.chessprogramming.org/Kindergarten_Bitboards
    ///
    /// Given a square and occupancy masked for rank, diagonal or anti-diagonal (note: not a file!)
    /// return an attack bitboard that considers blocking pieces
    fn kindergarten_attacks_base(&self, occupancy: Bitboard, mask: Bitboard, square: u8) -> Bitboard {
        let file = square % 8;
        let masked_occupancy = occupancy & mask;
        let occupancy_rank = ((masked_occupancy as u128 * B_FILE as u128) >> 58 & 0b111111) << 1;
        let rank_attacks = self.first_rank_attacks[occupancy_rank as usize][file as usize] as Bitboard;

        let mut filled_up_attacks = 0;
        for rank in 0..8 {
            filled_up_attacks |= rank_attacks << (rank * 8);
        }
        filled_up_attacks & mask
    }

    /// https://www.chessprogramming.org/Kindergarten_Bitboards
    fn kindergarten_attacks_file(&self, occupancy: Bitboard, mask: Bitboard, square: u8) -> Bitboard {
        let file = square % 8;
        let rank  = square / 8;

        let masked_occupancy = (occupancy & mask) >> file; // Shift occupancy to A file
        let occupancy_rank = ((masked_occupancy as u128 * DIAG_C2_H7 as u128) >> 58 & 0b111111) << 1;
        // Use reversed rank as index, since occupancy is reversed
        let rank_attacks = self.first_rank_attacks[occupancy_rank as usize][7 - rank as usize] as Bitboard;

        ((rank_attacks as u128 * DIAG_A1_H8 as u128) as Bitboard & H_FILE) >> (7 - file)
    }

    fn bishop(&self, occupancy: Bitboard, square: u8) -> Bitboard {
        let diagonal_mask =
            self.ray_attacks[Direction::NoEa as usize][square as usize] |
            self.ray_attacks[Direction::SoWe as usize][square as usize];
        let anti_diagonal_mask =
            self.ray_attacks[Direction::NoWe as usize][square as usize] |
            self.ray_attacks[Direction::SoEa as usize][square as usize];

        self.kindergarten_attacks_base(occupancy, diagonal_mask, square) | self.kindergarten_attacks_base(occupancy, anti_diagonal_mask, square)
    }

    fn rook(&self, occupancy: Bitboard, square: u8) -> Bitboard {
        let vertical =
            self.ray_attacks[Direction::Nort as usize][square as usize] |
            self.ray_attacks[Direction::Sout as usize][square as usize];
        let horizontal =
            self.ray_attacks[Direction::West as usize][square as usize] |
            self.ray_attacks[Direction::East as usize][square as usize];

        self.kindergarten_attacks_file(occupancy, vertical, square) | self.kindergarten_attacks_base(occupancy, horizontal, square)
    }

    fn queen(&self, occupancy: Bitboard, square: u8) -> Bitboard {
        self.rook(occupancy, square) | self.bishop(occupancy, square)
    }
}

#[cfg(test)]
mod tests {
    use crate::{bitboard::{pop_count, print}, board::Square};

    use super::*;

    #[test]
    fn test_pawn_attacks() {
        let attacks = Attacks::precompute_pawn_attacks();
        let square = Square::E4 as usize;

        let white_attacks = attacks[Color::White as usize][square];
        assert_eq!(white_attacks, 1 << Square::D5 as usize | 1 << Square::F5 as usize);
        print(white_attacks);

        assert_eq!(attacks[Color::White as usize][Square::H4 as usize], 1 << Square::G5 as usize);
        assert_eq!(attacks[Color::White as usize][Square::A4 as usize], 1 << Square::B5 as usize);

        assert_eq!(pop_count(attacks[Color::White as usize][Square::E8 as usize]), 0);
        assert_eq!(pop_count(attacks[Color::Black as usize][Square::E1 as usize]), 0);
    }

    #[test]
    fn test_pawn_pushes() {
        let pushes = Attacks::precompute_pawn_pushes();
        assert_eq!(pushes[Color::White as usize][Square::E4 as usize], 1 << Square::E5 as usize);
        assert_eq!(pushes[Color::White as usize][Square::A2 as usize], 1 << Square::A3 as usize | 1 << Square::A4 as usize);
        assert_eq!(pushes[Color::Black as usize][Square::E4 as usize], 1 << Square::E3 as usize);
        assert_eq!(pushes[Color::Black as usize][Square::H6 as usize], 1 << Square::H5 as usize | 1 << Square::H4 as usize);
    }

    #[test]
    fn test_knight_attacks() {
        let attacks = Attacks::precompute_knight_attacks();
        let e4_attacks = attacks[Square::E4 as usize];

        assert_ne!(e4_attacks & 1 << Square::G5 as usize, 0);
        assert_ne!(e4_attacks & 1 << Square::G3 as usize, 0);
        assert_ne!(e4_attacks & 1 << Square::C5 as usize, 0);
        assert_ne!(e4_attacks & 1 << Square::C3 as usize, 0);
        assert_ne!(e4_attacks & 1 << Square::D2 as usize, 0);
        assert_ne!(e4_attacks & 1 << Square::F2 as usize, 0);
        assert_ne!(e4_attacks & 1 << Square::D6 as usize, 0);
        assert_ne!(e4_attacks & 1 << Square::F6 as usize, 0);

        assert_eq!(e4_attacks & 1 << Square::E5 as usize, 0);
        assert_eq!(e4_attacks & 1 << Square::D4 as usize, 0);
        assert_eq!(e4_attacks & 1 << Square::A1 as usize, 0);

        assert_eq!(pop_count(attacks[Square::G1 as usize]), 3);
        assert_eq!(pop_count(attacks[Square::H8 as usize]), 2);
        print(e4_attacks);
    }

    #[test]
    fn test_king_attacks() {
        let attacks = Attacks::precompute_king_attacks();
        assert_eq!(pop_count(attacks[Square::E4 as usize]), 8);

        assert_eq!(pop_count(attacks[Square::A1 as usize]), 3);
        assert_eq!(pop_count(attacks[Square::A8 as usize]), 3);
        assert_eq!(pop_count(attacks[Square::H1 as usize]), 3);
        assert_eq!(pop_count(attacks[Square::H8 as usize]), 3);

        assert_eq!(pop_count(attacks[Square::E1 as usize]), 5);
        assert_eq!(pop_count(attacks[Square::E8 as usize]), 5);
        assert_eq!(pop_count(attacks[Square::A4 as usize]), 5);
        assert_eq!(pop_count(attacks[Square::H4 as usize]), 5);
        print(attacks[Square::E4 as usize]);
    }

    #[test]
    fn test_first_rank_attacks() {
        let attacks = Attacks::precompute_first_rank_attacks();
        //                   HGFEDCBA        HGFEDCBA
        assert_eq!(attacks[0b00010000][4], 0b11101111, "If rook is the only piece on a rank, it should be able to attack all rank");
        assert_eq!(attacks[0b00000000][4], 0b11101111, "Even with 0 occupancy rook should be able to attack all rank");
        assert_eq!(attacks[0b00000000][0], 0b11111110, "Even with 0 occupancy rook should be able to attack all rank");
        assert_eq!(attacks[0b00010001][4], 0b11101111, "If only other piece is on A rank, rook should be able to attack all rank");
        assert_eq!(attacks[0b10010000][4], 0b11101111, "If only other piece is on H rank, rook should be able to attack all rank");
        assert_eq!(attacks[0b10010001][4], 0b11101111, "If only other pieces are on A and H ranks, rook should be able to attack all rank");
        assert_eq!(attacks[0b00010100][4], 0b11101100);
        assert_eq!(attacks[0b01010100][4], 0b01101100);
        assert_eq!(attacks[0b01010010][4], 0b01101110);
        assert_eq!(attacks[0b00000010][4], 0b11101110);
        assert_eq!(attacks[0b01011000][5], 0b01010000);
    }

    #[test]
    fn test_ray_attacks() {
        let attacks = Attacks::precompute_ray_attacks();
        let square = Square::E4 as usize;
        let bitboard =
            attacks[0][square] |
            attacks[1][square] |
            attacks[2][square] |
            attacks[3][square] |
            attacks[4][square] |
            attacks[5][square] |
            attacks[6][square] |
            attacks[7][square];
        print(bitboard);
    }

    #[test]
    fn test_bishop_attacks() {
        let attacks = Attacks::new();
        let square = Square::E4 as u8;
        let occupancy =
            1 << Square::B7 as usize |
            1 << Square::B1 as usize |
            1 << Square::C2 as usize |
            1 << Square::F3 as usize;
        let bb = attacks.bishop(occupancy, square);

        assert_ne!(bb & 1 << Square::C2 as u8, 0);
        assert_eq!(bb & 1 << Square::B1 as u8, 0);
        assert_ne!(bb & 1 << Square::F3 as u8, 0);
        assert_eq!(bb & 1 << Square::G2 as u8, 0);
        assert_ne!(bb & 1 << Square::H7 as u8, 0);
        assert_ne!(bb & 1 << Square::B7 as u8, 0);
        assert_eq!(bb & 1 << Square::A8 as u8, 0);

        print(bb);
    }

    #[test]
    fn test_rook_attacks() {
        let attacks = Attacks::new();
        let square = Square::E4 as u8;
        let occupancy =
            1 << Square::B7 as usize |
            1 << Square::B1 as usize |
            1 << Square::C2 as usize |
            1 << Square::E3 as usize |
            1 << Square::F3 as usize;
        let bb = attacks.rook(occupancy, square);

        assert_ne!(bb & 1 << Square::E8 as u8, 0);
        assert_ne!(bb & 1 << Square::E7 as u8, 0);
        assert_ne!(bb & 1 << Square::E6 as u8, 0);
        assert_ne!(bb & 1 << Square::E5 as u8, 0);
        assert_ne!(bb & 1 << Square::E3 as u8, 0);
        assert_eq!(bb & 1 << Square::E2 as u8, 0);
        assert_eq!(bb & 1 << Square::E1 as u8, 0);
        assert_ne!(bb & 1 << Square::A4 as u8, 0);
        assert_ne!(bb & 1 << Square::H4 as u8, 0);
        assert_eq!(bb & 1 << Square::E4 as u8, 0);

        print(bb);
    }

    #[test]
    fn test_queen_attacks() {
        let attacks = Attacks::new();
        let square = Square::E4 as u8;
        let occupancy =
            1 << Square::B7 as usize |
            1 << Square::B1 as usize |
            1 << Square::C2 as usize |
            1 << Square::E3 as usize |
            1 << Square::F3 as usize;
        let bb = attacks.queen(occupancy, square);
        print(occupancy);
        print(bb);
    }
}