eschac/src/lookup.rs

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

macro_rules! loop_subsets {
    ($premask: ident, $subset: ident, $e: expr) => {{
        let mut $subset: u64 = 0;
        loop {
            $subset = $subset.wrapping_sub($premask) & $premask;
            $e
            if $subset == 0 {
                break;
            }
        }
    }};
}

macro_rules! by_color {
    ($c: ident, $e: expr) => {{
        ByColor([
            {
                let $c = Color::White;
                $e
            },
            {
                let $c = Color::Black;
                $e
            },
        ])
    }};
}

macro_rules! by_square {
    ($sq: ident, $init: expr, $e: expr) => {{
        let mut res = [$init; 64];
        let mut $sq: u8 = 0;
        while $sq < 64 {
            res[$sq as usize] = {
                let $sq = Square::new($sq).unwrap();
                $e
            };
            $sq += 1;
        }
        BySquare(res)
    }};
}

macro_rules! loop_bishop_directions {
    ($d: ident, $e: expr) => {{
        {
            let $d = Direction::NorthEast;
            $e
        }
        {
            let $d = Direction::NorthWest;
            $e
        }
        {
            let $d = Direction::SouthWest;
            $e
        }
        {
            let $d = Direction::SouthEast;
            $e
        }
    }};
}

macro_rules! loop_rook_directions {
    ($d: ident, $e: expr) => {{
        {
            let $d = Direction::East;
            $e
        }
        {
            let $d = Direction::North;
            $e
        }
        {
            let $d = Direction::West;
            $e
        }
        {
            let $d = Direction::South;
            $e
        }
    }};
}

macro_rules! loop_all_directions {
    ($d: ident, $e: expr) => {{
        loop_bishop_directions!($d, $e);
        loop_rook_directions!($d, $e);
    }};
}

macro_rules! by_direction {
    ($d: ident, $e: expr) => {{
        let mut res = [Bitboard(0); 8];
        let mut $d: u8 = 0;
        while $d < 8 {
            res[$d as usize] = {
                let $d = Direction::from_index($d).unwrap();
                $e
            };
            $d += 1;
        }
        ByDirection(res)
    }};
}

static RAYS: BySquare<ByDirection<Bitboard>> = by_square!(
    square,
    ByDirection([Bitboard(0); 8]),
    by_direction!(direction, {
        let mut square = square;
        let mut res = 0;
        while let Some(x) = square.trans(direction) {
            square = x;
            res |= square.bitboard().0;
        }
        Bitboard(res)
    })
);

static LINES: BySquare<BySquare<Bitboard>> = by_square!(a, BySquare([Bitboard(0); 64]), {
    by_square!(b, Bitboard(0), {
        let mut res = Bitboard(0);
        loop_all_directions!(d, {
            let r = *RAYS.get_const(a).get_const(d);
            if r.contains(b) {
                res = r;
            }
        });
        res
    })
});

static SEGMENTS: BySquare<BySquare<Bitboard>> = by_square!(a, BySquare([Bitboard(0); 64]), {
    by_square!(b, Bitboard(0), {
        let mut res = 0;
        loop_all_directions!(d, {
            let r = *RAYS.get_const(a).get_const(d);
            if r.contains(b) {
                res = r.0 & !RAYS.get_const(b).get_const(d).0;
            }
        });
        Bitboard(res | b.bitboard().0)
    })
});

static KING_MOVES: BySquare<Bitboard> = by_square!(sq, Bitboard(0), {
    let mut res = 0;
    loop_all_directions!(d, {
        if let Some(x) = sq.trans(d) {
            res |= x.bitboard().0;
        }
    });
    Bitboard(res)
});

