-- |
-- Module      : Data.ByteArray.Pack
-- License     : BSD-Style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : unknown
--
-- Simple Byte Array packer
--
-- Simple example:
--
-- > > flip pack 20 $ putWord8 0x41 >> putByteString "BCD" >> putWord8 0x20 >> putStorable (42 :: Word32)
-- > Right (ABCD *\NUL\NUL\NUL")
--
-- Original code from <https://hackage.haskell.org/package/bspack>
-- generalized and adapted to run on 'memory', and spellchecked / tweaked. (2015-05)
-- Copyright (c) 2014 Nicolas DI PRIMA
--
module Data.ByteArray.Pack
    ( Packer
    , Result(..)
    , fill
    , pack
      -- * Operations
      -- ** put
    , putWord8
    , putWord16
    , putWord32
    , putStorable
    , putBytes
    , fillList
    , fillUpWith
      -- ** skip
    , skip
    , skipStorable
    ) where

import           Data.Word
import           Foreign.Ptr
import           Foreign.Storable
import           Data.Memory.Internal.Imports ()
import           Data.Memory.Internal.Compat
import           Data.Memory.PtrMethods
import           Data.ByteArray.Pack.Internal
import           Data.ByteArray (ByteArray, ByteArrayAccess, MemView(..))
import qualified Data.ByteArray as B

-- | Fill a given sized buffer with the result of the Packer action
fill :: ByteArray byteArray => Int -> Packer a -> Either String byteArray
fill :: forall byteArray a.
ByteArray byteArray =>
Int -> Packer a -> Either String byteArray
fill Int
len Packer a
packing = IO (Either String byteArray) -> Either String byteArray
forall a. IO a -> a
unsafeDoIO (IO (Either String byteArray) -> Either String byteArray)
-> IO (Either String byteArray) -> Either String byteArray
forall a b. (a -> b) -> a -> b
$ do
    (Result a
val, byteArray
out) <- Int -> (Ptr Word8 -> IO (Result a)) -> IO (Result a, byteArray)
forall ba p a. ByteArray ba => Int -> (Ptr p -> IO a) -> IO (a, ba)
forall p a. Int -> (Ptr p -> IO a) -> IO (a, byteArray)
B.allocRet Int
len ((Ptr Word8 -> IO (Result a)) -> IO (Result a, byteArray))
-> (Ptr Word8 -> IO (Result a)) -> IO (Result a, byteArray)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> Packer a -> MemView -> IO (Result a)
forall a. Packer a -> MemView -> IO (Result a)
runPacker_ Packer a
packing (Ptr Word8 -> Int -> MemView
MemView Ptr Word8
ptr Int
len)
    case Result a
val of 
        PackerMore a
_ (MemView Ptr Word8
_ Int
r)
            | Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0    -> Either String byteArray -> IO (Either String byteArray)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String byteArray -> IO (Either String byteArray))
-> Either String byteArray -> IO (Either String byteArray)
forall a b. (a -> b) -> a -> b
$ byteArray -> Either String byteArray
forall a b. b -> Either a b
Right byteArray
out
            | Bool
otherwise -> Either String byteArray -> IO (Either String byteArray)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String byteArray -> IO (Either String byteArray))
-> Either String byteArray -> IO (Either String byteArray)
forall a b. (a -> b) -> a -> b
$ String -> Either String byteArray
forall a b. a -> Either a b
Left (String
"remaining unpacked bytes " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
r String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" at the end of buffer")
        PackerFail String
err  -> Either String byteArray -> IO (Either String byteArray)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String byteArray -> IO (Either String byteArray))
-> Either String byteArray -> IO (Either String byteArray)
forall a b. (a -> b) -> a -> b
$ String -> Either String byteArray
forall a b. a -> Either a b
Left String
err

-- | Pack the given packer into the given bytestring
pack :: ByteArray byteArray => Packer a -> Int -> Either String byteArray
pack :: forall byteArray a.
ByteArray byteArray =>
Packer a -> Int -> Either String byteArray
pack Packer a
packing Int
len = Int -> Packer a -> Either String byteArray
forall byteArray a.
ByteArray byteArray =>
Int -> Packer a -> Either String byteArray
fill Int
len Packer a
packing
{-# DEPRECATED pack "use fill instead" #-}

fillUpWithWord8' :: Word8 -> Packer ()
fillUpWithWord8' :: Word8 -> Packer ()
fillUpWithWord8' Word8
w = (MemView -> IO (Result ())) -> Packer ()
forall a. (MemView -> IO (Result a)) -> Packer a
Packer ((MemView -> IO (Result ())) -> Packer ())
-> (MemView -> IO (Result ())) -> Packer ()
forall a b. (a -> b) -> a -> b
$ \(MemView Ptr Word8
ptr Int
size) -> do
    Ptr Word8 -> Word8 -> Int -> IO ()
memSet Ptr Word8
ptr Word8
w Int
size
    Result () -> IO (Result ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Result () -> IO (Result ())) -> Result () -> IO (Result ())
forall a b. (a -> b) -> a -> b
$ () -> MemView -> Result ()
forall a. a -> MemView -> Result a
PackerMore () (Ptr Word8 -> Int -> MemView
MemView (Ptr Word8
ptr Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size) Int
0)

-- | Put a storable from the current position in the stream
putStorable :: Storable storable => storable -> Packer ()
putStorable :: forall storable. Storable storable => storable -> Packer ()
putStorable storable
s = Int -> (Ptr Word8 -> IO ()) -> Packer ()
forall a. Int -> (Ptr Word8 -> IO a) -> Packer a
actionPacker (storable -> Int
forall a. Storable a => a -> Int
sizeOf storable
s) (\Ptr Word8
ptr -> Ptr storable -> storable -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Ptr Word8 -> Ptr storable
forall a b. Ptr a -> Ptr b
castPtr Ptr Word8
ptr) storable
s)

