QuadBitboard.hs

module Game.Chess.Internal.QuadBitboard (
  -- * The QuadBitboard data type
  QuadBitboard
, occupied, black, white
, pawns, knights, bishops, rooks, queens, kings
, wPawns, wKnights, wBishops, wRooks, wQueens, wKings
, bPawns, bKnights, bBishops, bRooks, bQueens, bKings
, insufficientMaterial
, toString
  -- * Square codes
, Word4(..)
, pattern NoPiece
, pattern WhitePawn, pattern WhiteKnight, pattern WhiteBishop
, pattern WhiteRook, pattern WhiteQueen, pattern WhiteKing
, pattern BlackPawn, pattern BlackKnight, pattern BlackBishop
, pattern BlackRook, pattern BlackQueen, pattern BlackKing
  -- * Construction
, empty, standard
  -- * Access
, (!), setNibble
  -- * Transformations
  -- ** Normal moves
, move, move'
  -- ** Castling
, whiteKingsideCastle, whiteQueensideCastle
, blackKingsideCastle, blackQueensideCastle
  -- ** En passant
, enPassant
  -- ** Promotion
, whitePromotion, blackPromotion, whitePromotion', blackPromotion'
) where

import Control.Applicative (liftA2)
import Data.Bifunctor (first)
import Data.Binary ( Binary(put, get) )
import Data.Bits
    ( Bits((.|.), (.&.), xor, complement
          , unsafeShiftR, unsafeShiftL, testBit, setBit, clearBit, popCount)
    , FiniteBits(..) )
import Data.Char (ord, toLower)
import Data.Ix ( Ix(inRange) )
import Data.List ( groupBy, intercalate )
import Data.String ( IsString(..) )
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as M
import Data.Vector.Unboxed (Vector, MVector, Unbox)
import Data.Word (Word8, Word64)
import Foreign.Storable
import Game.Chess.Internal.Square (toIndex, toRF, Sq(A8))
import GHC.Enum
    ( boundedEnumFrom,
      boundedEnumFromThen,
      predError,
      succError,
      toEnumError )
import GHC.Exts (IsList(Item, fromList, toList))
import GHC.Ptr ( castPtr, plusPtr )
import Numeric (showHex)

