module Utils.TranspositionTable where

import           AppPrelude

import           Models.Command
import           Models.Move
import           Models.Position              (ZKey (..))
import           Models.Score
import           Utils.Board

import           Data.Vector.Storable.Mutable (IOVector)
import qualified Data.Vector.Storable.Mutable as Vector
import           Data.Word
import           Foreign.Storable.Generic
import           Test.QuickCheck


type TTable = IOVector StorableTEntry


data TEntry = TEntry {
  TEntry -> ZKey
zobristKey :: ZKey,
  TEntry -> Maybe Move
bestMove   :: Maybe Move,
  TEntry -> Score
score      :: Score,
  TEntry -> Age
depth      :: Depth,
  TEntry -> NodeType
nodeType   :: NodeType,
  TEntry -> Age
age        :: Age
} deriving (TEntry -> TEntry -> Bool
(TEntry -> TEntry -> Bool)
-> (TEntry -> TEntry -> Bool) -> Eq TEntry
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TEntry -> TEntry -> Bool
== :: TEntry -> TEntry -> Bool
$c/= :: TEntry -> TEntry -> Bool
/= :: TEntry -> TEntry -> Bool
Eq, Int -> TEntry -> ShowS
[TEntry] -> ShowS
TEntry -> String
(Int -> TEntry -> ShowS)
-> (TEntry -> String) -> ([TEntry] -> ShowS) -> Show TEntry
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TEntry -> ShowS
showsPrec :: Int -> TEntry -> ShowS
$cshow :: TEntry -> String
show :: TEntry -> String
$cshowList :: [TEntry] -> ShowS
showList :: [TEntry] -> ShowS
Show, (forall x. TEntry -> Rep TEntry x)
-> (forall x. Rep TEntry x -> TEntry) -> Generic TEntry
forall x. Rep TEntry x -> TEntry
forall x. TEntry -> Rep TEntry x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. TEntry -> Rep TEntry x
from :: forall x. TEntry -> Rep TEntry x
$cto :: forall x. Rep TEntry x -> TEntry
to :: forall x. Rep TEntry x -> TEntry
Generic)

instance Arbitrary TEntry where
  arbitrary :: Gen TEntry
arbitrary = ZKey -> Maybe Move -> Score -> Age -> NodeType -> Age -> TEntry
TEntry (ZKey -> Maybe Move -> Score -> Age -> NodeType -> Age -> TEntry)
-> Gen ZKey
-> Gen (Maybe Move -> Score -> Age -> NodeType -> Age -> TEntry)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen ZKey
forall a. Arbitrary a => Gen a
arbitrary Gen (Maybe Move -> Score -> Age -> NodeType -> Age -> TEntry)
-> Gen (Maybe Move)
-> Gen (Score -> Age -> NodeType -> Age -> TEntry)
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen (Maybe Move)
forall a. Arbitrary a => Gen a
arbitrary
                     Gen (Score -> Age -> NodeType -> Age -> TEntry)
-> Gen Score -> Gen (Age -> NodeType -> Age -> TEntry)
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen Score
forall a. Arbitrary a => Gen a
arbitrary Gen (Age -> NodeType -> Age -> TEntry)
-> Gen Age -> Gen (NodeType -> Age -> TEntry)
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen Age
forall a. Arbitrary a => Gen a
arbitrary
                     Gen (NodeType -> Age -> TEntry)
-> Gen NodeType -> Gen (Age -> TEntry)
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen NodeType
forall a. Arbitrary a => Gen a
arbitrary Gen (Age -> TEntry) -> Gen Age -> Gen TEntry
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Age, Age) -> Gen Age
forall a. Random a => (a, a) -> Gen a
choose (Age
0, Age
31)

data StorableTEntry = StorableTEntry {
  StorableTEntry -> ZKey
zobristKey :: ZKey,
  StorableTEntry -> Word64
info       :: Word64
} deriving (forall x. StorableTEntry -> Rep StorableTEntry x)
-> (forall x. Rep StorableTEntry x -> StorableTEntry)
-> Generic StorableTEntry
forall x. Rep StorableTEntry x -> StorableTEntry
forall x. StorableTEntry -> Rep StorableTEntry x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. StorableTEntry -> Rep StorableTEntry x
from :: forall x. StorableTEntry -> Rep StorableTEntry x
$cto :: forall x. Rep StorableTEntry x -> StorableTEntry
to :: forall x. Rep StorableTEntry x -> StorableTEntry
Generic

