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.IO.ByteString as ByteStringIO
import qualified PtrPoker.IO.Prim as PrimIO
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 (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 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 =
      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 51.54 ns  ( +- 3.083 ns  )
-- > textUtf8/ptr-poker/latin/10              mean 132.8 ns  ( +- 14.75 ns  )
-- > textUtf8/ptr-poker/latin/100             mean 860.6 ns  ( +- 66.61 ns  )
-- > textUtf8/ptr-poker/greek/1               mean 106.4 ns  ( +- 19.28 ns  )
-- > textUtf8/ptr-poker/greek/10              mean 498.4 ns  ( +- 8.022 ns  )
-- > textUtf8/ptr-poker/greek/100             mean 4.462 μs  ( +- 31.58 ns  )
-- > textUtf8/text/latin/1                    mean 52.77 ns  ( +- 3.311 ns  )
-- > textUtf8/text/latin/10                   mean 206.1 ns  ( +- 26.78 ns  )
-- > textUtf8/text/latin/100                  mean 1.337 μs  ( +- 43.34 ns  )
-- > textUtf8/text/greek/1                    mean 88.22 ns  ( +- 1.119 ns  )
-- > textUtf8/text/greek/10                   mean 475.2 ns  ( +- 21.15 ns  )
-- > textUtf8/text/greek/100                  mean 4.252 μs  ( +- 64.33 ns  )
{-# INLINEABLE 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