data QuadBitboard = QBB { black, pbq, nbk, rqk :: {-# UNPACK #-} !Word64 }
                    deriving Eq

occupied, pnr, white, pawns, knights, bishops, rooks, queens, kings :: QuadBitboard -> Word64
occupied = liftA2 (.|.) pbq $ liftA2 (.|.) nbk rqk
pnr      = liftA2  xor  pbq $ liftA2  xor  nbk rqk
white    = liftA2  xor  occupied black
pawns    = liftA2 (.&.) pnr pbq
knights  = liftA2 (.&.) pnr nbk
bishops  = liftA2 (.&.) pbq nbk
rooks    = liftA2 (.&.) pnr rqk
queens   = liftA2 (.&.) pbq rqk
kings    = liftA2 (.&.) nbk rqk

wPawns, wKnights, wBishops, wRooks, wQueens, wKings :: QuadBitboard -> Word64
wPawns   = liftA2 (.&.) pawns (complement . black)
wKnights = liftA2 (.&.) knights (complement . black)
wBishops = liftA2 (.&.) bishops (complement . black)
wRooks   = liftA2 (.&.) rooks (complement . black)
wQueens  = liftA2 (.&.) queens (complement . black)
wKings   = liftA2 (.&.) kings (complement . black)

bPawns, bKnights, bBishops, bRooks, bQueens, bKings :: QuadBitboard -> Word64
bPawns   = liftA2 (.&.) pawns black
bKnights = liftA2 (.&.) knights black
bBishops = liftA2 (.&.) bishops black
bRooks   = liftA2 (.&.) rooks black
bQueens  = liftA2 (.&.) queens black
bKings   = liftA2 (.&.) kings black

{-# INLINE pnr #-}
{-# INLINE occupied #-}
{-# INLINE white #-}
{-# INLINE pawns #-}
{-# INLINE knights #-}
{-# INLINE bishops #-}
{-# INLINE rooks #-}
{-# INLINE queens #-}
{-# INLINE kings #-}
{-# INLINE wPawns #-}
{-# INLINE wKnights #-}
{-# INLINE wBishops #-}
{-# INLINE wRooks #-}
{-# INLINE wQueens #-}
{-# INLINE wKings #-}
{-# INLINE bPawns #-}
{-# INLINE bKnights #-}
{-# INLINE bBishops #-}
{-# INLINE bRooks #-}
{-# INLINE bQueens #-}
{-# INLINE bKings #-}

pattern NoPiece :: Word4
pattern NoPiece     = 0

pattern WhitePawn, WhiteKnight, WhiteBishop, WhiteRook, WhiteQueen, WhiteKing
  :: Word4
pattern WhitePawn   = 2
pattern WhiteKnight = 4
pattern WhiteBishop = 6
pattern WhiteRook   = 8
pattern WhiteQueen  = 10
pattern WhiteKing   = 12

pattern BlackPawn, BlackKnight, BlackBishop, BlackRook, BlackQueen, BlackKing
  :: Word4
pattern BlackPawn   = 3
pattern BlackKnight = 5
pattern BlackBishop = 7
pattern BlackRook   = 9
pattern BlackQueen  = 11
pattern BlackKing   = 13

instance IsList QuadBitboard where
  type Item QuadBitboard = (Sq, Word4)
  fromList = mconcat . map (uncurry singleton . first toIndex)
  toList qbb = go maxBound [] where
    go sq xs
      | sq /= minBound = go (pred sq) xs'
      | otherwise      = xs'
     where nb = qbb ! toIndex sq
           xs' | nb /= NoPiece = (sq, nb) : xs
               | otherwise     = xs

empty, standard :: QuadBitboard
empty = QBB 0 0 0 0
standard = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR"

-- | bitwise XOR
instance Semigroup QuadBitboard where
  {-# INLINE (<>) #-}
  QBB b0 b1 b2 b3 <> QBB b0' b1' b2' b3' =
    QBB (b0 `xor` b0') (b1 `xor` b1') (b2 `xor` b2') (b3 `xor` b3')

instance Monoid QuadBitboard where
  mempty = empty

insufficientMaterial :: QuadBitboard -> Bool
insufficientMaterial qbb@QBB{black, pbq, nbk, rqk} =
  noPawnsNorQueens && eachSideHasOneKing && noRooks &&
  (oneSideHasAtMostOneMinorPiece || opposingBishopsOnEquallyColoredSquares)
 where
  eachSideHasOneKing = popCount (wKings qbb) == 1 && popCount (bKings qbb) == 1
  noPawnsNorQueens = pbq `xor` bishops qbb == 0
  noRooks = popCount rqk == 2
  oneSideHasAtMostOneMinorPiece =
    (popCount (nbk .&. complement black) == 1 && atMostOneMinorPiece black) ||
    (popCount (nbk .&. black) == 1 && atMostOneMinorPiece (complement black))
  opposingBishopsOnEquallyColoredSquares =
    popCount (knights qbb) == 0 &&
    popCount (nbk .&. black) == 2 && popCount (nbk .&. complement black) == 2 &&
    even (countTrailingZeros (wBishops qbb)) ==
    even (countTrailingZeros (bBishops qbb))
  atMostOneMinorPiece mask = popCount (nbk .&. mask) `elem` [1,2]

instance Storable QuadBitboard where
  sizeOf _ = 32
  alignment _ = 8
  peek p = QBB <$> peek (castPtr p) <*> peek (castPtr $ p `plusPtr` 8) <*> peek (castPtr $ p `plusPtr` 16) <*> peek (castPtr $ p `plusPtr` 24)
  poke p QBB{black, pbq, nbk, rqk} = do
    poke (castPtr $ p) black
    poke (castPtr $ p `plusPtr` 8) pbq
    poke (castPtr $ p `plusPtr` 16) nbk
    poke (castPtr $ p `plusPtr` 24) rqk

newtype instance MVector s QuadBitboard = MV_QuadBitboard (MVector s (Word64, Word64, Word64, Word64))
newtype instance Vector    QuadBitboard = V_QuadBitboard (Vector (Word64, Word64, Word64, Word64))
instance Unbox QuadBitboard

instance M.MVector MVector QuadBitboard where
  basicLength (MV_QuadBitboard v) = M.basicLength v
  basicUnsafeSlice i n (MV_QuadBitboard v) = MV_QuadBitboard $ M.basicUnsafeSlice i n v
  basicOverlaps (MV_QuadBitboard v1) (MV_QuadBitboard v2) = M.basicOverlaps v1 v2
  basicUnsafeNew n = MV_QuadBitboard <$> M.basicUnsafeNew n
  basicInitialize (MV_QuadBitboard v) = M.basicInitialize v
  basicUnsafeReplicate n (QBB b0 b1 b2 b3) = MV_QuadBitboard <$> M.basicUnsafeReplicate n (b0, b1, b2, b3)
  basicUnsafeRead (MV_QuadBitboard v) i = f <$> M.basicUnsafeRead v i where
    f (b0, b1, b2, b3) = QBB b0 b1 b2 b3
  basicUnsafeWrite (MV_QuadBitboard v) i (QBB b0 b1 b2 b3) = M.basicUnsafeWrite v i (b0, b1, b2, b3)
  basicClear (MV_QuadBitboard v) = M.basicClear v
  basicSet (MV_QuadBitboard v) (QBB b0 b1 b2 b3) = M.basicSet v (b0, b1, b2, b3)
  basicUnsafeCopy (MV_QuadBitboard v1) (MV_QuadBitboard v2) = M.basicUnsafeCopy v1 v2
  basicUnsafeMove (MV_QuadBitboard v1) (MV_QuadBitboard v2) = M.basicUnsafeMove v1 v2
  basicUnsafeGrow (MV_QuadBitboard v) n = MV_QuadBitboard <$> M.basicUnsafeGrow v n

instance G.Vector Vector QuadBitboard where
  {-# INLINE basicUnsafeIndexM #-}
  basicUnsafeFreeze (MV_QuadBitboard v) = V_QuadBitboard <$> G.basicUnsafeFreeze v
  basicUnsafeThaw (V_QuadBitboard v) = MV_QuadBitboard <$> G.basicUnsafeThaw v
  basicLength (V_QuadBitboard v) = G.basicLength v
  basicUnsafeSlice i n (V_QuadBitboard v) = V_QuadBitboard $ G.basicUnsafeSlice  i n v
  basicUnsafeIndexM (V_QuadBitboard v) i
    = f <$> G.basicUnsafeIndexM v i where
    f (b0, b1, b2, b3) = QBB b0 b1 b2 b3
  basicUnsafeCopy (MV_QuadBitboard mv) (V_QuadBitboard v) = G.basicUnsafeCopy mv v
  elemseq _ (QBB b0 b1 b2 b3) z
    = G.elemseq (undefined :: Vector a) b0
    $ G.elemseq (undefined :: Vector a) b1
    $ G.elemseq (undefined :: Vector a) b2
    $ G.elemseq (undefined :: Vector a) b3
    z

newtype Word4 = W4 Word8
              deriving (Bits, Eq, Integral, Ix, Num, Ord, Read, Real)

instance Show Word4 where
  show NoPiece     = "NoPiece"
  show WhiteKing   = "WhiteKing"
  show WhitePawn   = "WhitePawn"
  show WhiteKnight = "WhiteKnight"
  show WhiteBishop = "WhiteBishop"
  show WhiteRook   = "WhiteRook"
  show WhiteQueen  = "WhiteQueen"
  show BlackKing   = "BlackKing"
  show BlackPawn   = "BlackPawn"
  show BlackKnight = "BlackKnight"
  show BlackBishop = "BlackBishop"
  show BlackRook   = "BlackRook"
  show BlackQueen  = "BlackQueen"
  show (W4 n)      = "W4 " <> show n

instance Bounded Word4 where
  minBound = 0
  maxBound = 0xF

instance Enum Word4 where
  succ x | x /= maxBound = x + 1
         | otherwise     = succError "Word4"
  pred x | x /= minBound = x - 1
         | otherwise     = predError "Word4"
  toEnum i | i >= 0 && i <= fromIntegral (maxBound::Word4) = W4 $ fromIntegral i
           | otherwise     = toEnumError "Word4" i (minBound::Word4, maxBound::Word4)
  fromEnum (W4 x) = fromIntegral x
  enumFrom        = boundedEnumFrom
  enumFromThen    = boundedEnumFromThen

instance FiniteBits Word4 where
  finiteBitSize _ = 4
  countLeadingZeros (W4 x) = countLeadingZeros x - 4
  countTrailingZeros (W4 x) = countTrailingZeros x

-- | law: singleton i x ! i = x where inRange (0,63) i && inRange (0,15) x
{-# INLINE singleton #-}
singleton :: Int -> Word4 -> QuadBitboard
singleton !sq !nb = QBB (f 0) (f 1) (f 2) (f 3) where
  !b = 1 `unsafeShiftL` sq
  f !n = fromIntegral ((nb `unsafeShiftR` n) .&. 1) * b

(!) :: QuadBitboard -> Int -> Word4
(!) QBB{..} sq = fromIntegral $ f black 0 .|. f pbq 1 .|. f nbk 2 .|. f rqk 3 where
  f !bb !n = ((bb `unsafeShiftR` sq) .&. 1) `unsafeShiftL` n

setNibble :: Bits nibble => QuadBitboard -> Int -> nibble -> QuadBitboard
setNibble QBB{..} sq nb = QBB (f 0 black) (f 1 pbq) (f 2 nbk) (f 3 rqk) where
  f n | nb `testBit` n = (`setBit` sq)
      | otherwise      = (`clearBit` sq)

instance Binary QuadBitboard where
  get = QBB <$> get <*> get <*> get <*> get
  put QBB{..} = put black *> put pbq *> put nbk *> put rqk

instance IsString QuadBitboard where
  fromString = go (toRF A8) mempty where
    go _ !qbb "" = qbb
    go (!r,_) qbb ('/':xs) = go (r - 1, 0) qbb xs
    go rf@(!r,!f) !qbb (x:xs)
      | inRange ('1','8') x = go (r, f + ord x - ord '0') qbb xs
      | otherwise = go (r, f + 1) (qbb <> singleton (toIndex rf) nb) xs where
        nb = case x of
          'P' -> WhitePawn
          'N' -> WhiteKnight
          'B' -> WhiteBishop
          'R' -> WhiteRook
          'Q' -> WhiteQueen
          'K' -> WhiteKing
          'p' -> BlackPawn
          'n' -> BlackKnight
          'b' -> BlackBishop
          'r' -> BlackRook
          'q' -> BlackQueen
          'k' -> BlackKing
          _ -> error $ "QuadBitBoard.fromString: Illegal FEN character " <> show x

instance Show QuadBitboard where
  show QBB{..} =
     "QBB {black = 0x" <> showHex black
    (", pbq = 0x" <> showHex pbq
    (", nbk = 0x" <> showHex nbk
    (", rqk = 0x" <> showHex rqk "}")))

toString :: QuadBitboard -> String
toString qbb = intercalate "/" $ rank <$> [7, 6..0] where
  rank r = concatMap countEmpty . groupBy spaces $ charAt r <$> [0..7]
  countEmpty xs | head xs == spc = show $ length xs
                | otherwise      = xs
  spaces x y = x == spc && x == y
  charAt r f = maybe spc (if odd nb then toLower else id) $
    lookup (nb `div` 2) $ zip [1..] "PNBRQK"
   where nb = qbb ! (r*8+f)
  spc = ' '

-- | Move a nibble.  Note that this function, while convenient, isn't very
-- fast as it needs to lookup the source nibble value.
move :: QuadBitboard -> Int -> Int -> QuadBitboard
move qbb fromSq toSq = qbb <> move' fromSq (qbb ! fromSq) toSq (qbb ! toSq)

move' :: Int -> Word4 -> Int -> Word4 -> QuadBitboard
move' fromSq fromCode toSq toCode =
  singleton fromSq fromCode <> singleton toSq (fromCode `xor` toCode)

whiteKingsideCastle, whiteQueensideCastle, blackKingsideCastle, blackQueensideCastle
  :: QuadBitboard
whiteKingsideCastle  = move' 4 WhiteKing 6 NoPiece <> move' 7 WhiteRook 5 NoPiece
whiteQueensideCastle = move' 4 WhiteKing 2 NoPiece <> move' 0 WhiteRook 3 NoPiece
blackKingsideCastle  = move' 60 BlackKing 62 NoPiece <> move' 63 BlackRook 61 NoPiece
blackQueensideCastle = move' 60 BlackKing 58 NoPiece <> move' 56 BlackRook 59 NoPiece

enPassant :: Int -> Int -> QuadBitboard
enPassant fromSq toSq
  | fromSq < toSq
  = move' fromSq WhitePawn toSq NoPiece <> singleton (toSq-8) BlackPawn
  | otherwise
  = move' fromSq BlackPawn toSq NoPiece <> singleton (toSq+8) WhitePawn

whitePromotion, blackPromotion :: QuadBitboard -> Int -> Int -> Word4 -> QuadBitboard
whitePromotion qbb fromSq toSq promoCode =
  qbb <> whitePromotion' fromSq toSq (qbb ! toSq) promoCode
blackPromotion qbb fromSq toSq promoCode =
  qbb <> blackPromotion' fromSq toSq (qbb ! toSq) promoCode

whitePromotion', blackPromotion' :: Int -> Int -> Word4 -> Word4 -> QuadBitboard
whitePromotion' fromSq toSq toCode promoCode =
  singleton fromSq WhitePawn <> singleton toSq (toCode `xor` promoCode)
blackPromotion' fromSq toSq toCode promoCode =
  singleton fromSq BlackPawn <> singleton toSq (toCode `xor` promoCode)