static KNIGHT_MOVES: BySquare<Bitboard> = by_square!(s, Bitboard(0), {
    let mut res = Bitboard(0);
    if let Some(s) = s.trans(Direction::North) {
        if let Some(s) = s.trans(Direction::NorthEast) {
            res.insert(s);
        }
        if let Some(s) = s.trans(Direction::NorthWest) {
            res.insert(s)
        };
    }
    if let Some(s) = s.trans(Direction::West) {
        if let Some(s) = s.trans(Direction::NorthWest) {
            res.insert(s)
        };
        if let Some(s) = s.trans(Direction::SouthWest) {
            res.insert(s)
        };
    }
    if let Some(s) = s.trans(Direction::South) {
        if let Some(s) = s.trans(Direction::SouthWest) {
            res.insert(s)
        };
        if let Some(s) = s.trans(Direction::SouthEast) {
            res.insert(s)
        };
    }
    if let Some(s) = s.trans(Direction::East) {
        if let Some(s) = s.trans(Direction::SouthEast) {
            res.insert(s)
        };
        if let Some(s) = s.trans(Direction::NorthEast) {
            res.insert(s)
        };
    }
    res
});

static PAWN_ATTACKS: ByColor<BySquare<Bitboard>> = {
    by_color!(color, {
        let direction = match color {
            Color::White => Direction::North,
            Color::Black => Direction::South,
        };
        by_square!(square, Bitboard(0), {
            let mut res = Bitboard(0);
            if let Some(square) = square.trans(direction) {
                if let Some(s) = square.trans(Direction::East) {
                    res.insert(s)
                };
                if let Some(s) = square.trans(Direction::West) {
                    res.insert(s)
                };
            }
            res
        })
    })
};

const fn blocked_ray(square: Square, direction: Direction, blockers: Bitboard) -> Bitboard {
    let ray = *RAYS.get_const(square).get_const(direction);
    let blockers = Bitboard(blockers.0 & ray.0);
    let square = if (direction as u8) < 4 {
        Bitboard::first(&blockers)
    } else {
        Bitboard::last(&blockers)
    };
    match square {
        None => ray,
        Some(square) => Bitboard(ray.0 & !RAYS.get_const(square).get_const(direction).0),
    }
}

#[inline]
pub(crate) fn ray(square: Square, direction: Direction) -> Bitboard {
    *RAYS.get(square).get(direction)
}
#[inline]
pub(crate) fn line(a: Square, b: Square) -> Bitboard {
    *LINES.get(a).get(b)
}
#[inline]
pub(crate) fn segment(a: Square, b: Square) -> Bitboard {
    *SEGMENTS.get(a).get(b)
}
#[inline]
pub(crate) fn king(square: Square) -> Bitboard {
    *KING_MOVES.get(square)
}
#[inline]
pub(crate) fn knight(square: Square) -> Bitboard {
    *KNIGHT_MOVES.get(square)
}
#[inline]
pub(crate) fn pawn_attack(color: Color, square: Square) -> Bitboard {
    *PAWN_ATTACKS.get(color).get(square)
}

const fn bishop_premask(square: Square) -> Bitboard {
    let mut premask = 0;
    loop_bishop_directions!(direction, {
        premask |= RAYS.get_const(square).get_const(direction).0;
    });
    premask &= !Rank::First.bitboard().0;
    premask &= !Rank::Eighth.bitboard().0;
    premask &= !File::A.bitboard().0;
    premask &= !File::H.bitboard().0;
    Bitboard(premask)
}

const fn rook_premask(square: Square) -> Bitboard {
    let rays = RAYS.get_const(square);
    let mut premask = 0;
    premask |= rays.get_const(Direction::East).0 & !File::H.bitboard().0;
    premask |= rays.get_const(Direction::North).0 & !Rank::Eighth.bitboard().0;
    premask |= rays.get_const(Direction::West).0 & !File::A.bitboard().0;
    premask |= rays.get_const(Direction::South).0 & !Rank::First.bitboard().0;
    Bitboard(premask)
}

