{-|
Description : Utilities for packing stuff
Maintainer  : srk <srk@48.io>
|-}
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 :: 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 :: 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 :: (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 (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 :: Get a -> Get [a]
getMany Get a
parser = do
  Int
count <- Get Int
forall a. Integral a => Get a
getInt
  Int -> Get a -> Get [a]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
count Get a
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 :: 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
  Int64
len <- Get Int64
forall a. Integral a => Get a
getInt
  ByteString
st  <- Int64 -> Get ByteString
getLazyByteString Int64
len
  Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int64
len Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`mod` Int64
8 Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Int64
0) (Get () -> Get ()) -> Get () -> Get ()
forall a b. (a -> b) -> a -> b
$ do
    [Word8]
pads <- Int -> Get [Word8]
unpad (Int -> Get [Word8]) -> Int -> Get [Word8]
forall a b. (a -> b) -> a -> b
$ Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int) -> Int64 -> Int
forall a b. (a -> b) -> a -> b
$ Int64
8 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- (Int64
len Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`mod` Int64
8)
    Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ((Word8 -> Bool) -> [Word8] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0) [Word8]
pads) (Get () -> Get ()) -> Get () -> Get ()
forall a b. (a -> b) -> a -> b
$ String -> Get ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get ()) -> String -> Get ()
forall a b. (a -> b) -> a -> b
$ String
"No zeroes" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> (ByteString, Int64, [Word8]) -> String
forall b a. (Show a, IsString b) => a -> b
show (ByteString
st, Int64
len, [Word8]
pads)
  ByteString -> Get ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Get ByteString) -> ByteString -> Get ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
forall l s. LazyStrict l s => l -> s
toStrict ByteString
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