{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances #-}
module Data.BinaryState where

import Control.Monad
import qualified Control.Monad.State as State
import qualified Data.Binary as Binary
import qualified Data.Binary.Put as Put
import qualified Data.Binary.Get as Get
import qualified Data.ByteString.Lazy as B
import Data.Word
import Data.Int

type PutState s a = State.StateT s Put.PutM a
type GetState s a = State.StateT s Binary.Get a

class BinaryState s a where
  put :: a -> PutState s ()
  get :: GetState s a

instance (Binary.Binary a) => BinaryState () a where
  put :: a -> PutState () ()
put a
x = a -> PutState () ()
forall a s. Binary a => a -> PutState s ()
putZ a
x
  get :: GetState () a
get = GetState () a
forall a s. Binary a => GetState s a
getZ

putZ :: (Binary.Binary a) => a -> PutState s ()
putZ :: a -> PutState s ()
putZ a
x = PutM () -> PutState s ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
State.lift (a -> PutM ()
forall t. Binary t => t -> PutM ()
Binary.put a
x)

getZ :: (Binary.Binary a) => GetState s a
getZ :: GetState s a
getZ = Get a -> GetState s a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
State.lift Get a
forall t. Binary t => Get t
Binary.get

------------------------------------------------

encodeS :: (BinaryState s a) => s -> a -> B.ByteString
encodeS :: s -> a -> ByteString
encodeS s
s a
a = PutM () -> ByteString
Put.runPut (PutM () -> ByteString) -> PutM () -> ByteString
forall a b. (a -> b) -> a -> b
$ StateT s PutM () -> s -> PutM ()
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
State.evalStateT (a -> StateT s PutM ()
forall s a. BinaryState s a => a -> PutState s ()
put a
a) s
s

decodeS :: (BinaryState s a) => s -> B.ByteString -> a
decodeS :: s -> ByteString -> a
decodeS s
s ByteString
str = Get a -> ByteString -> a
forall a. Get a -> ByteString -> a
Get.runGet (StateT s Get a -> s -> Get a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
State.evalStateT StateT s Get a
forall s a. BinaryState s a => GetState s a
get s
s) ByteString
str

decodeWith :: GetState s a -> s -> B.ByteString -> a
decodeWith :: GetState s a -> s -> ByteString -> a
decodeWith GetState s a
getter s
s ByteString
str =
  let (a
x,ByteString
_,ByteOffset
_) = Get a -> ByteString -> ByteOffset -> (a, ByteString, ByteOffset)
forall a.
Get a -> ByteString -> ByteOffset -> (a, ByteString, ByteOffset)
Get.runGetState (GetState s a -> s -> Get a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
State.evalStateT GetState s a
getter s
s) ByteString
str ByteOffset
0
  in  a
x

encodeWith :: (a -> PutState s ()) -> s -> a -> B.ByteString
encodeWith :: (a -> PutState s ()) -> s -> a -> ByteString
encodeWith a -> PutState s ()
putter s
s a
a = PutM () -> ByteString
Put.runPut (PutM () -> ByteString) -> PutM () -> ByteString
forall a b. (a -> b) -> a -> b
$ PutState s () -> s -> PutM ()
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
State.evalStateT (a -> PutState s ()
putter a
a) s
s

encodeFile :: BinaryState s a => FilePath -> s -> a -> IO ()
encodeFile :: FilePath -> s -> a -> IO ()
encodeFile FilePath
f s
s a
v = FilePath -> ByteString -> IO ()
B.writeFile FilePath
f (s -> a -> ByteString
forall s a. BinaryState s a => s -> a -> ByteString
encodeS s
s a
v)

decodeFile :: BinaryState s a => FilePath -> s -> IO a
decodeFile :: FilePath -> s -> IO a
decodeFile FilePath
f s
s = (ByteString -> a) -> IO ByteString -> IO a
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (s -> ByteString -> a
forall s a. BinaryState s a => s -> ByteString -> a
decodeS s
s) (FilePath -> IO ByteString
B.readFile FilePath
f)

decodeFile' :: BinaryState s a => FilePath -> s -> IO (a, s)
decodeFile' :: FilePath -> s -> IO (a, s)
decodeFile' FilePath
path s
s = do
  ByteString
str <- FilePath -> IO ByteString
B.readFile FilePath
path
  let getter :: Get (a, s)
getter = StateT s Get a -> s -> Get (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
State.runStateT StateT s Get a
forall s a. BinaryState s a => GetState s a
get s
s
  (a, s) -> IO (a, s)
forall (m :: * -> *) a. Monad m => a -> m a
return ((a, s) -> IO (a, s)) -> (a, s) -> IO (a, s)
forall a b. (a -> b) -> a -> b
$ Get (a, s) -> ByteString -> (a, s)
forall a. Get a -> ByteString -> a
Get.runGet Get (a, s)
getter ByteString
str

------------------------------------------------

getByte :: GetState s Word8
getByte :: GetState s Word8
getByte = Get Word8 -> GetState s Word8
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
State.lift Get Word8
Binary.getWord8

liftOffset :: (Binary.Binary a) => Integer -> (a -> Binary.Put) -> a -> PutState Integer ()
liftOffset :: Integer -> (a -> PutM ()) -> a -> PutState Integer ()
liftOffset Integer
d a -> PutM ()
fn a
x = (Integer -> Integer) -> PutState Integer ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
State.modify (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
d) PutState Integer () -> PutState Integer () -> PutState Integer ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> PutM () -> PutState Integer ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
State.lift (a -> PutM ()
fn a
x)

putByte :: Word8 -> PutState Integer ()
putByte :: Word8 -> PutState Integer ()
putByte Word8
x = Integer -> (Word8 -> PutM ()) -> Word8 -> PutState Integer ()
forall a.
Binary a =>
Integer -> (a -> PutM ()) -> a -> PutState Integer ()
liftOffset Integer
1 Word8 -> PutM ()
Put.putWord8 Word8
x

isEmpty :: GetState s Bool
isEmpty :: GetState s Bool
isEmpty = Get Bool -> GetState s Bool
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
State.lift Get Bool
Get.isEmpty

skip :: Int -> GetState s ()
skip :: Int -> GetState s ()
skip Int
n = Get () -> GetState s ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
State.lift (Int -> Get ()
Get.skip Int
n)

getOffset :: PutState Integer Integer
getOffset :: PutState Integer Integer
getOffset = PutState Integer Integer
forall s (m :: * -> *). MonadState s m => m s
State.get

bytesRead :: GetState s Int64
bytesRead :: GetState s ByteOffset
bytesRead = Get ByteOffset -> GetState s ByteOffset
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
State.lift Get ByteOffset
Get.bytesRead

--------------------------------------------------

instance BinaryState Integer Word8 where
  put :: Word8 -> PutState Integer ()
put Word8
x = Word8 -> PutState Integer ()
putByte Word8
x
  get :: GetState Integer Word8
get = GetState Integer Word8
forall a s. Binary a => GetState s a
getZ

instance BinaryState Integer Word16 where
  put :: Word16 -> PutState Integer ()
put Word16
x = Integer -> (Word16 -> PutM ()) -> Word16 -> PutState Integer ()
forall a.
Binary a =>
Integer -> (a -> PutM ()) -> a -> PutState Integer ()
liftOffset Integer
2 Word16 -> PutM ()
forall t. Binary t => t -> PutM ()
Binary.put Word16
x
  get :: GetState Integer Word16
get = GetState Integer Word16
forall a s. Binary a => GetState s a
getZ

instance BinaryState Integer Word32 where
  put :: Word32 -> PutState Integer ()
put Word32
x = Integer -> (Word32 -> PutM ()) -> Word32 -> PutState Integer ()
forall a.
Binary a =>
Integer -> (a -> PutM ()) -> a -> PutState Integer ()
liftOffset Integer
4 Word32 -> PutM ()
forall t. Binary t => t -> PutM ()
Binary.put Word32
x
  get :: GetState Integer Word32
get = GetState Integer Word32
forall a s. Binary a => GetState s a
getZ

instance (BinaryState s a, BinaryState s b) => BinaryState s (a,b) where
  put :: (a, b) -> PutState s ()
put (a
x,b
y) = a -> PutState s ()
forall s a. BinaryState s a => a -> PutState s ()
put a
x PutState s () -> PutState s () -> PutState s ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> b -> PutState s ()
forall s a. BinaryState s a => a -> PutState s ()
put b
y
  get :: GetState s (a, b)
get = do
    a
x <- GetState s a
forall s a. BinaryState s a => GetState s a
get
    b
y <- GetState s b
forall s a. BinaryState s a => GetState s a
get
    (a, b) -> GetState s (a, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
x,b
y)

--------------------------------------------------

-- instance (Binary.Binary a, Storable a) => BinaryState Integer a where
--   put x = liftOffset (fromIntegral $ sizeOf x) Binary.put x
--   get = getZ