instance GStorable StorableTEntry


tTableSize :: (?opts :: EngineOptions) => Word64
tTableSize :: (?opts::EngineOptions) => Word64
tTableSize | ?opts::EngineOptions
EngineOptions
?opts.hashSize Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
0 = Word64
0
           | Bool
otherwise          = Int -> Word64
toBoard Int
bits
  where
    bits :: Int
bits = Word64 -> Int
msb (Word16 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ?opts::EngineOptions
EngineOptions
?opts.hashSize) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
16


hashZKey :: (?opts :: EngineOptions) => ZKey -> Int
hashZKey :: (?opts::EngineOptions) => ZKey -> Int
hashZKey (ZKey Word64
zKey) =
  Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
zKey Word64 -> Word64 -> Word64
forall a. Integral a => a -> a -> a
% Word64
(?opts::EngineOptions) => Word64
tTableSize)


encodeTEntry :: TEntry -> StorableTEntry
encodeTEntry :: TEntry -> StorableTEntry
encodeTEntry TEntry {Score
Maybe Move
Age
NodeType
ZKey
$sel:zobristKey:TEntry :: TEntry -> ZKey
$sel:bestMove:TEntry :: TEntry -> Maybe Move
$sel:score:TEntry :: TEntry -> Score
$sel:depth:TEntry :: TEntry -> Age
$sel:nodeType:TEntry :: TEntry -> NodeType
$sel:age:TEntry :: TEntry -> Age
zobristKey :: ZKey
bestMove :: Maybe Move
score :: Score
depth :: Age
nodeType :: NodeType
age :: Age
..} = StorableTEntry {
  $sel:zobristKey:StorableTEntry :: ZKey
zobristKey = ZKey
zobristKey,
  $sel:info:StorableTEntry :: Word64
info = Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
bestMoveN
    Word64 -> Word64 -> Word64
.| Word16 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Score -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Score
score :: Word16) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
<< Int
32
    Word64 -> Word64 -> Word64
.| Age -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Age
depth Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
<< Int
48
    Word64 -> Word64 -> Word64
.| Age -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Age
nodeTypeN Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
<< Int
56
    Word64 -> Word64 -> Word64
.| Age -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Age
age Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
<< Int
58
  }
  where
    StorableMove Word32
bestMoveN = Maybe Move -> StorableMove
encodeMove Maybe Move
bestMove
    NodeType Age
nodeTypeN     = NodeType
nodeType


decodeTEntry :: StorableTEntry -> Maybe TEntry
decodeTEntry :: StorableTEntry -> Maybe TEntry
decodeTEntry StorableTEntry {Word64
ZKey
$sel:zobristKey:StorableTEntry :: StorableTEntry -> ZKey
$sel:info:StorableTEntry :: StorableTEntry -> Word64
zobristKey :: ZKey
info :: Word64
..}
  | Word64 -> Int -> Bool
testSquare Word64
info Int
63 = Maybe TEntry
forall a. Maybe a
Nothing
  | Bool
otherwise          = TEntry -> Maybe TEntry
forall a. a -> Maybe a
Just TEntry {
      $sel:zobristKey:TEntry :: ZKey
zobristKey = ZKey
zobristKey
      , $sel:bestMove:TEntry :: Maybe Move
bestMove   = StorableMove -> Maybe Move
decodeMove (StorableMove -> Maybe Move) -> StorableMove -> Maybe Move
forall a b. (a -> b) -> a -> b
$ Word64 -> StorableMove
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
info
      , $sel:score:TEntry :: Score
score      = Word64 -> Score
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
info Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
>> Int
32)
      , $sel:depth:TEntry :: Age
depth      = Word64 -> Age
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
info Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
>> Int
48)
      , $sel:nodeType:TEntry :: NodeType
nodeType   = Word64 -> NodeType
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word64
info Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
>> Int
56) Word64 -> Word64 -> Word64
& Word64
3)
      , $sel:age:TEntry :: Age
age        = Word64 -> Age
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word64
info Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
>> Int
58) Word64 -> Word64 -> Word64
& Word64
31)
  }


lookupEntry
  :: (?tTable :: TTable, ?opts :: EngineOptions) => ZKey -> IO (Maybe TEntry)
