module PtrPoker.Write
where

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


{-|
Execute Write, producing strict ByteString.
-}
{-# INLINABLE 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 (IO (Ptr Word8) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Ptr Word8) -> IO ())
-> (Ptr Word8 -> IO (Ptr Word8)) -> Ptr Word8 -> IO ()
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rSize) (Poke
lPoke Poke -> Poke -> Poke
forall a. Semigroup a => a -> a -> a
<> Poke
rPoke)
  {-# INLINE sconcat #-}
  sconcat :: NonEmpty Write -> Write
sconcat =
    NonEmpty Write -> Write
forall (f :: * -> *). Foldable f => f Write -> Write
concat

instance Monoid Write where
  {-# INLINE mempty #-}
  mempty :: Write
mempty =
    Int -> Poke -> Write
Write Int
0 Poke
forall a. Monoid a => a
mempty
  {-# INLINE mconcat #-}
  mconcat :: [Write] -> Write
mconcat =
    [Write] -> Write
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 (ByteString -> Write) -> (String -> ByteString) -> String -> Write
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 writes.
-}
{-# INLINE concat #-}
concat :: Foldable f => f Write -> Write
concat :: f Write -> Write
concat f Write
f =
  Int -> Poke -> Write
Write
    ((Int -> Write -> Int) -> Int -> f Write -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\ Int
a Write
b -> Int
a Int -> Int -> Int
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 -> (Ptr Word8 -> Write -> IO (Ptr Word8))
-> Ptr Word8 -> f Write -> 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 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 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 =
      Int -> (Ptr Word8 -> IO ()) -> Poke
forall a. Int -> (Ptr Word8 -> IO a) -> Poke
Poke.sizedReverse Int
size (CULLong -> Ptr Word8 -> IO ()
Ffi.revPokeUInt64 (Word64 -> CULLong
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 (Word64 -> Write) -> (Word -> Word64) -> Word -> Write
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Word -> Word64
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 =
      Int -> (Ptr Word8 -> IO ()) -> Poke
forall a. Int -> (Ptr Word8 -> IO a) -> Poke
Poke.sizedReverse Int
size (CLLong -> Ptr Word8 -> IO ()
Ffi.revPokeInt64 (Int64 -> CLLong
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 (Int64 -> Write) -> (Int -> Int64) -> Int -> Write
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int -> Int64
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 Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
    then Word8 -> Write
word8 Word8
48
    else if Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
a
      then Write
"NaN"
      else if Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
a
        then if Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0
          then Write
"-Infinity"
          else Write
"Infinity"
        else if Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0
          then Word8 -> Write
word8 Word8
45 Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> ByteString -> Write
byteString (Double -> ByteString
ByteString.double (Double -> 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 Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
a Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
a
    then Word8 -> Write
word8 Word8
48
    else if Double
a Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0
      then Word8 -> Write
word8 Word8
45 Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> ByteString -> Write
byteString (Double -> ByteString
ByteString.double (Double -> 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 (ByteString -> Write)
-> (Scientific -> ByteString) -> Scientific -> Write
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) ((ByteString -> Poke) -> ByteString -> Poke
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               mean 57.06 ns  ( +- 3.283 ns  )
> textUtf8/ptr-poker/latin/10              mean 214.1 ns  ( +- 8.601 ns  )
> textUtf8/ptr-poker/latin/100             mean 1.536 μs  ( +- 75.03 ns  )
> textUtf8/ptr-poker/greek/1               mean 85.98 ns  ( +- 5.038 ns  )
> textUtf8/ptr-poker/greek/10              mean 482.1 ns  ( +- 12.38 ns  )
> textUtf8/ptr-poker/greek/100             mean 4.398 μs  ( +- 33.94 ns  )
> textUtf8/text/latin/1                    mean 60.28 ns  ( +- 3.517 ns  )
> textUtf8/text/latin/10                   mean 201.6 ns  ( +- 8.118 ns  )
> textUtf8/text/latin/100                  mean 1.323 μs  ( +- 51.25 ns  )
> textUtf8/text/greek/1                    mean 99.14 ns  ( +- 1.264 ns  )
> textUtf8/text/greek/10                   mean 483.4 ns  ( +- 5.844 ns  )
> textUtf8/text/greek/100                  mean 4.238 μs  ( +- 40.55 ns  )
-}
{-# INLINABLE textUtf8 #-}
textUtf8 :: Text -> Write
textUtf8 :: Text -> Write
textUtf8 =
  ByteString -> Write
byteString (ByteString -> Write) -> (Text -> ByteString) -> Text -> Write
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Text -> ByteString
ByteString.textUtf8