static MAGICS: (BySquare<Magic>, BySquare<Magic>, usize) = {
    let mut len: usize = 0;
    (
        by_square!(
            square,
            Magic {
                premask: Bitboard(0),
                factor: 0,
                offset: 0
            },
            {
                let premask = bishop_premask(square).0;
                let factor = bishop_factor(square);
                let mut i = usize::MAX;
                let mut j = 0;
                loop_subsets!(premask, blockers, {
                    let cur = hash(BISHOP_SHR, factor, Bitboard(blockers | !premask));
                    if cur < i {
                        i = cur;
                    }
                    if cur > j {
                        j = cur;
                    }
                });
                let offset = len as isize - i as isize;
                len += j - i + 1;
                Magic {
                    premask: Bitboard(!premask),
                    factor,
                    offset,
                }
            }
        ),
        by_square!(
            square,
            Magic {
                premask: Bitboard(0),
                factor: 0,
                offset: 0
            },
            {
                let premask = rook_premask(square).0;
                let factor = rook_factor(square);
                let mut i = usize::MAX;
                let mut j = 0;
                loop_subsets!(premask, blockers, {
                    let cur = hash(ROOK_SHR, factor, Bitboard(blockers | !premask));
                    if cur < i {
                        i = cur;
                    }
                    if cur > j {
                        j = cur;
                    }
                });
                let offset = len as isize - i as isize;
                len += j - i + 1;
                Magic {
                    premask: Bitboard(!premask),
                    factor,
                    offset,
                }
            }
        ),
        len,
    )
};

#[derive(Clone, Copy)]
struct Magic {
    premask: Bitboard,
    factor: u64,
    offset: isize,
}

#[allow(long_running_const_eval)]
static MAGIC_TABLE: [Bitboard; MAGICS.2] = {
    let mut table = [Bitboard(0); MAGICS.2];
    by_square!(square, (), {
        let Magic {
            premask,
            factor,
            offset,
        } = *MAGICS.0.get_const(square);
        let premask = !premask.0;
        loop_subsets!(premask, blockers, {
            let index = (hash(BISHOP_SHR, factor, Bitboard(blockers | !premask)) as isize + offset)
                as usize;
            let mut res = 0;
            loop_bishop_directions!(direction, {
                res |= blocked_ray(square, direction, Bitboard(blockers)).0;
            });
            assert!(table[index].0 == 0 || table[index].0 == res);
            table[index] = Bitboard(res);
        });
    });
    by_square!(square, (), {
        let Magic {
            premask,
            factor,
            offset,
        } = *MAGICS.1.get_const(square);
        let premask = !premask.0;
        loop_subsets!(premask, blockers, {
            let index =
                (offset + hash(ROOK_SHR, factor, Bitboard(blockers | !premask)) as isize) as usize;
            let mut res = 0;
            loop_rook_directions!(direction, {
                res |= blocked_ray(square, direction, Bitboard(blockers)).0;
            });
            assert!(table[index].0 == 0 || table[index].0 == res);
            table[index] = Bitboard(res);
        });
    });
    table
};

#[inline]
pub(crate) fn bishop(square: Square, blockers: Bitboard) -> Bitboard {
    unsafe { magic_aux(BISHOP_SHR, *(MAGICS.0).get(square), blockers) }
}

#[inline]
pub(crate) fn rook(square: Square, blockers: Bitboard) -> Bitboard {
    unsafe { magic_aux(ROOK_SHR, *(MAGICS.1).get(square), blockers) }
}

#[inline]
unsafe fn magic_aux(shr: u8, magic: Magic, blockers: Bitboard) -> Bitboard {
    let Magic {
        premask,
        factor,
        offset,
    } = magic;
    debug_assert!(MAGIC_TABLE
        .get(
            (hash(shr, factor, blockers | premask) as isize)
                .checked_add(offset)
                .unwrap() as usize
        )
        .is_some());
    *MAGIC_TABLE.get_unchecked(
        ((hash(shr, factor, blockers | premask) as isize).unchecked_add(offset)) as usize,
    )
}

#[inline]
const fn hash(shr: u8, factor: u64, x: Bitboard) -> usize {
    (x.0.wrapping_mul(factor) >> shr) as usize
}

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