module PtrPoker.Poke where

import qualified PtrPoker.Compat.ByteString as ByteStringCompat
import qualified PtrPoker.Compat.Text as TextCompat
import qualified PtrPoker.Ffi as Ffi
import qualified PtrPoker.IO.Prim as PrimIO
import PtrPoker.Prelude hiding (concat)

{-# 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
p -> Ptr Word8 -> IO (Ptr Word8)
lIO Ptr Word8
p forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Ptr Word8 -> IO (Ptr Word8)
rIO)
  sconcat :: NonEmpty Poke -> Poke
sconcat =
    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 forall (m :: * -> *) a. Monad m => a -> m a
return
  mconcat :: [Poke] -> Poke
mconcat =
    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 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. IsString a => String -> a
fromString

-- |
-- Concatenate a foldable of pokes.
{-# INLINE [1] concat #-}
concat :: (Foldable f) => f Poke -> Poke
concat :: forall (f :: * -> *). Foldable f => f Poke -> Poke
concat f Poke
pokers =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke (\Ptr Word8
p -> 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 forall a b. (a -> b) -> a -> b
$ forall a. a -> a
inline ByteString -> Ptr Word8 -> IO (Ptr Word8)
ByteStringCompat.poke 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 forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> 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 forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> 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 forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> 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 forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> 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 forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> 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 forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> 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 forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> 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 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. 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 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. 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 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. 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 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. 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 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. 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 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- |
-- Encode Text in UTF8.
{-# INLINE textUtf8 #-}
textUtf8 :: Text -> Poke
textUtf8 :: Text -> Poke
textUtf8 = (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Text -> Ptr Word8 -> IO (Ptr Word8)
TextCompat.pokeInUtf8

-- * 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 (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 (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 (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 (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 (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 (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 (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 (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 (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 (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 forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr ->
    Double -> Ptr Word8 -> IO CInt
Ffi.pokeDouble Double
a Ptr Word8
ptr
      forall a b. a -> (a -> b) -> b
& forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
ptr forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. 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 :: forall a. Int -> (Ptr Word8 -> IO a) -> Poke
sizedReverse Int
size Ptr Word8 -> IO a
action =
  (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr ->
    let afterPtr :: Ptr Word8
afterPtr =
          forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
ptr Int
size
     in Ptr Word8 -> IO a
action Ptr Word8
afterPtr forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Ptr Word8
afterPtr