-- | Put a Byte Array from the current position in the stream
--
-- If the ByteArray is null, then do nothing
putBytes :: ByteArrayAccess ba => ba -> Packer ()
putBytes :: forall ba. ByteArrayAccess ba => ba -> Packer ()
putBytes ba
bs
    | Int
neededLength Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = () -> Packer ()
forall a. a -> Packer a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    | Bool
otherwise         =
        Int -> (Ptr Word8 -> IO ()) -> Packer ()
forall a. Int -> (Ptr Word8 -> IO a) -> Packer a
actionPacker Int
neededLength ((Ptr Word8 -> IO ()) -> Packer ())
-> (Ptr Word8 -> IO ()) -> Packer ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dstPtr -> ba -> (Ptr Word8 -> IO ()) -> IO ()
forall ba p a. ByteArrayAccess ba => ba -> (Ptr p -> IO a) -> IO a
forall p a. ba -> (Ptr p -> IO a) -> IO a
B.withByteArray ba
bs ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
srcPtr ->
            Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memCopy Ptr Word8
dstPtr Ptr Word8
srcPtr Int
neededLength
  where
    neededLength :: Int
neededLength = ba -> Int
forall ba. ByteArrayAccess ba => ba -> Int
B.length ba
bs

-- | Skip some bytes from the current position in the stream
skip :: Int -> Packer ()
skip :: Int -> Packer ()
skip Int
n = Int -> (Ptr Word8 -> IO ()) -> Packer ()
forall a. Int -> (Ptr Word8 -> IO a) -> Packer a
actionPacker Int
n (\Ptr Word8
_ -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())

-- | Skip the size of a storable from the current position in the stream
skipStorable :: Storable storable => storable -> Packer ()
skipStorable :: forall storable. Storable storable => storable -> Packer ()
skipStorable = Int -> Packer ()
skip (Int -> Packer ()) -> (storable -> Int) -> storable -> Packer ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. storable -> Int
forall a. Storable a => a -> Int
sizeOf

-- | Fill up from the current position in the stream to the end
--
-- It is equivalent to:
--
-- > fillUpWith s == fillList (repeat s)
--
fillUpWith :: Storable storable => storable -> Packer ()
fillUpWith :: forall storable. Storable storable => storable -> Packer ()
fillUpWith storable
s = [storable] -> Packer ()
forall storable. Storable storable => [storable] -> Packer ()
fillList ([storable] -> Packer ()) -> [storable] -> Packer ()
forall a b. (a -> b) -> a -> b
$ storable -> [storable]
forall a. a -> [a]
repeat storable
s
{-# RULES "fillUpWithWord8" forall s . fillUpWith s = fillUpWithWord8' s #-}
{-# NOINLINE fillUpWith #-}

-- | Will put the given storable list from the current position in the stream
-- to the end.
--
-- This function will fail with not enough storage if the given storable can't
-- be written (not enough space)
--
-- Example:
--
-- > > pack (fillList $ [1..] :: Word8) 9
-- > "\1\2\3\4\5\6\7\8\9"
-- > > pack (fillList $ [1..] :: Word32) 4
-- > "\1\0\0\0"
-- > > pack (fillList $ [1..] :: Word32) 64
-- > .. <..succesful..>
-- > > pack (fillList $ [1..] :: Word32) 1
-- > .. <.. not enough space ..>
-- > > pack (fillList $ [1..] :: Word32) 131
-- > .. <.. not enough space ..>
--
fillList :: Storable storable => [storable] -> Packer ()
fillList :: forall storable. Storable storable => [storable] -> Packer ()
fillList []     = () -> Packer ()
forall a. a -> Packer a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
fillList (storable
x:[storable]
xs) = storable -> Packer ()
forall storable. Storable storable => storable -> Packer ()
putStorable storable
x Packer () -> Packer () -> Packer ()
forall a b. Packer a -> Packer b -> Packer b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [storable] -> Packer ()
forall storable. Storable storable => [storable] -> Packer ()
fillList [storable]
xs

------------------------------------------------------------------------------
-- Common packer                                                            --
------------------------------------------------------------------------------

-- | put Word8 in the current position in the stream
putWord8 :: Word8 -> Packer ()
putWord8 :: Word8 -> Packer ()
putWord8 = Word8 -> Packer ()
forall storable. Storable storable => storable -> Packer ()
putStorable
{-# INLINE putWord8 #-}

-- | put Word16 in the current position in the stream
-- /!\ use Host Endianness
putWord16 :: Word16 -> Packer ()
putWord16 :: Word16 -> Packer ()
putWord16 = Word16 -> Packer ()
forall storable. Storable storable => storable -> Packer ()
putStorable
{-# INLINE putWord16 #-}

-- | put Word32 in the current position in the stream
-- /!\ use Host Endianness
putWord32 :: Word32 -> Packer ()
putWord32 :: Word32 -> Packer ()
putWord32 = Word32 -> Packer ()
forall storable. Storable storable => storable -> Packer ()
putStorable
{-# INLINE putWord32 #-}