eschac/src/lookup.rs

//! Lookup tables initialisation.
//!
//! Move generation in eschac requires about 1MB of precomputed lookup tables.

use crate::bitboard::*;
use crate::board::*;
use crate::magics::*;
use crate::rays::*;

pub(crate) use init::InitialisedLookup;

/// Forces the initialisation of the lookup tables.
///
/// It is not necessary to call this function, as lookup tables are initialised lazily, but it can
/// be used to ensure that they are initialised before a given time.
pub fn init() {
    InitialisedLookup::init();
}

pub(crate) struct Lookup {
    rays: Rays,
    lines: BySquare<BySquare<Bitboard>>,
    segments: BySquare<BySquare<Bitboard>>,
    pawn_attacks: ByColor<BySquare<Bitboard>>,
    king_moves: BySquare<Bitboard>,
    knight_moves: BySquare<Bitboard>,
    pub(crate) magics: Magics,
}

impl Lookup {
    #[inline]
    pub(crate) fn line(&self, a: Square, b: Square) -> Bitboard {
        *self.lines.get(a).get(b)
    }

    #[inline]
    pub(crate) fn segment(&self, a: Square, b: Square) -> Bitboard {
        *self.segments.get(a).get(b)
    }

    #[inline]
    pub(crate) fn ray(&self, square: Square, direction: Direction) -> Bitboard {
        self.rays.ray(square, direction)
    }

    #[inline]
    pub(crate) fn king(&self, square: Square) -> Bitboard {
        *self.king_moves.get(square)
    }

    #[inline]
    pub(crate) fn knight(&self, square: Square) -> Bitboard {
        *self.knight_moves.get(square)
    }

    #[inline]
    pub(crate) fn pawn_attack(&self, color: Color, square: Square) -> Bitboard {
        *self.pawn_attacks.get(color).get(square)
    }

    #[inline]
    pub(crate) fn bishop(&self, square: Square, blockers: Bitboard) -> Bitboard {
        self.magics.bishop(square, blockers)
    }

    #[inline]
    pub(crate) fn rook(&self, square: Square, blockers: Bitboard) -> Bitboard {
        self.magics.rook(square, blockers)
    }

    /// `role != Pawn`
    #[inline]
    pub(crate) fn targets(&self, role: Role, from: Square, blockers: Bitboard) -> Bitboard {
        match role {
            Role::Pawn => unreachable!(),
            Role::Knight => self.knight(from),
            Role::Bishop => self.bishop(from, blockers),
            Role::Rook => self.rook(from, blockers),
            Role::Queen => self.bishop(from, blockers) | self.rook(from, blockers),
            Role::King => self.king(from),
        }
    }

    pub(crate) fn compute() -> Self {
        let rays = Rays::new();

        let lines = BySquare::new(|a| {
            BySquare::new(|b| {
                for d in Direction::all() {
                    let r = rays.ray(a, d);
                    if r.contains(b) {
                        return r;
                    }
                }
                Bitboard::new()
            })
        });

        let segments = BySquare::new(|a| {
            BySquare::new(|b| {
                for d in Direction::all() {
                    let r = rays.ray(a, d);
                    if r.contains(b) {
                        return r & !rays.ray(b, d);
                    }
                }
                b.bitboard()
            })
        });

        let pawn_attacks = ByColor::new(|color| {
            let direction = match color {
                Color::White => Direction::North,
                Color::Black => Direction::South,
            };
            BySquare::new(|square| {
                let mut res = Bitboard::new();
                if let Some(square) = square.trans(direction) {
                    square.trans(Direction::East).map(|s| res.insert(s));
                    square.trans(Direction::West).map(|s| res.insert(s));
                }
                res
            })
        });

        let king_moves = BySquare::new(|square| {
            let mut res = Bitboard::new();
            for direction in Direction::all() {
                if let Some(x) = square.trans(direction) {
                    res |= x.bitboard();
                }
            }
            res
        });

        let knight_moves = BySquare::new(|s| {
            let mut res = Bitboard::new();
            if let Some(s) = s.trans(Direction::North) {
                s.trans(Direction::NorthEast).map(|s| res.insert(s));
                s.trans(Direction::NorthWest).map(|s| res.insert(s));
            }
            if let Some(s) = s.trans(Direction::West) {
                s.trans(Direction::NorthWest).map(|s| res.insert(s));
                s.trans(Direction::SouthWest).map(|s| res.insert(s));
            }
            if let Some(s) = s.trans(Direction::South) {
                s.trans(Direction::SouthWest).map(|s| res.insert(s));
                s.trans(Direction::SouthEast).map(|s| res.insert(s));
            }
            if let Some(s) = s.trans(Direction::East) {
                s.trans(Direction::SouthEast).map(|s| res.insert(s));
                s.trans(Direction::NorthEast).map(|s| res.insert(s));
            }
            res
        });

        let magics = Magics::compute(&rays);

        Self {
            rays,
            lines,
            segments,
            pawn_attacks,
            king_moves,
            knight_moves,
            magics,
        }
    }
}

mod init {
    use std::{mem::MaybeUninit, sync::LazyLock};

    use super::Lookup;

    static mut LOOKUP: MaybeUninit<Lookup> = MaybeUninit::uninit();

    #[allow(static_mut_refs)]
    static INIT: LazyLock<()> = LazyLock::new(|| unsafe {
        LOOKUP.write(Lookup::compute());
    });

    #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub(crate) struct InitialisedLookup(());

    impl InitialisedLookup {
        #[inline]
        pub(crate) fn init() -> Self {
            LazyLock::force(&INIT);
            Self(())
        }
    }

    impl std::ops::Deref for InitialisedLookup {
        type Target = Lookup;
        #[allow(static_mut_refs)]
        #[inline]
        fn deref(&self) -> &Self::Target {
            unsafe { LOOKUP.assume_init_ref() }
        }
    }
}