{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE KindSignatures             #-}
-- |
--
-- Module      : Raaz.Core.Transfer
-- Description : Type safe transfer of bytes.
-- Copyright   : (c) Piyush P Kurur, 2019
-- License     : Apache-2.0 OR BSD-3-Clause
-- Maintainer  : Piyush P Kurur <ppk@iitpkd.ac.in>
-- Stability   : experimental
--

module Raaz.Core.Transfer
       ( -- * Transfer actions.
         -- $transfer$
         Transfer, ReadFrom, WriteTo
       , consume, consumeStorable, consumeParse
       , writeEncodable
       , write, writeStorable, writeVector, writeStorableVector
       , writeBytes
       , padWrite, prependWrite, glueWrites
       , writeByteString
       , transferSize
       , skip
       ) where


import qualified Data.Vector.Generic       as G
import           Foreign.Storable          ( Storable, poke )

import           Raaz.Core.Transfer.Unsafe
import           Raaz.Core.Prelude
import           Raaz.Core.Parse.Unsafe
import           Raaz.Core.Parse     hiding (skip)
import           Raaz.Core.Types.Endian
import           Raaz.Core.Types.Pointer
import           Raaz.Core.Encode


-- | The transfer @skip l@ skip ahead by an offset @l@. If it is a
-- read, it does not read the next @l@ positions. If it is a write it
-- does not mutate the next @l@ positions.
skip :: LengthUnit l => l -> Transfer t
skip :: l -> Transfer t
skip = (l -> (Ptr Word8 -> IO ()) -> Transfer t)
-> (Ptr Word8 -> IO ()) -> l -> Transfer t
forall a b c. (a -> b -> c) -> b -> a -> c
flip l -> (Ptr Word8 -> IO ()) -> Transfer t
forall u (t :: Mode).
LengthUnit u =>
u -> (Ptr Word8 -> IO ()) -> Transfer t
unsafeMakeTransfer Ptr Word8 -> IO ()
forall b. b -> IO ()
doNothing
       where doNothing :: b -> IO ()
doNothing = IO () -> b -> IO ()
forall a b. a -> b -> a
const (IO () -> b -> IO ()) -> IO () -> b -> IO ()
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()


-------------------------- Monoids for consuming stuff ------------------------------------

-- | Given a parser @p :: Parser a@ for parsing @a@ and @act :: a -> m
-- b@ consuming a, @consumeParse p act@, gives a reader that parses a
-- from the input buffer passing it to the action act.
consumeParse ::  Parser a -> (a -> IO b) -> ReadFrom
consumeParse :: Parser a -> (a -> IO b) -> ReadFrom
consumeParse Parser a
p a -> IO b
action = BYTES Int -> (Ptr Word8 -> IO ()) -> ReadFrom
forall u (t :: Mode).
LengthUnit u =>
u -> (Ptr Word8 -> IO ()) -> Transfer t
unsafeMakeTransfer (Parser a -> BYTES Int
forall a. Parser a -> BYTES Int
parseWidth Parser a
p) ((Ptr Word8 -> IO ()) -> ReadFrom)
-> (Ptr Word8 -> IO ()) -> ReadFrom
forall a b. (a -> b) -> a -> b
$
                        Parser a -> Ptr Word8 -> IO a
forall (ptr :: * -> *) a b.
Pointer ptr =>
Parser a -> ptr b -> IO a
unsafeRunParser Parser a
p (Ptr Word8 -> IO a) -> (a -> IO ()) -> Ptr Word8 -> IO ()
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> IO b -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO b -> IO ()) -> (a -> IO b) -> a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> IO b
action

-- | Reads @a@ from the buffer and supplies it to the action. The
-- value read is independent of the endianness of the underlying.
consume :: EndianStore a
        => (a -> IO b)
        -> ReadFrom
consume :: (a -> IO b) -> ReadFrom
consume = Parser a -> (a -> IO b) -> ReadFrom
forall a b. Parser a -> (a -> IO b) -> ReadFrom
consumeParse Parser a
forall a. EndianStore a => Parser a
parse

-- | Similar to @consume@ but does not take care of adjusting for
-- endianness. Use therefore limited to internal buffers.
consumeStorable :: Storable a
                => (a -> IO b)
                -> ReadFrom
consumeStorable :: (a -> IO b) -> ReadFrom
consumeStorable = Parser a -> (a -> IO b) -> ReadFrom
forall a b. Parser a -> (a -> IO b) -> ReadFrom
consumeParse Parser a
forall a. Storable a => Parser a
parseStorable

-------------------------- Monoid for writing stuff --------------------------------------