lookupEntry :: (?tTable::TTable, ?opts::EngineOptions) =>
ZKey -> IO (Maybe TEntry)
lookupEntry !ZKey
zKey
  | Word64
(?opts::EngineOptions) => Word64
tTableSize Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0 = Maybe TEntry -> IO (Maybe TEntry)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe TEntry
forall a. Maybe a
Nothing
  | Bool
otherwise      = do
  Maybe TEntry
entry <- StorableTEntry -> Maybe TEntry
decodeTEntry (StorableTEntry -> Maybe TEntry)
-> IO StorableTEntry -> IO (Maybe TEntry)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) StorableTEntry -> Int -> IO StorableTEntry
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
Vector.unsafeRead ?tTable::TTable
TTable
MVector (PrimState IO) StorableTEntry
?tTable ((?opts::EngineOptions) => ZKey -> Int
ZKey -> Int
hashZKey ZKey
zKey)
  Maybe TEntry -> IO (Maybe TEntry)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((TEntry -> Bool) -> Maybe TEntry -> Maybe TEntry
forall a. (a -> Bool) -> Maybe a -> Maybe a
maybeFilter ((ZKey -> ZKey -> Bool
forall a. Eq a => a -> a -> Bool
== ZKey
zKey) (ZKey -> Bool) -> (TEntry -> ZKey) -> TEntry -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (.zobristKey)) Maybe TEntry
entry)


lookupScore
  :: (?tTable :: TTable, ?opts :: EngineOptions)
  => Score -> Score -> Depth -> ZKey -> IO (Maybe (Score, Maybe Move))
lookupScore :: (?tTable::TTable, ?opts::EngineOptions) =>
Score -> Score -> Age -> ZKey -> IO (Maybe (Score, Maybe Move))
lookupScore !Score
alpha !Score
beta !Age
depth !ZKey
zKey = do
  Maybe TEntry
entry <- (TEntry -> Bool) -> Maybe TEntry -> Maybe TEntry
forall a. (a -> Bool) -> Maybe a -> Maybe a
maybeFilter ((Age -> Age -> Bool
forall a. Ord a => a -> a -> Bool
>= Age
depth) (Age -> Bool) -> (TEntry -> Age) -> TEntry -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (.depth)) (Maybe TEntry -> Maybe TEntry)
-> IO (Maybe TEntry) -> IO (Maybe TEntry)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (?tTable::TTable, ?opts::EngineOptions) =>
ZKey -> IO (Maybe TEntry)
ZKey -> IO (Maybe TEntry)
lookupEntry ZKey
zKey
  let !score :: Maybe Score
score = TEntry -> Maybe Score
getScore (TEntry -> Maybe Score) -> Maybe TEntry -> Maybe Score
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Maybe TEntry
entry
      !bestMove :: Maybe Move
bestMove = (.bestMove) (TEntry -> Maybe Move) -> Maybe TEntry -> Maybe Move
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Maybe TEntry
entry
  Maybe (Score, Maybe Move) -> IO (Maybe (Score, Maybe Move))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((, Maybe Move
bestMove) (Score -> (Score, Maybe Move))
-> Maybe Score -> Maybe (Score, Maybe Move)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Score
score)
  where
    getScore :: TEntry -> Maybe Score
getScore TEntry {Score
$sel:score:TEntry :: TEntry -> Score
score :: Score
score, NodeType
$sel:nodeType:TEntry :: TEntry -> NodeType
nodeType :: NodeType
nodeType} =
      case NodeType
nodeType of
        NodeType
PV                   -> Score -> Maybe Score
forall a. a -> Maybe a
Just Score
score
        NodeType
Cut | Score
score Score -> Score -> Bool
forall a. Ord a => a -> a -> Bool
>= Score
beta  -> Score -> Maybe Score
forall a. a -> Maybe a
Just Score
beta
        NodeType
All | Score
score Score -> Score -> Bool
forall a. Ord a => a -> a -> Bool
<= Score
alpha -> Score -> Maybe Score
forall a. a -> Maybe a
Just Score
alpha
        NodeType
_                    -> Maybe Score
forall a. Maybe a
Nothing


lookupBestMove
  :: (?tTable :: TTable, ?opts :: EngineOptions) => ZKey -> IO (Maybe Move)
