{-# LANGUAGE CPP #-}

module PtrPoker.Poke where

import qualified Data.Text.Array as TextArray
import qualified Data.Text.Internal as TextInternal
import qualified PtrPoker.Ffi as Ffi
import qualified PtrPoker.IO.ByteString as ByteStringIO
import qualified PtrPoker.IO.Prim as PrimIO
import PtrPoker.Prelude hiding (concat)
import qualified PtrPoker.Text as Text

{-# RULES
"foldMap" forall f foldable.
  foldMap f foldable =
    Poke $ \p -> foldM (\p (Poke poker) -> poker p) p foldable
  #-}

-- |
-- Abstraction over an IO action,
-- which takes a pointer, populates it and
-- produces a pointer right after the populated data.
newtype Poke = Poke {Poke -> Ptr Word8 -> IO (Ptr Word8)
pokePtr :: Ptr Word8 -> IO (Ptr Word8)}

instance Semigroup Poke where
  {-# INLINE [1] (<>) #-}
  Poke Ptr Word8 -> IO (Ptr Word8)
lIO <> :: Poke -> Poke -> Poke
<> Poke Ptr Word8 -> IO (Ptr Word8)
rIO =
    (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (Ptr Word8 -> IO (Ptr Word8)
lIO (Ptr Word8 -> IO (Ptr Word8))
-> (Ptr Word8 -> IO (Ptr Word8)) -> Ptr Word8 -> IO (Ptr Word8)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Ptr Word8 -> IO (Ptr Word8)
rIO)
  sconcat :: NonEmpty Poke -> Poke
sconcat =
    NonEmpty Poke -> Poke
forall (f :: * -> *). Foldable f => f Poke -> Poke
concat

instance Monoid Poke where
  {-# INLINE [1] mempty #-}
  mempty :: Poke
mempty =
    (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke Ptr Word8 -> IO (Ptr Word8)
forall (m :: * -> *) a. Monad m => a -> m a
return
  mconcat :: [Poke] -> Poke
mconcat =
    [Poke] -> Poke
forall (f :: * -> *). Foldable f => f Poke -> Poke
concat

-- |
-- Reuses the IsString instance of 'ByteString'.
instance IsString Poke where
  fromString :: String -> Poke
fromString = ByteString -> Poke
byteString (ByteString -> Poke) -> (String -> ByteString) -> String -> Poke
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. String -> ByteString
forall a. IsString a => String -> a
fromString

-- |
-- Concatenate a foldable of pokes.
{-# INLINE [1] concat #-}
concat :: Foldable f => f Poke -> Poke
concat :: f Poke -> Poke
concat f Poke
pokers =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (\Ptr Word8
p -> (Ptr Word8 -> Poke -> IO (Ptr Word8))
-> Ptr Word8 -> f Poke -> IO (Ptr Word8)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Ptr Word8
p (Poke Ptr Word8 -> IO (Ptr Word8)
io) -> Ptr Word8 -> IO (Ptr Word8)
io Ptr Word8
p) Ptr Word8
p f Poke
pokers)

-- |
-- Efficiently copy the contents of ByteString using @memcpy@.
{-# INLINE byteString #-}
byteString :: ByteString -> Poke
byteString :: ByteString -> Poke
byteString ByteString
bs =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke ((Ptr Word8 -> IO (Ptr Word8)) -> Poke)
-> (Ptr Word8 -> IO (Ptr Word8)) -> Poke
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> (Ptr Word8 -> ByteString -> IO (Ptr Word8))
-> Ptr Word8 -> ByteString -> IO (Ptr Word8)
forall a. a -> a
inline Ptr Word8 -> ByteString -> IO (Ptr Word8)
ByteStringIO.pokeByteString Ptr Word8
ptr ByteString
bs

-- |
-- Encode Word8 as byte, incrementing the pointer by 1.
{-# INLINE [1] word8 #-}
word8 :: Word8 -> Poke
word8 :: Word8 -> Poke
word8 Word8
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (\Ptr Word8
p -> Ptr Word8 -> Word8 -> IO ()
PrimIO.pokeWord8 Ptr Word8
p Word8
a IO () -> Ptr Word8 -> IO (Ptr Word8)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
p Int
1)

-- |
-- Encode Word16 in Little-endian.
{-# INLINE [1] lWord16 #-}
lWord16 :: Word16 -> Poke
lWord16 :: Word16 -> Poke
lWord16 Word16
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (\Ptr Word8
p -> Ptr Word8 -> Word16 -> IO ()
PrimIO.pokeLEWord16 Ptr Word8
p Word16
a IO () -> Ptr Word8 -> IO (Ptr Word8)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
p Int
2)

-- |
-- Encode Word16 in Big-endian.
{-# INLINE [1] bWord16 #-}
bWord16 :: Word16 -> Poke
bWord16 :: Word16 -> Poke
bWord16 Word16
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (\Ptr Word8
p -> Ptr Word8 -> Word16 -> IO ()
PrimIO.pokeBEWord16 Ptr Word8
p Word16
a IO () -> Ptr Word8 -> IO (Ptr Word8)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
p Int
2)

-- |
-- Encode Word32 in Little-endian.
{-# INLINE [1] lWord32 #-}
lWord32 :: Word32 -> Poke
lWord32 :: Word32 -> Poke
lWord32 Word32
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (\Ptr Word8
p -> Ptr Word8 -> Word32 -> IO ()
PrimIO.pokeLEWord32 Ptr Word8
p Word32
a IO () -> Ptr Word8 -> IO (Ptr Word8)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
p Int
4)

-- |
-- Encode Word32 in Big-endian.
{-# INLINE [1] bWord32 #-}
bWord32 :: Word32 -> Poke
bWord32 :: Word32 -> Poke
bWord32 Word32
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (\Ptr Word8
p -> Ptr Word8 -> Word32 -> IO ()
PrimIO.pokeBEWord32 Ptr Word8
p Word32
a IO () -> Ptr Word8 -> IO (Ptr Word8)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
p Int
4)

-- |
-- Encode Word64 in Little-endian.
{-# INLINE [1] lWord64 #-}
lWord64 :: Word64 -> Poke
lWord64 :: Word64 -> Poke
lWord64 Word64
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (\Ptr Word8
p -> Ptr Word8 -> Word64 -> IO ()
PrimIO.pokeLEWord64 Ptr Word8
p Word64
a IO () -> Ptr Word8 -> IO (Ptr Word8)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
p Int
8)

-- |
-- Encode Word64 in Big-endian.
{-# INLINE [1] bWord64 #-}
bWord64 :: Word64 -> Poke
bWord64 :: Word64 -> Poke
bWord64 Word64
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (\Ptr Word8
p -> Ptr Word8 -> Word64 -> IO ()
PrimIO.pokeBEWord64 Ptr Word8
p Word64
a IO () -> Ptr Word8 -> IO (Ptr Word8)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
p Int
8)

-- |
-- Encode Int16 in Little-endian.
{-# INLINE lInt16 #-}
lInt16 :: Int16 -> Poke
lInt16 :: Int16 -> Poke
lInt16 = Word16 -> Poke
lWord16 (Word16 -> Poke) -> (Int16 -> Word16) -> Int16 -> Poke
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int16 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- |
-- Encode Int16 in Big-endian.
{-# INLINE bInt16 #-}
bInt16 :: Int16 -> Poke
bInt16 :: Int16 -> Poke
bInt16 = Word16 -> Poke
bWord16 (Word16 -> Poke) -> (Int16 -> Word16) -> Int16 -> Poke
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int16 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- |
-- Encode Int32 in Little-endian.
{-# INLINE lInt32 #-}
lInt32 :: Int32 -> Poke
lInt32 :: Int32 -> Poke
lInt32 = Word32 -> Poke
lWord32 (Word32 -> Poke) -> (Int32 -> Word32) -> Int32 -> Poke
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int32 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- |
-- Encode Int32 in Big-endian.
{-# INLINE bInt32 #-}
bInt32 :: Int32 -> Poke
bInt32 :: Int32 -> Poke
bInt32 = Word32 -> Poke
bWord32 (Word32 -> Poke) -> (Int32 -> Word32) -> Int32 -> Poke
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int32 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- |
-- Encode Int64 in Little-endian.
{-# INLINE lInt64 #-}
lInt64 :: Int64 -> Poke
lInt64 :: Int64 -> Poke
lInt64 = Word64 -> Poke
lWord64 (Word64 -> Poke) -> (Int64 -> Word64) -> Int64 -> Poke
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- |
-- Encode Int64 in Big-endian.
{-# INLINE bInt64 #-}
bInt64 :: Int64 -> Poke
bInt64 :: Int64 -> Poke
bInt64 = Word64 -> Poke
bWord64 (Word64 -> Poke) -> (Int64 -> Word64) -> Int64 -> Poke
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- |
-- Encode Text in UTF8.
{-# INLINE textUtf8 #-}
textUtf8 :: Text -> Poke
#if MIN_VERSION_text(2,0,0)
textUtf8 (TextInternal.Text arr off len) =
  Poke (\p -> do
    stToIO $ TextArray.copyToPointer arr off p len
    pure (plusPtr p len))
#else
textUtf8 :: Text -> Poke
textUtf8 = (ByteArray# -> Int -> Int -> Poke) -> Text -> Poke
forall x. (ByteArray# -> Int -> Int -> x) -> Text -> x
Text.destruct ((ByteArray# -> Int -> Int -> Poke) -> Text -> Poke)
-> (ByteArray# -> Int -> Int -> Poke) -> Text -> Poke
forall a b. (a -> b) -> a -> b
$ \ByteArray#
arr Int
off Int
len ->
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (\Ptr Word8
p -> Ptr Word8 -> ByteArray# -> CSize -> CSize -> IO (Ptr Word8)
Ffi.encodeText Ptr Word8
p ByteArray#
arr (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
off) (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len))
#endif

-- * ASCII integers

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

-- |
-- Encode Int8 as a signed ASCII decimal.
{-# INLINE [1] int8AsciiDec #-}
int8AsciiDec :: Int8 -> Poke
int8AsciiDec :: Int8 -> Poke
int8AsciiDec Int8
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (CInt -> Ptr Word8 -> IO (Ptr Word8)
Ffi.pokeIntInDec (Int8 -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
a))

-- |
-- Encode Int16 as a signed ASCII decimal.
{-# INLINE [1] int16AsciiDec #-}
int16AsciiDec :: Int16 -> Poke
int16AsciiDec :: Int16 -> Poke
int16AsciiDec Int16
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (CInt -> Ptr Word8 -> IO (Ptr Word8)
Ffi.pokeIntInDec (Int16 -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
a))

-- |
-- Encode Int32 as a signed ASCII decimal.
{-# INLINE [1] int32AsciiDec #-}
int32AsciiDec :: Int32 -> Poke
int32AsciiDec :: Int32 -> Poke
int32AsciiDec Int32
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (CInt -> Ptr Word8 -> IO (Ptr Word8)
Ffi.pokeIntInDec (Int32 -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
a))

-- |
-- Encode Int64 as a signed ASCII decimal.
{-# INLINE [1] int64AsciiDec #-}
int64AsciiDec :: Int64 -> Poke
int64AsciiDec :: Int64 -> Poke
int64AsciiDec Int64
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (CLLong -> Ptr Word8 -> IO (Ptr Word8)
Ffi.pokeLongLongIntInDec (Int64 -> CLLong
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
a))

-- |
-- Encode Int as a signed ASCII decimal.
{-# INLINE [1] intAsciiDec #-}
intAsciiDec :: Int -> Poke
intAsciiDec :: Int -> Poke
intAsciiDec Int
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (CLLong -> Ptr Word8 -> IO (Ptr Word8)
Ffi.pokeLongLongIntInDec (Int -> CLLong
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
a))

-- |
-- Encode Word8 as an unsigned ASCII decimal.
{-# INLINE [1] word8AsciiDec #-}
word8AsciiDec :: Word8 -> Poke
word8AsciiDec :: Word8 -> Poke
word8AsciiDec Word8
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (CUInt -> Ptr Word8 -> IO (Ptr Word8)
Ffi.pokeUIntInDec (Word8 -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
a))

-- |
-- Encode Word16 as an unsigned ASCII decimal.
{-# INLINE [1] word16AsciiDec #-}
word16AsciiDec :: Word16 -> Poke
word16AsciiDec :: Word16 -> Poke
word16AsciiDec Word16
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (CUInt -> Ptr Word8 -> IO (Ptr Word8)
Ffi.pokeUIntInDec (Word16 -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
a))

-- |
-- Encode Word32 as an unsigned ASCII decimal.
{-# INLINE [1] word32AsciiDec #-}
word32AsciiDec :: Word32 -> Poke
word32AsciiDec :: Word32 -> Poke
word32AsciiDec Word32
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (CUInt -> Ptr Word8 -> IO (Ptr Word8)
Ffi.pokeUIntInDec (Word32 -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
a))

-- |
-- Encode Word64 as an unsigned ASCII decimal.
{-# INLINE [1] word64AsciiDec #-}
word64AsciiDec :: Word64 -> Poke
word64AsciiDec :: Word64 -> Poke
word64AsciiDec Word64
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (CULLong -> Ptr Word8 -> IO (Ptr Word8)
Ffi.pokeLongLongUIntInDec (Word64 -> CULLong
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
a))

-- |
-- Encode Word as an unsigned ASCII decimal.
{-# INLINE [1] wordAsciiDec #-}
wordAsciiDec :: Word -> Poke
wordAsciiDec :: Word -> Poke
wordAsciiDec Word
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (CULLong -> Ptr Word8 -> IO (Ptr Word8)
Ffi.pokeLongLongUIntInDec (Word -> CULLong
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
a))

-- |
-- Encode Double as a signed ASCII decimal.
{-# INLINE doubleAsciiDec #-}
doubleAsciiDec :: Double -> Poke
doubleAsciiDec :: Double -> Poke
doubleAsciiDec Double
a =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke ((Ptr Word8 -> IO (Ptr Word8)) -> Poke)
-> (Ptr Word8 -> IO (Ptr Word8)) -> Poke
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr ->
    Double -> Ptr Word8 -> IO CInt
Ffi.pokeDouble Double
a Ptr Word8
ptr
      IO CInt -> (IO CInt -> IO (Ptr Word8)) -> IO (Ptr Word8)
forall a b. a -> (a -> b) -> b
& (CInt -> Ptr Word8) -> IO CInt -> IO (Ptr Word8)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
ptr (Int -> Ptr Word8) -> (CInt -> Int) -> CInt -> Ptr Word8
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral)

-- * Low level

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

-- |
-- Having the amount of bytes to be written precomputed,
-- executes an action,
-- which fills the pointer going downward,
-- starting from the pointer that follows the chunk.
-- I.e., you have to decrement the pointer
-- before writing the first byte,
-- decrement it again before writing the second byte and so on.
{-# INLINE sizedReverse #-}
sizedReverse :: Int -> (Ptr Word8 -> IO a) -> Poke
sizedReverse :: Int -> (Ptr Word8 -> IO a) -> Poke
sizedReverse Int
size Ptr Word8 -> IO a
action =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke ((Ptr Word8 -> IO (Ptr Word8)) -> Poke)
-> (Ptr Word8 -> IO (Ptr Word8)) -> Poke
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr ->
    let afterPtr :: Ptr Word8
afterPtr =
          Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
ptr Int
size
     in Ptr Word8 -> IO a
action Ptr Word8
afterPtr IO a -> Ptr Word8 -> IO (Ptr Word8)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ptr Word8
afterPtr