-- | The expression @`writeStorable` a@ gives a write action that
-- stores a value @a@ in machine endian. The type of the value @a@ has
-- to be an instance of `Storable`. This should be used when we want
-- to talk with C functions and not when talking to the outside world
-- (otherwise this could lead to endian confusion). To take care of
-- endianness use the `write` combinator.
writeStorable :: Storable a => a -> WriteTo
writeStorable :: a -> WriteTo
writeStorable a
a = BYTES Int -> (Ptr Word8 -> IO ()) -> WriteTo
forall u (t :: Mode).
LengthUnit u =>
u -> (Ptr Word8 -> IO ()) -> Transfer t
unsafeMakeTransfer (Proxy a -> BYTES Int
forall a. Storable a => Proxy a -> BYTES Int
sizeOf (Proxy a -> BYTES Int) -> Proxy a -> BYTES Int
forall a b. (a -> b) -> a -> b
$ a -> Proxy a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a) Ptr Word8 -> IO ()
forall a. Ptr a -> IO ()
pokeIt
  where pokeIt :: Ptr a -> IO ()
pokeIt = (Ptr a -> a -> IO ()) -> a -> Ptr a -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke a
a (Ptr a -> IO ()) -> (Ptr a -> Ptr a) -> Ptr a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr a -> Ptr a
forall (ptr :: * -> *) a b. Pointer ptr => ptr a -> ptr b
castPointer
-- | The expression @`write` a@ gives a write action that stores a
-- value @a@. One needs the type of the value @a@ to be an instance of
-- `EndianStore`. Proper endian conversion is done irrespective of
-- what the machine endianness is. The man use of this write is to
-- serialize data for the consumption of the outside world.
write :: EndianStore a => a -> WriteTo
write :: a -> WriteTo
write a
a = BYTES Int -> (Ptr Word8 -> IO ()) -> WriteTo
forall u (t :: Mode).
LengthUnit u =>
u -> (Ptr Word8 -> IO ()) -> Transfer t
unsafeMakeTransfer (Proxy a -> BYTES Int
forall a. Storable a => Proxy a -> BYTES Int
sizeOf (Proxy a -> BYTES Int) -> Proxy a -> BYTES Int
forall a b. (a -> b) -> a -> b
$ a -> Proxy a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a) ((Ptr Word8 -> IO ()) -> WriteTo)
-> (Ptr Word8 -> IO ()) -> WriteTo
forall a b. (a -> b) -> a -> b
$ (Ptr Word8 -> a -> IO ()) -> a -> Ptr Word8 -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Ptr a -> a -> IO ()
forall w. EndianStore w => Ptr w -> w -> IO ()
store (Ptr a -> a -> IO ())
-> (Ptr Word8 -> Ptr a) -> Ptr Word8 -> a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr Word8 -> Ptr a
forall (ptr :: * -> *) a b. Pointer ptr => ptr a -> ptr b
castPointer) a
a


-- | Write any encodable elements
writeEncodable :: Encodable a => a -> WriteTo
writeEncodable :: a -> WriteTo
writeEncodable = ByteString -> WriteTo
writeByteString (ByteString -> WriteTo) -> (a -> ByteString) -> a -> WriteTo
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
forall a. Encodable a => a -> ByteString
toByteString