lookupBestMove :: (?tTable::TTable, ?opts::EngineOptions) => ZKey -> IO (Maybe Move)
lookupBestMove !ZKey
zKey = do
  Maybe TEntry
entry <- (?tTable::TTable, ?opts::EngineOptions) =>
ZKey -> IO (Maybe TEntry)
ZKey -> IO (Maybe TEntry)
lookupEntry ZKey
zKey
  Maybe Move -> IO (Maybe Move)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((.bestMove) (TEntry -> Maybe Move) -> Maybe TEntry -> Maybe Move
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Maybe TEntry
entry)


insert
  :: (?tTable :: TTable, ?opts :: EngineOptions) => ZKey -> TEntry -> IO ()
insert :: (?tTable::TTable, ?opts::EngineOptions) => ZKey -> TEntry -> IO ()
insert !ZKey
zKey !TEntry
newEntry
 | Word64
(?opts::EngineOptions) => Word64
tTableSize Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0 = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
 | Bool
otherwise       = do
  Maybe TEntry
entry <- StorableTEntry -> Maybe TEntry
decodeTEntry (StorableTEntry -> Maybe TEntry)
-> IO StorableTEntry -> IO (Maybe TEntry)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) StorableTEntry -> Int -> IO StorableTEntry
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
Vector.unsafeRead ?tTable::TTable
TTable
MVector (PrimState IO) StorableTEntry
?tTable Int
hashKey
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe TEntry -> TEntry -> Bool
isStaleEntry Maybe TEntry
entry TEntry
newEntry)
    (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) StorableTEntry
-> Int -> StorableTEntry -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
Vector.unsafeWrite ?tTable::TTable
TTable
MVector (PrimState IO) StorableTEntry
?tTable Int
hashKey (TEntry -> StorableTEntry
encodeTEntry TEntry
newEntry)
  where
    hashKey :: Int
hashKey = (?opts::EngineOptions) => ZKey -> Int
ZKey -> Int
hashZKey ZKey
zKey


isStaleEntry :: Maybe TEntry -> TEntry -> Bool
isStaleEntry :: Maybe TEntry -> TEntry -> Bool
isStaleEntry (Just TEntry
entry) TEntry
newEntry =
  TEntry
newEntry.age Age -> Age -> Age
forall a. Num a => a -> a -> a
- TEntry
entry.age Age -> Age -> Bool
forall a. Ord a => a -> a -> Bool
> Age
1
  Bool -> Bool -> Bool
|| TEntry
newEntry.depth Age -> Age -> Bool
forall a. Ord a => a -> a -> Bool
>= TEntry
entry.depth
  Bool -> Bool -> Bool
&& (TEntry
newEntry.nodeType NodeType -> NodeType -> Bool
forall a. Eq a => a -> a -> Bool
== NodeType
PV Bool -> Bool -> Bool
|| TEntry
entry.nodeType NodeType -> NodeType -> Bool
forall a. Eq a => a -> a -> Bool
/= NodeType
PV)
isStaleEntry Maybe TEntry
Nothing TEntry
_ =
  Bool
True


emptyTEntry :: StorableTEntry
emptyTEntry :: StorableTEntry
emptyTEntry = StorableTEntry {
  $sel:zobristKey:StorableTEntry :: ZKey
zobristKey = ZKey
0,
  $sel:info:StorableTEntry :: Word64
info       = Int -> Word64
toBoard Int
63
}


create ::  EngineOptions -> IO TTable
create :: EngineOptions -> IO TTable
create EngineOptions
opts = Int -> StorableTEntry -> IO (MVector (PrimState IO) StorableTEntry)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> a -> m (MVector (PrimState m) a)
Vector.replicate (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
(?opts::EngineOptions) => Word64
tTableSize) StorableTEntry
emptyTEntry
  where
    ?opts = ?opts::EngineOptions
EngineOptions
opts


reset :: TTable -> IO ()
reset :: TTable -> IO ()
reset = (TTable -> StorableTEntry -> IO ())
-> StorableTEntry -> TTable -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip TTable -> StorableTEntry -> IO ()
MVector (PrimState IO) StorableTEntry -> StorableTEntry -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> a -> m ()
Vector.set StorableTEntry
emptyTEntry


clear :: TTable -> IO ()
clear :: TTable -> IO ()
clear = TTable -> IO ()
MVector (PrimState IO) StorableTEntry -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> m ()
Vector.clear