{-# LANGUAGE RecordWildCards #-}

-- safe module, only export the safe bits (no @Write(..)@!!)
module Bytezap.Write
  ( Write(writeLength, writeOp)
  , LengthType(ExactLength, MaxLength)
  , runWriteBS, runWriteBSUptoN
  , prim, byteString, byteArray#, replicateByte
  ) where

import Bytezap.Write.Internal

import Bytezap.Poke qualified as Poke
import Bytezap.Poke ( Poke )
import Raehik.Compat.Data.Primitive.Types
import GHC.Exts
import Data.ByteString qualified as BS
import Data.Word ( Word8 )

runWriteBS :: Write ExactLength RealWorld -> BS.ByteString
runWriteBS :: Write 'ExactLength RealWorld -> ByteString
runWriteBS = (Int -> Poke RealWorld -> ByteString)
-> Write 'ExactLength RealWorld -> ByteString
forall a s (lt :: LengthType).
(Int -> Poke s -> a) -> Write lt s -> a
runWriteWith Int -> Poke RealWorld -> ByteString
Poke.unsafeRunPokeBS

runWriteBSUptoN :: Write MaxLength RealWorld -> BS.ByteString
runWriteBSUptoN :: Write 'MaxLength RealWorld -> ByteString
runWriteBSUptoN = (Int -> Poke RealWorld -> ByteString)
-> Write 'MaxLength RealWorld -> ByteString
forall a s (lt :: LengthType).
(Int -> Poke s -> a) -> Write lt s -> a
runWriteWith Int -> Poke RealWorld -> ByteString
Poke.unsafeRunPokeBSUptoN

-- | Helper for writing 'Write' runners.
runWriteWith :: forall a s lt. (Int -> Poke s -> a) -> Write lt s -> a
runWriteWith :: forall a s (lt :: LengthType).
(Int -> Poke s -> a) -> Write lt s -> a
runWriteWith Int -> Poke s -> a
runPoke (Write Int
l Poke s
p) = Int -> Poke s -> a
runPoke Int
l Poke s
p

prim :: forall a s. Prim' a => a -> Write ExactLength s
prim :: forall a s. Prim' a => a -> Write 'ExactLength s
prim a
a = Int -> Poke s -> Write 'ExactLength s
forall (lt :: LengthType) s. Int -> Poke s -> Write lt s
Write (a -> Int
forall a. Prim a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) (a -> Poke s
forall a s. Prim' a => a -> Poke s
Poke.prim a
a)

byteString :: BS.ByteString -> Write ExactLength RealWorld
byteString :: ByteString -> Write 'ExactLength RealWorld
byteString ByteString
bs = Int -> Poke RealWorld -> Write 'ExactLength RealWorld
forall (lt :: LengthType) s. Int -> Poke s -> Write lt s
Write (ByteString -> Int
BS.length ByteString
bs) (ByteString -> Poke RealWorld
Poke.byteString ByteString
bs)

byteArray# :: ByteArray# -> Int# -> Int# -> Write ExactLength s
byteArray# :: forall s. ByteArray# -> Int# -> Int# -> Write 'ExactLength s
byteArray# ByteArray#
ba# Int#
baos# Int#
balen# = Write{Int
Poke s
writeLength :: Int
writeOp :: Poke s
writeLength :: Int
writeOp :: Poke s
..}
  where
    writeLength :: Int
writeLength = Int# -> Int
I# Int#
balen#
    writeOp :: Poke s
writeOp     = ByteArray# -> Int# -> Int# -> Poke s
forall s. ByteArray# -> Int# -> Int# -> Poke s
Poke.byteArray# ByteArray#
ba# Int#
baos# Int#
balen#

-- | essentially memset
replicateByte :: Int -> Word8 -> Write ExactLength RealWorld
replicateByte :: Int -> Word8 -> Write 'ExactLength RealWorld
replicateByte Int
len Word8
byte = Int -> Poke RealWorld -> Write 'ExactLength RealWorld
forall (lt :: LengthType) s. Int -> Poke s -> Write lt s
Write Int
len (Int -> Word8 -> Poke RealWorld
Poke.replicateByte Int
len Word8
byte)