-- | The vector version of `writeStorable`.
writeStorableVector :: (Storable a, G.Vector v a) => v a -> WriteTo
{-# INLINE writeStorableVector #-}
writeStorableVector :: v a -> WriteTo
writeStorableVector = (WriteTo -> a -> WriteTo) -> WriteTo -> v a -> WriteTo
forall (v :: * -> *) b a.
Vector v b =>
(a -> b -> a) -> a -> v b -> a
G.foldl' WriteTo -> a -> WriteTo
forall a. Storable a => WriteTo -> a -> WriteTo
foldFunc WriteTo
forall a. Monoid a => a
mempty
  where foldFunc :: WriteTo -> a -> WriteTo
foldFunc WriteTo
w a
a =  WriteTo
w WriteTo -> WriteTo -> WriteTo
forall a. Semigroup a => a -> a -> a
<> a -> WriteTo
forall a. Storable a => a -> WriteTo
writeStorable a
a

{-

TODO: This function can be slow due to the fact that each time we use
the semi-direct product, we incur a cost due to the lambda being not
lifted.

-}

-- | The vector version of `write`.
writeVector :: (EndianStore a, G.Vector v a) => v a -> WriteTo
{-# INLINE writeVector #-}
{- TODO: improve this using the fact that the size is known -}
writeVector :: v a -> WriteTo
writeVector = (WriteTo -> a -> WriteTo) -> WriteTo -> v a -> WriteTo
forall (v :: * -> *) b a.
Vector v b =>
(a -> b -> a) -> a -> v b -> a
G.foldl' WriteTo -> a -> WriteTo
forall a. EndianStore a => WriteTo -> a -> WriteTo
foldFunc WriteTo
forall a. Monoid a => a
mempty
  where foldFunc :: WriteTo -> a -> WriteTo
foldFunc WriteTo
w a
a =  WriteTo
w WriteTo -> WriteTo -> WriteTo
forall a. Semigroup a => a -> a -> a
<> a -> WriteTo
forall a. EndianStore a => a -> WriteTo
write a
a
{- TODO: Same as in writeStorableVector -}


-- | The combinator @writeBytes b n@ writes @b@ as the next @n@
-- consecutive bytes.
writeBytes :: LengthUnit n
           => Word8   -- ^ Byte to write
           -> n       -- ^ How much to write
           -> WriteTo
writeBytes :: Word8 -> n -> WriteTo
writeBytes Word8
w8 n
n = n -> (Ptr Word8 -> IO ()) -> WriteTo
forall u (t :: Mode).
LengthUnit u =>
u -> (Ptr Word8 -> IO ()) -> Transfer t
unsafeMakeTransfer n
n Ptr Word8 -> IO ()
forall (ptr :: * -> *) a. Pointer ptr => ptr a -> IO ()
memsetIt
  where memsetIt :: ptr a -> IO ()
memsetIt ptr a
cptr = ptr a -> Word8 -> n -> IO ()
forall l (ptr :: * -> *) a.
(LengthUnit l, Pointer ptr) =>
ptr a -> Word8 -> l -> IO ()
memset ptr a
cptr Word8
w8 n
n

-- | The combinator @glueWrites w n hdr ftr@ is equivalent to @hdr <>
-- glue <> ftr@ where the write @glue@ writes just enough bytes @w@ so
-- that the total length is aligned to the boundary @n@.
glueWrites :: LengthUnit n
           => Word8    -- ^ The bytes to use in the glue
           -> n        -- ^ The length boundary to align to.
           -> WriteTo  -- ^ The header write
           -> WriteTo  -- ^ The footer write
           -> WriteTo
glueWrites :: Word8 -> n -> WriteTo -> WriteTo -> WriteTo
glueWrites Word8
w8 n
n WriteTo
hdr WriteTo
ftr = WriteTo
hdr WriteTo -> WriteTo -> WriteTo
forall a. Semigroup a => a -> a -> a
<> Word8 -> BYTES Int -> WriteTo
forall n. LengthUnit n => Word8 -> n -> WriteTo
writeBytes Word8
w8 BYTES Int
lglue WriteTo -> WriteTo -> WriteTo
forall a. Semigroup a => a -> a -> a
<> WriteTo
ftr
  where lhead :: BYTES Int
lhead   = WriteTo -> BYTES Int
forall (t :: Mode). Transfer t -> BYTES Int
transferSize WriteTo
hdr
        lfoot :: BYTES Int
lfoot   = WriteTo -> BYTES Int
forall (t :: Mode). Transfer t -> BYTES Int
transferSize WriteTo
ftr
        lexceed :: BYTES Int
lexceed = (BYTES Int
lhead BYTES Int -> BYTES Int -> BYTES Int
forall a. Num a => a -> a -> a
+ BYTES Int
lfoot) BYTES Int -> BYTES Int -> BYTES Int
forall a. Integral a => a -> a -> a
`rem` BYTES Int
nBytes  -- bytes exceeding the boundary.
        lglue :: BYTES Int
lglue   = if BYTES Int
lexceed BYTES Int -> BYTES Int -> Bool
forall a. Ord a => a -> a -> Bool
> BYTES Int
0 then BYTES Int
nBytes BYTES Int -> BYTES Int -> BYTES Int
forall a. Num a => a -> a -> a
- BYTES Int
lexceed else BYTES Int
0
        nBytes :: BYTES Int
nBytes  = n -> BYTES Int
forall u. LengthUnit u => u -> BYTES Int
inBytes n
n

-- | The write action @prependWrite w n wr@ is wr pre-pended with the byte @w@ so that the total length
-- ends at a multiple of @n@.
prependWrite  :: LengthUnit n
              => Word8     -- ^ the byte to pre-pend with.
              -> n         -- ^ the length to align the message to
              -> WriteTo  -- ^ the message that needs pre-pending
              -> WriteTo
prependWrite :: Word8 -> n -> WriteTo -> WriteTo
prependWrite Word8
w8 n
n = Word8 -> n -> WriteTo -> WriteTo -> WriteTo
forall n.
LengthUnit n =>
Word8 -> n -> WriteTo -> WriteTo -> WriteTo
glueWrites Word8
w8 n
n WriteTo
forall a. Monoid a => a
mempty

-- | The write action @padWrite w n wr@ is wr padded with the byte @w@ so that the total length
-- ends at a multiple of @n@.
padWrite :: LengthUnit n
         => Word8     -- ^ the padding byte to use
         -> n         -- ^ the length to align message to
         -> WriteTo   -- ^ the message that needs padding
         -> WriteTo
padWrite :: Word8 -> n -> WriteTo -> WriteTo
padWrite Word8
w8 n
n = (WriteTo -> WriteTo -> WriteTo) -> WriteTo -> WriteTo -> WriteTo
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Word8 -> n -> WriteTo -> WriteTo -> WriteTo
forall n.
LengthUnit n =>
Word8 -> n -> WriteTo -> WriteTo -> WriteTo
glueWrites Word8
w8 n
n) WriteTo
forall a. Monoid a => a
mempty


-------------  Reading stuff  -----------------------------------