{-|
Description : Utilities for packing stuff
Maintainer  : srk <[email protected]>
|-}
module System.Nix.Store.Remote.Binary where

import           Data.Binary.Get
import           Data.Binary.Put
import qualified Data.ByteString.Lazy          as BSL

putInt :: Integral a => a -> Put
putInt :: forall a. Integral a => a -> Put
putInt = Word64 -> Put
putWord64le (Word64 -> Put) -> (a -> Word64) -> a -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

getInt :: Integral a => Get a
getInt :: forall a. Integral a => Get a
getInt = Word64 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> a) -> Get Word64 -> Get a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word64
getWord64le

putMany :: Foldable t => (a -> Put) -> t a -> Put
putMany :: forall (t :: * -> *) a. Foldable t => (a -> Put) -> t a -> Put
putMany a -> Put
printer t a
xs = do
  Int -> Put
forall a. Integral a => a -> Put
putInt (t a -> Int
forall a. t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
xs)
  (a -> Put) -> t a -> Put
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ a -> Put
printer t a
xs

getMany :: Get a -> Get [a]
getMany :: forall a. Get a -> Get [a]
getMany Get a
parser = do
  count <- Get Int
forall a. Integral a => Get a
getInt
  replicateM count parser

-- length prefixed string packing with padding to 8 bytes
putByteStringLen :: BSL.ByteString -> Put
putByteStringLen :: ByteString -> Put
putByteStringLen ByteString
x = do
  Int -> Put
forall a. Integral a => a -> Put
putInt Int
len
  ByteString -> Put
putLazyByteString ByteString
x
  Bool -> Put -> Put
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
8 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (Put -> Put) -> Put -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Put
pad (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
8)
 where
  len :: Int
  len :: Int
len = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int) -> Int64 -> Int
forall a b. (a -> b) -> a -> b
$ ByteString -> Int64
BSL.length ByteString
x
  pad :: Int -> Put
pad Int
count = Int -> Put -> Put
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
count (Word8 -> Put
putWord8 Word8
0)

putByteStrings :: Foldable t => t BSL.ByteString -> Put
putByteStrings :: forall (t :: * -> *). Foldable t => t ByteString -> Put
putByteStrings = (ByteString -> Put) -> t ByteString -> Put
forall (t :: * -> *) a. Foldable t => (a -> Put) -> t a -> Put
putMany ByteString -> Put
putByteStringLen

getByteStringLen :: Get ByteString
getByteStringLen :: Get ByteString
getByteStringLen = do
  len <- Get Int64
forall a. Integral a => Get a
getInt
  st  <- getLazyByteString len
  when (len `mod` 8 /= 0) $ do
    pads <- unpad $ fromIntegral $ 8 - (len `mod` 8)
    unless (all (== 0) pads) $ fail $ "No zeroes" <> show (st, len, pads)
  pure $ toStrict st
  where unpad :: Int -> Get [Word8]
unpad Int
x = Int -> Get Word8 -> Get [Word8]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
x Get Word8
getWord8

getByteStrings :: Get [ByteString]
getByteStrings :: Get [ByteString]
getByteStrings = Get ByteString -> Get [ByteString]
forall a. Get a -> Get [a]
getMany Get ByteString
getByteStringLen