module PtrPoker.Write where

import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Internal as ByteString
import qualified PtrPoker.ByteString as ByteString
import qualified PtrPoker.Ffi as Ffi
import qualified PtrPoker.Poke as Poke
import PtrPoker.Prelude hiding (concat)
import qualified PtrPoker.Size as Size

-- |
-- Execute Write, producing strict ByteString.
{-# INLINEABLE writeToByteString #-}
writeToByteString :: Write -> ByteString
writeToByteString :: Write -> ByteString
writeToByteString Write {Int
Poke
writePoke :: Write -> Poke
writeSize :: Write -> Int
writePoke :: Poke
writeSize :: Int
..} =
  Int -> (Ptr Word8 -> IO ()) -> ByteString
ByteString.unsafeCreate Int
writeSize (forall (f :: * -> *) a. Functor f => f a -> f ()
void forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Poke -> Ptr Word8 -> IO (Ptr Word8)
Poke.pokePtr Poke
writePoke)

-- |
-- Specification of how many bytes to allocate and how to populate them.
--
-- Useful for creating strict bytestrings and tasks like that.
data Write = Write
  { Write -> Int
writeSize :: Int,
    Write -> Poke
writePoke :: Poke.Poke
  }

instance Semigroup Write where
  {-# INLINE (<>) #-}
  Write Int
lSize Poke
lPoke <> :: Write -> Write -> Write
<> Write Int
rSize Poke
rPoke =
    Int -> Poke -> Write
Write (Int
lSize forall a. Num a => a -> a -> a
+ Int
rSize) (Poke
lPoke forall a. Semigroup a => a -> a -> a
<> Poke
rPoke)
  {-# INLINE sconcat #-}
  sconcat :: NonEmpty Write -> Write
sconcat =
    forall (f :: * -> *). Foldable f => f Write -> Write
concat

instance Monoid Write where
  {-# INLINE mempty #-}
  mempty :: Write
mempty =
    Int -> Poke -> Write
Write Int
0 forall a. Monoid a => a
mempty
  {-# INLINE mconcat #-}
  mconcat :: [Write] -> Write
mconcat =
    forall (f :: * -> *). Foldable f => f Write -> Write
concat

-- |
-- Reuses the IsString instance of 'ByteString'.
instance IsString Write where
  {-# INLINE fromString #-}
  fromString :: String -> Write
fromString =
    ByteString -> Write
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 writes.
{-# INLINE concat #-}
concat :: (Foldable f) => f Write -> Write
concat :: forall (f :: * -> *). Foldable f => f Write -> Write
concat f Write
f =
  Int -> Poke -> Write
Write
    (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
a Write
b -> Int
a forall a. Num a => a -> a -> a
+ Write -> Int
writeSize Write
b) Int
0 f Write
f)
    ((Ptr Word8 -> IO (Ptr Word8)) -> Poke
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 Write
write -> Poke -> Ptr Word8 -> IO (Ptr Word8)
Poke.pokePtr (Write -> Poke
writePoke Write
write) Ptr Word8
p) Ptr Word8
p f Write
f))

-- |
-- Render Word8 as byte.
{-# INLINE word8 #-}
word8 :: Word8 -> Write
word8 :: Word8 -> Write
word8 Word8
a =
  Int -> Poke -> Write
Write Int
1 (Word8 -> Poke
Poke.word8 Word8
a)

-- |
-- Render Word16 in Little-endian.
{-# INLINE lWord16 #-}
lWord16 :: Word16 -> Write
lWord16 :: Word16 -> Write
lWord16 Word16
a =
  Int -> Poke -> Write
Write Int
2 (Word16 -> Poke
Poke.lWord16 Word16
a)

-- |
-- Render Word16 in Big-endian.
{-# INLINE bWord16 #-}
bWord16 :: Word16 -> Write
bWord16 :: Word16 -> Write
bWord16 Word16
a =
  Int -> Poke -> Write
Write Int
2 (Word16 -> Poke
Poke.bWord16 Word16
a)

-- |
-- Render Word32 in Little-endian.
{-# INLINE lWord32 #-}
lWord32 :: Word32 -> Write
lWord32 :: Word32 -> Write
lWord32 Word32
a =
  Int -> Poke -> Write
Write Int
4 (Word32 -> Poke
Poke.lWord32 Word32
a)

-- |
-- Render Word32 in Big-endian.
{-# INLINE bWord32 #-}
bWord32 :: Word32 -> Write
bWord32 :: Word32 -> Write
bWord32 Word32
a =
  Int -> Poke -> Write
Write Int
4 (Word32 -> Poke
Poke.bWord32 Word32
a)

-- |
-- Render Word64 in Little-endian.
{-# INLINE lWord64 #-}
lWord64 :: Word64 -> Write
lWord64 :: Word64 -> Write
lWord64 Word64
a =
  Int -> Poke -> Write
Write Int
8 (Word64 -> Poke
Poke.lWord64 Word64
a)

-- |
-- Render Word64 in Big-endian.
{-# INLINE bWord64 #-}
bWord64 :: Word64 -> Write
bWord64 :: Word64 -> Write
bWord64 Word64
a =
  Int -> Poke -> Write
Write Int
8 (Word64 -> Poke
Poke.bWord64 Word64
a)

-- |
-- Render Int16 in Little-endian.
{-# INLINE lInt16 #-}
lInt16 :: Int16 -> Write
lInt16 :: Int16 -> Write
lInt16 Int16
a =
  Int -> Poke -> Write
Write Int
2 (Int16 -> Poke
Poke.lInt16 Int16
a)

-- |
-- Render Int16 in Big-endian.
{-# INLINE bInt16 #-}
bInt16 :: Int16 -> Write
bInt16 :: Int16 -> Write
bInt16 Int16
a =
  Int -> Poke -> Write
Write Int
2 (Int16 -> Poke
Poke.bInt16 Int16
a)

-- |
-- Render Int32 in Little-endian.
{-# INLINE lInt32 #-}
lInt32 :: Int32 -> Write
lInt32 :: Int32 -> Write
lInt32 Int32
a =
  Int -> Poke -> Write
Write Int
4 (Int32 -> Poke
Poke.lInt32 Int32
a)

-- |
-- Render Int32 in Big-endian.
{-# INLINE bInt32 #-}
bInt32 :: Int32 -> Write
bInt32 :: Int32 -> Write
bInt32 Int32
a =
  Int -> Poke -> Write
Write Int
4 (Int32 -> Poke
Poke.bInt32 Int32
a)

-- |
-- Render Int64 in Little-endian.
{-# INLINE lInt64 #-}
lInt64 :: Int64 -> Write
lInt64 :: Int64 -> Write
lInt64 Int64
a =
  Int -> Poke -> Write
Write Int
8 (Int64 -> Poke
Poke.lInt64 Int64
a)

-- |
-- Render Int64 in Big-endian.
{-# INLINE bInt64 #-}
bInt64 :: Int64 -> Write
bInt64 :: Int64 -> Write
bInt64 Int64
a =
  Int -> Poke -> Write
Write Int
8 (Int64 -> Poke
Poke.bInt64 Int64
a)

-- |
-- Render Word64 in ASCII decimal.
{-# INLINE word64AsciiDec #-}
word64AsciiDec :: Word64 -> Write
word64AsciiDec :: Word64 -> Write
word64AsciiDec Word64
a =
  Int -> Poke -> Write
Write Int
size Poke
poke
  where
    size :: Int
size =
      Word64 -> Int
Size.word64AsciiDec Word64
a
    poke :: Poke
poke =
      forall a. Int -> (Ptr Word8 -> IO a) -> Poke
Poke.sizedReverse Int
size (CULLong -> Ptr Word8 -> IO ()
Ffi.revPokeUInt64 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
a))

-- |
-- Render Word in ASCII decimal.
{-# INLINE wordAsciiDec #-}
wordAsciiDec :: Word -> Write
wordAsciiDec :: Word -> Write
wordAsciiDec =
  Word64 -> Write
word64AsciiDec 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

-- |
-- Render Int64 in ASCII decimal.
{-# INLINE int64AsciiDec #-}
int64AsciiDec :: Int64 -> Write
int64AsciiDec :: Int64 -> Write
int64AsciiDec Int64
a =
  Int -> Poke -> Write
Write Int
size Poke
poke
  where
    size :: Int
size =
      Int64 -> Int
Size.int64AsciiDec Int64
a
    poke :: Poke
poke =
      forall a. Int -> (Ptr Word8 -> IO a) -> Poke
Poke.sizedReverse Int
size (CLLong -> Ptr Word8 -> IO ()
Ffi.revPokeInt64 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
a))

-- |
-- Render Int in ASCII decimal.
{-# INLINE intAsciiDec #-}
intAsciiDec :: Int -> Write
intAsciiDec :: Int -> Write
intAsciiDec =
  Int64 -> Write
int64AsciiDec 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

-- |
-- Render double interpreting non-real values,
-- such as @NaN@, @Infinity@, @-Infinity@,
-- as is.
{-# INLINE doubleAsciiDec #-}
doubleAsciiDec :: Double -> Write
doubleAsciiDec :: Double -> Write
doubleAsciiDec Double
a =
  if Double
a forall a. Eq a => a -> a -> Bool
== Double
0
    then Word8 -> Write
word8 Word8
48
    else
      if forall a. RealFloat a => a -> Bool
isNaN Double
a
        then Write
"NaN"
        else
          if forall a. RealFloat a => a -> Bool
isInfinite Double
a
            then
              if Double
a forall a. Ord a => a -> a -> Bool
< Double
0
                then Write
"-Infinity"
                else Write
"Infinity"
            else
              if Double
a forall a. Ord a => a -> a -> Bool
< Double
0
                then Word8 -> Write
word8 Word8
45 forall a. Semigroup a => a -> a -> a
<> ByteString -> Write
byteString (Double -> ByteString
ByteString.double (forall a. Num a => a -> a
negate Double
a))
                else ByteString -> Write
byteString (Double -> ByteString
ByteString.double Double
a)

-- |
-- Render double interpreting non real values,
-- such as @NaN@, @Infinity@, @-Infinity@,
-- as zero.
{-# INLINE zeroNonRealDoubleAsciiDec #-}
zeroNonRealDoubleAsciiDec :: Double -> Write
zeroNonRealDoubleAsciiDec :: Double -> Write
zeroNonRealDoubleAsciiDec Double
a =
  if Double
a forall a. Eq a => a -> a -> Bool
== Double
0 Bool -> Bool -> Bool
|| forall a. RealFloat a => a -> Bool
isNaN Double
a Bool -> Bool -> Bool
|| forall a. RealFloat a => a -> Bool
isInfinite Double
a
    then Word8 -> Write
word8 Word8
48
    else
      if Double
a forall a. Ord a => a -> a -> Bool
< Double
0
        then Word8 -> Write
word8 Word8
45 forall a. Semigroup a => a -> a -> a
<> ByteString -> Write
byteString (Double -> ByteString
ByteString.double (forall a. Num a => a -> a
negate Double
a))
        else ByteString -> Write
byteString (Double -> ByteString
ByteString.double Double
a)

-- |
-- Render Scientific in ASCII decimal.
{-# INLINE scientificAsciiDec #-}
scientificAsciiDec :: Scientific -> Write
scientificAsciiDec :: Scientific -> Write
scientificAsciiDec =
  ByteString -> Write
byteString forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Scientific -> ByteString
ByteString.scientific

-- |
-- Efficiently copy the contents of ByteString using @memcpy@.
{-# INLINE byteString #-}
byteString :: ByteString -> Write
byteString :: ByteString -> Write
byteString ByteString
a =
  Int -> Poke -> Write
Write (ByteString -> Int
ByteString.length ByteString
a) (forall a. a -> a
inline ByteString -> Poke
Poke.byteString ByteString
a)

-- |
-- Render Text in UTF8.
--
-- Does pretty much the same as 'Data.Text.Encoding.encodeUtf8',
-- both implementation and performance-wise,
-- while allowing you to avoid redundant @memcpy@
-- compared to @('byteString' . 'Data.Text.Encoding.encodeUtf8')@.
--
-- Following are the benchmark results comparing the performance of
-- @('writeToByteString' . 'textUtf8')@ with
-- @Data.Text.Encoding.'Data.Text.Encoding.encodeUtf8'@
-- on inputs in Latin and Greek (requiring different number of surrogate bytes).
-- The results show that they are quite similar.
--
-- === __Benchmark results__
--
-- > textUtf8/ptr-poker/latin/1   25.61 ns
-- > textUtf8/ptr-poker/latin/10  31.59 ns
-- > textUtf8/ptr-poker/latin/100 121.5 ns
-- > textUtf8/ptr-poker/greek/1   28.54 ns
-- > textUtf8/ptr-poker/greek/10  41.97 ns
-- > textUtf8/ptr-poker/greek/100 250.3 ns
-- > textUtf8/text/latin/1        22.84 ns
-- > textUtf8/text/latin/10       31.10 ns
-- > textUtf8/text/latin/100      118.2 ns
-- > textUtf8/text/greek/1        25.80 ns
-- > textUtf8/text/greek/10       40.80 ns
-- > textUtf8/text/greek/100      293.1 ns
{-# INLINEABLE textUtf8 #-}
textUtf8 :: Text -> Write
textUtf8 :: Text -> Write
textUtf8 Text
a =
  Int -> Poke -> Write
Write (Text -> Int
Size.textUtf8 Text
a) (Text -> Poke
Poke.textUtf8 Text
a)