{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell       #-}

-- | Utilities for reading from / writing to the filesystem
module Test.Goblin.Persist where

import           Control.Arrow (first)
import qualified Data.Binary as Binary
import           Data.Bits (Bits(..), FiniteBits(..))
import qualified Data.ByteString.Lazy as BL
import           Data.List.Extra (chunksOf)
import           Data.Word (Word64)
import           Language.Haskell.TH (Q, Exp, runIO, stringE)
import           Moo.GeneticAlgorithm.Types (Population)


-- | Decode a `Population Bool` from a lazy ByteString.
decodePopulation :: BL.ByteString -> Population Bool
decodePopulation :: ByteString -> Population Bool
decodePopulation ByteString
bs =
  -- The added Int tells us how much padding we must remove.
  let intermediate :: [(([Word64],Int),Double)]
      intermediate :: [(([Word64], Int), Double)]
intermediate = ByteString -> [(([Word64], Int), Double)]
forall a. Binary a => ByteString -> a
Binary.decode ByteString
bs
   in ((([Word64], Int), Double) -> Phenotype Bool)
-> [(([Word64], Int), Double)] -> Population Bool
forall a b. (a -> b) -> [a] -> [b]
map ((([Word64], Int) -> [Bool])
-> (([Word64], Int), Double) -> Phenotype Bool
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first ([Word64], Int) -> [Bool]
splitter) [(([Word64], Int), Double)]
intermediate

-- | Encode a `Population Bool` to a lazy ByteString.
encodePopulation :: Population Bool -> BL.ByteString
encodePopulation :: Population Bool -> ByteString
encodePopulation Population Bool
pop =
  -- The added Int tells us how much padding we must remove.
  let intermediate :: [(([Word64],Int),Double)]
      intermediate :: [(([Word64], Int), Double)]
intermediate = (Phenotype Bool -> (([Word64], Int), Double))
-> Population Bool -> [(([Word64], Int), Double)]
forall a b. (a -> b) -> [a] -> [b]
map (([Bool] -> ([Word64], Int))
-> Phenotype Bool -> (([Word64], Int), Double)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first [Bool] -> ([Word64], Int)
grouper) Population Bool
pop
   in [(([Word64], Int), Double)] -> ByteString
forall a. Binary a => a -> ByteString
Binary.encode [(([Word64], Int), Double)]
intermediate

-- | Load a Population from a file and return the first (highest
-- scoring) genome.
readFirstGenomeFromFile :: FilePath -> IO [Bool]
readFirstGenomeFromFile :: FilePath -> IO [Bool]
readFirstGenomeFromFile FilePath
filePath =
  (Phenotype Bool -> [Bool]
forall a b. (a, b) -> a
fst (Phenotype Bool -> [Bool])
-> (Population Bool -> Phenotype Bool) -> Population Bool -> [Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Population Bool -> Phenotype Bool
forall a. [a] -> a
head) (Population Bool -> [Bool]) -> IO (Population Bool) -> IO [Bool]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> IO (Population Bool)
readPopulationFromFile FilePath
filePath

-- | Read a Population from a file.
readPopulationFromFile :: FilePath -> IO (Population Bool)
readPopulationFromFile :: FilePath -> IO (Population Bool)
readPopulationFromFile FilePath
filePath =
  ByteString -> Population Bool
decodePopulation (ByteString -> Population Bool)
-> IO ByteString -> IO (Population Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> IO ByteString
BL.readFile FilePath
filePath

-- | Write a Population to a file.
writePopulationToFile :: FilePath -> Population Bool -> IO ()
writePopulationToFile :: FilePath -> Population Bool -> IO ()
writePopulationToFile FilePath
filePath Population Bool
pop =
  FilePath -> ByteString -> IO ()
BL.writeFile FilePath
filePath (Population Bool -> ByteString
encodePopulation Population Bool
pop)

-- | Read a file at compile-time, and splice in the `show` of its ByteString
-- as a String in the source file.
loadBestPopToShownByteString :: FilePath -> Q Exp
loadBestPopToShownByteString :: FilePath -> Q Exp
loadBestPopToShownByteString FilePath
fp = do
  FilePath -> Q Exp
stringE (FilePath -> Q Exp)
-> (ByteString -> FilePath) -> ByteString -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> FilePath
forall a. Show a => a -> FilePath
show (ByteString -> Q Exp) -> Q ByteString -> Q Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (IO ByteString -> Q ByteString
forall a. IO a -> Q a
runIO (IO ByteString -> Q ByteString) -> IO ByteString -> Q ByteString
forall a b. (a -> b) -> a -> b
$ do
    ByteString
bs <- FilePath -> IO ByteString
BL.readFile FilePath
fp
    let best :: Phenotype Bool
best = Population Bool -> Phenotype Bool
forall a. [a] -> a
head (ByteString -> Population Bool
decodePopulation ByteString
bs)
    ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Population Bool -> ByteString
encodePopulation [Phenotype Bool
best]))

-- | Splice in a genome as a `show`n ByteString, and decode it at runtime.
-- This is less safe than inlining the whole list of `Bool`s into source code,
-- but results in less source bloat.
loadGoblinDataFromFilePath :: FilePath
                           -> Q Exp
loadGoblinDataFromFilePath :: FilePath -> Q Exp
loadGoblinDataFromFilePath FilePath
fp = [|
  let popStr = $(loadBestPopToShownByteString fp)
      genome = case decodePopulation (read popStr) of
                 [] -> error "sigGenChain: impossible"
                 (x,_):_ -> x
   in mkEmptyGoblin genome
   |]

-- | Convert an Integral into a little-endian binary representation.
integralToBits :: (FiniteBits a) => a -> [Bool]
integralToBits :: a -> [Bool]
integralToBits a
x = (Int -> Bool) -> [Int] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (a -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit a
x) [Int
0 .. a -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize a
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

-- | Convert from a little-endian binary representation to an Integral.
integralFromBits :: forall a. (Integral a, FiniteBits a) => [Bool] -> a
integralFromBits :: [Bool] -> a
integralFromBits [Bool]
bits =
  if [Bool] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Bool]
bits Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
numBits Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
     then FilePath -> a
forall a. HasCallStack => FilePath -> a
error (FilePath
"length of bits is not a multiple of " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> Int -> FilePath
forall a. Show a => a -> FilePath
show Int
numBits)
     else [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Bool -> Int -> a) -> [Bool] -> [Int] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Bool
b Int
i -> if Bool
b then a
2a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^Int
i else a
0)
              [Bool]
bits
              [Int
0 .. Int
numBits])
 where
  numBits :: Int
