eschac/src/bitboard.rs

//! Sets of squares.

use crate::board::*;

use std::iter::ExactSizeIterator;
use std::iter::FusedIterator;

/// A set of squares.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Bitboard(pub(crate) u64);

impl Bitboard {
    /// Returns an empty bitboard.
    #[inline]
    pub fn new() -> Self {
        Self(0)
    }

    /// Returns `true` if the bitboard is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.0 == 0
    }

    /// Returns the square in the bitboard with the smallest index.
    #[inline]
    pub fn first(&self) -> Option<Square> {
        let mask = self.0;
        match mask {
            0 => None,
            _ => Some(unsafe { Square::new_unchecked(mask.trailing_zeros() as u8) }),
        }
    }

    /// Removes the square in the bitboard with the smallest index and returns it, or `None` if the
    /// bitboard is empty.
    #[inline]
    pub fn pop(&mut self) -> Option<Square> {
        let Self(ref mut mask) = self;
        let square = match mask {
            0 => None,
            _ => Some(unsafe { Square::new_unchecked(mask.trailing_zeros() as u8) }),
        };
        *mask &= mask.wrapping_sub(1);
        square
    }

    /// Inserts a square in the bitboard.
    #[inline]
    pub fn insert(&mut self, square: Square) {
        self.0 |= 1 << square as u8;
    }

    /// Returns `true` if the bitboard contains the given square.
    #[inline]
    pub fn contains(&self, square: Square) -> bool {
        self.0 & (1 << square as u8) != 0
    }

    /// Returns the mirror of the bitboard (see [`Setup::mirror`](crate::setup::Setup::mirror)).
    #[inline]
    pub fn mirror(self) -> Bitboard {
        let [a, b, c, d, e, f, g, h] = self.0.to_le_bytes();
        Self(u64::from_le_bytes([h, g, f, e, d, c, b, a]))
    }

    #[inline]
    pub(crate) fn trans(&self, direction: Direction) -> Self {
        match direction {
            Direction::North => Self(self.0 << 8),
            Direction::NorthEast => Self(self.0 << 9) & !File::A.bitboard(),
            Direction::East => Self(self.0 << 1) & !File::A.bitboard(),
            Direction::SouthEast => Self(self.0 >> 7) & !File::A.bitboard(),
            Direction::South => Self(self.0 >> 8),
            Direction::SouthWest => Self(self.0 >> 9) & !File::H.bitboard(),
            Direction::West => Self(self.0 >> 1) & !File::H.bitboard(),
            Direction::NorthWest => Self(self.0 << 7) & !File::H.bitboard(),
        }
    }
}

impl std::ops::BitOr for Bitboard {
    type Output = Self;
    #[inline]
    fn bitor(self, rhs: Self) -> Self::Output {
        Self(self.0 | rhs.0)
    }
}
impl std::ops::BitAnd for Bitboard {
    type Output = Self;
    #[inline]
    fn bitand(self, rhs: Self) -> Self::Output {
        Self(self.0 & rhs.0)
    }
}
impl std::ops::BitXor for Bitboard {
    type Output = Self;
    #[inline]
    fn bitxor(self, rhs: Self) -> Self::Output {
        Self(self.0 ^ rhs.0)
    }
}
impl std::ops::BitOrAssign for Bitboard {
    #[inline]
    fn bitor_assign(&mut self, rhs: Self) {
        self.0 |= rhs.0;
    }
}
impl std::ops::BitAndAssign for Bitboard {
    #[inline]
    fn bitand_assign(&mut self, rhs: Self) {
        self.0 &= rhs.0;
    }
}
impl std::ops::BitXorAssign for Bitboard {
    #[inline]
    fn bitxor_assign(&mut self, rhs: Self) {
        self.0 ^= rhs.0;
    }
}
impl std::ops::Not for Bitboard {
    type Output = Self;
    #[inline]
    fn not(self) -> Self::Output {
        Self(!self.0)
    }
}

impl Iterator for Bitboard {
    type Item = Square;
    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.pop()
    }
    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.len();
        (len, Some(len))
    }
    #[inline]
    fn for_each<F>(self, mut f: F)
    where
        Self: Sized,
        F: FnMut(Self::Item),
    {
        let mut mask = self.0;
        while mask != 0 {
            f(unsafe { Square::new_unchecked(mask.trailing_zeros() as u8) });
            mask &= mask.wrapping_sub(1);
        }
    }
    #[inline]
    fn fold<B, F>(self, init: B, mut f: F) -> B
    where
        Self: Sized,
        F: FnMut(B, Self::Item) -> B,
    {
        let mut mask = self.0;
        let mut acc = init;
        while mask != 0 {
            acc = f(acc, unsafe {
                Square::new_unchecked(mask.trailing_zeros() as u8)
            });
            mask &= mask.wrapping_sub(1);
        }
        acc
    }
}
impl FusedIterator for Bitboard {}
impl ExactSizeIterator for Bitboard {
    #[inline]
    fn len(&self) -> usize {
        self.0.count_ones() as usize
    }
}

pub(crate) trait BitboardIterExt {
    fn reduce_or(self) -> Bitboard;
}
impl<T> BitboardIterExt for T
where
    T: Iterator<Item = Bitboard>,
{
    #[inline]
    fn reduce_or(self) -> Bitboard {
        self.fold(Bitboard(0), |a, b| a | b)
    }
}