makebook

src/Game/Chess.hs

@@ -27,7 +27,8 @@ module Game.Chess (
   -- * Chess moves
 , Ply(..)
   -- ** Converting from/to algebraic notation
-, strictSAN, relaxedSAN, fromSAN, toSAN, unsafeToSAN, fromUCI, toUCI, fromPolyglot
+, strictSAN, relaxedSAN, fromSAN, toSAN, unsafeToSAN, fromUCI, toUCI
+, fromPolyglot, toPolyglot
   -- ** Move generation
 , legalPlies
   -- ** Executing moves
@@ -508,6 +509,18 @@ fromPolyglot pos pl@(unpack -> (from, to, _)) = case color pos of
           -> from `move` C8
   _ -> pl
 
+toPolyglot :: Position -> Ply -> Ply
+toPolyglot pos pl@(unpack -> (from, to, _)) = case color pos of
+  White | from == toIndex E1 && canCastleKingside pos && to == toIndex G1
+          -> from `move` H1
+        | from == toIndex E1 && canCastleQueenside pos && to == toIndex C1
+          -> from `move` A1
+  Black | from == toIndex E8 && canCastleKingside pos && to == toIndex G8
+          -> from `move` H8
+        | from == toIndex E8 && canCastleQueenside pos && to == toIndex C8
+          -> from `move` A8
+  _ -> pl
+
 -- | Parse a move in the format used by the Universal Chess Interface protocol.
 fromUCI :: Position -> String -> Maybe Ply
 fromUCI pos (fmap (splitAt 2) . splitAt 2 -> (from, (to, promo)))

src/Game/Chess/PGN.hs

@@ -208,8 +208,9 @@ moveDoc ro pos (o,ts) = (fillSep $ go pos True ts <> [pretty o]) <> line where
   nag n = "$" <> pretty n
 
 weightedForest :: PGN -> Forest (Rational, Ply)
-weightedForest (PGN games) = merge . concatMap rate . map snd $ filter ok games where
-  ok (ts, (o, _)) = Nothing == lookup "FEN" ts && o /= Undecided
+weightedForest (PGN games) = merge . concatMap rate $ snd <$> filter ok games
+ where
+  ok (tags, (o, _)) = isNothing (lookup "FEN" tags) && o /= Undecided
   rate (o, ts) = f startpos <$> trunk ts where
     w c | o == Win c = 1
         | o == Win (opponent c) = -1
@@ -219,10 +220,9 @@ weightedForest (PGN games) = merge . concatMap rate . map snd $ filter ok games
   trunk [] = []
   trunk (x:_) = [x { subForest = trunk (subForest x)}]
   merge [] = []
-  merge ((Node a ts):xs) =
+  merge ((Node a ts) : xs) =
       sortOn (Down . fst . rootLabel)
     $ Node (w, snd a) (merge $ ts ++ concatMap subForest good) : merge bad
    where
-    (good, bad) = partition (eq a . rootLabel) xs where eq a b = snd a == snd b
-    w = fst a + sum (map (fst . rootLabel) good)
-
+    (good, bad) = partition (eq a . rootLabel) xs where eq x y = snd x == snd y
+    w = fst a + sum (fst . rootLabel <$> good)

src/Game/Chess/Polyglot/Book.hs

@@ -2,12 +2,13 @@
 {-# LANGUAGE TemplateHaskell #-}
 module Game.Chess.Polyglot.Book (
   PolyglotBook
-, fromByteString
 , defaultBook, twic
-, readPolyglotFile
+, fromByteString, toByteString
+, readPolyglotFile, writePolyglotFile
 , bookPly
 , bookPlies
 , bookForest
+, findPosition
 ) where
 
 import Control.Arrow
@@ -18,13 +19,16 @@ import Data.ByteString (ByteString)
 import qualified Data.ByteString.Internal as BS
 import qualified Data.ByteString as BS
 import Data.FileEmbed
+import Data.List
+import Data.Ord
 import qualified Data.Vector.Storable as VS
 import Data.Tree
 import Data.Word
-import Foreign.ForeignPtr (plusForeignPtr)
+import Foreign.ForeignPtr (castForeignPtr, plusForeignPtr)
 import Foreign.Ptr (castPtr)
 import Foreign.Storable
 import Game.Chess
+import Game.Chess.PGN
 import Game.Chess.Polyglot.Hash
 import GHC.Ptr
 import System.Random (RandomGen)
@@ -36,6 +40,10 @@ data BookEntry = BookEntry {
 , learn :: {-# UNPACK #-} !Word32
 } deriving (Eq, Show)
 
+instance Ord BookEntry where
+  compare (BookEntry k1 _ w1 _) (BookEntry k2 _ w2 _) =
+    compare k1 k2 <> compare (Down w1) (Down w2)
+
 instance Storable BookEntry where
   sizeOf _ = 16
   alignment _ = alignment (undefined :: Word64)
@@ -58,9 +66,7 @@ peekBE ptr = go ptr 0 (sizeOf (undefined :: a)) where
 pokeBE :: forall a. (Bits a, Integral a, Num a, Storable a) => Ptr Word8 -> a -> IO ()
 pokeBE p x = go x (sizeOf x) where
   go _ 0 = pure ()
-  go !x !n = do
-    pokeElemOff p (n-1) (fromIntegral x)
-    go (x `shiftR` 8) (n-1)
+  go !x !n = pokeElemOff p (n-1) (fromIntegral x) *> go (x `shiftR` 8) (n-1)
 
 defaultBook, twic :: PolyglotBook
 defaultBook = twic
@@ -79,9 +85,35 @@ fromByteString bs = Book v where
   (fptr, off, len) = BS.toForeignPtr bs
   elemSize = sizeOf (undefined `asTypeOf` VS.head v)
 
+toByteString :: PolyglotBook -> ByteString
+toByteString (Book v) = BS.fromForeignPtr (castForeignPtr fptr) off (len * elemSize)
+ where
+  (fptr, off, len) = VS.unsafeToForeignPtr v
+  elemSize = sizeOf (undefined `asTypeOf` VS.head v)
+
 readPolyglotFile :: FilePath -> IO PolyglotBook
 readPolyglotFile = fmap fromByteString . BS.readFile
 
+writePolyglotFile :: FilePath -> PolyglotBook -> IO ()
+writePolyglotFile fp = BS.writeFile fp . toByteString
+
+fromList :: [BookEntry] -> PolyglotBook
+fromList = Book . VS.fromList . sort
+
+toList :: PolyglotBook -> [BookEntry]
+toList (Book v) = VS.toList v
+
+makeBook :: PGN -> PolyglotBook
+makeBook = fromList . concatMap (foldTree f . annot startpos) . weightedForest
+ where
+  annot pos (Node a ts) =
+    Node (pos, a) $ annot (unsafeDoPly pos (snd a)) <$> ts
+  f (pos, (w, pl)) xs
+    | w > 0
+    = BookEntry (hashPosition pos) (toPolyglot pos pl) (floor w) 0 : concat xs
+    | otherwise
+    = concat xs
+
 bookForest :: PolyglotBook -> Position -> Forest Ply
 bookForest b p = tree <$> bookPlies b p where
   tree pl = Node pl . bookForest b $ unsafeDoPly p pl

src/Game/Chess/QuadBitboard.hs

@@ -45,12 +45,12 @@ data QuadBitboard = QBB { black :: {-# UNPACK #-} !Word64
                         , rqk :: {-# UNPACK #-} !Word64
                         } deriving (Eq)
 
-occupied, pnr, white, pawns, knights, bishops, rooks, queens, kings
-  :: QuadBitboard -> Word64
+occupied, pnr, white :: QuadBitboard -> Word64
 occupied QBB{pbq, nbk, rqk} = pbq  .|.  nbk  .|.  rqk
 pnr      QBB{pbq, nbk, rqk} = pbq `xor` nbk `xor` rqk
-white = liftA2 xor occupied black
+white                       = liftA2 xor occupied black
 
+pawns, knights, bishops, rooks, queens, kings :: QuadBitboard -> Word64
 pawns   = liftA2 (.&.) pnr pbq
 knights = liftA2 (.&.) pnr nbk
 bishops = liftA2 (.&.) pbq nbk
@@ -81,6 +81,7 @@ bKings   = liftA2 (.&.) kings black
 {-# INLINE knights #-}
 {-# INLINE bishops #-}
 {-# INLINE rooks #-}
+{-# INLINE queens #-}
 {-# INLINE kings #-}
 {-# INLINE wPawns #-}
 {-# INLINE wKnights #-}
@@ -123,7 +124,7 @@ instance FiniteBits Word4 where
   countTrailingZeros (W4 x) = countTrailingZeros x
 
 pattern NoPiece :: Word4
-pattern NoPiece = 0
+pattern NoPiece     = 0
 
 pattern WhitePawn, WhiteKnight, WhiteBishop, WhiteRook, WhiteQueen, WhiteKing
   :: Word4
@@ -157,7 +158,7 @@ square (bit -> b) nb = QBB (f 0) (f 1) (f 2) (f 3) where
 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)
+      | otherwise      = (`clearBit` sq)
 
 instance Binary QuadBitboard where
   get = QBB <$> get <*> get <*> get <*> get