numBits = a -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (a
forall a. HasCallStack => a
undefined :: a)

-- | Returns the padded list, plus the number of padding bits added.
padBits :: Int -> [Bool] -> ([Bool], Int)
padBits :: Int -> [Bool] -> ([Bool], Int)
padBits Int
multiple [Bool]
bs = ( [Bool]
bs [Bool] -> [Bool] -> [Bool]
forall a. [a] -> [a] -> [a]
++ Int -> Bool -> [Bool]
forall a. Int -> a -> [a]
replicate Int
delta Bool
False
                      , Int
delta )
 where
  delta :: Int
delta = Int
multiple Int -> Int -> Int
forall a. Num a => a -> a -> a
- ([Bool] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Bool]
bs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
multiple)

-- | Split a list of `Word64`s, little-endian style, into their requisite
-- bits. The `Int` describes the amount of padding to drop, since we must
-- necessarily pad up to a 64-bit multiple.
splitter :: ([Word64],Int) -> [Bool]
splitter :: ([Word64], Int) -> [Bool]
splitter ([Word64]
xs, Int
delta) =
  let bits :: [Bool]
bits = [[Bool]] -> [Bool]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Word64 -> [Bool]
forall a. FiniteBits a => a -> [Bool]
integralToBits (Word64 -> [Bool]) -> [Word64] -> [[Bool]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Word64]
xs)
   in Int -> [Bool] -> [Bool]
forall a. Int -> [a] -> [a]
take ([Bool] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Bool]
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
delta) [Bool]
bits

-- | Group a list of `Bool`s into a list of `Word64`s, little-endian style.
-- The `Int` describes the amount of padding added, since we must necessarily
-- pad up to a 64-bit multiple.
grouper :: [Bool] -> ([Word64], Int)
grouper :: [Bool] -> ([Word64], Int)
grouper [Bool]
bs =
  let ([Bool]
bits, Int
delta) = Int -> [Bool] -> ([Bool], Int)
padBits Int
64 [Bool]
bs
   in ( ([Bool] -> Word64) -> [[Bool]] -> [Word64]
forall a b. (a -> b) -> [a] -> [b]
map [Bool] -> Word64
forall a. (Integral a, FiniteBits a) => [Bool] -> a
integralFromBits (Int -> [Bool] -> [[Bool]]
forall a. HasCallStack => Int -> [a] -> [[a]]
chunksOf Int
64 [Bool]
bits)
      , Int
delta )