{-# LANGUAGE UnboxedTuples #-}

module Bytezap where

import GHC.Exts
import Data.ByteString qualified as B
import Data.ByteString.Internal qualified as B
import GHC.IO
import Data.Word

-- | TODO inner poke type
--
-- TODO can I change this to
--
-- @
-- Ptr Word8 -> IO (Ptr Word8)
-- @
--
-- without any performance loss? it's the same underneath newtypes and datas.
-- 'Ptr' is a data rather than a newtype, but IO is just a newtype.
--
-- I originally did this to beat ptr-poker, but idk. Now doubtful.
type Poke# = Addr# -> State# RealWorld -> (# State# RealWorld, Addr# #)

-- | Unboxed poke operation.
--
-- A newtype allows us a monoidal interface.
newtype Poke = Poke
  { -- | Write at an offset from an address and return the next offset.
    --
    -- The returned offset must be after the argument offset.
    --
    -- TODO I use that output order because it matches IO. Probs doesn't matter.
    Poke -> Poke#
unPoke :: Poke#
  }

-- | Construct a 'Poke'.
poke :: Poke# -> Poke
poke :: Poke# -> Poke
poke = Poke# -> Poke
Poke
{-# INLINE poke #-}

-- | Sequence two 'Poke's left-to-right.
instance Semigroup Poke where
    {-# INLINE (<>) #-}
    Poke Poke#
l <> :: Poke -> Poke -> Poke
<> Poke Poke#
r = Poke# -> Poke
Poke (Poke# -> Poke) -> Poke# -> Poke
forall a b. (a -> b) -> a -> b
$ \Addr#
addr# State# RealWorld
st# ->
        case Poke#
l Addr#
addr# State# RealWorld
st# of (# State# RealWorld
st'#, Addr#
addr'# #) -> Poke#
r Addr#
addr'# State# RealWorld
st'#

-- | The empty 'Poke' simply returns its arguments.
instance Monoid Poke where
    {-# INLINE mempty #-}
    mempty :: Poke
mempty = Poke# -> Poke
Poke (Poke# -> Poke) -> Poke# -> Poke
forall a b. (a -> b) -> a -> b
$ \Addr#
addr# State# RealWorld
st# -> (# State# RealWorld
st#, Addr#
addr# #)

-- | Allocate a buffer of the given size and run a 'Poke' over it.
--
-- The 'Poke' must fill the buffer exactly. If it goes under, you should get
-- some random garbage at the end. If it goes over, your computer will probably
-- explode.
runPoke :: Int -> Poke -> B.ByteString
runPoke :: Int -> Poke -> ByteString
runPoke Int
len = Int -> (Ptr Word8 -> IO ()) -> ByteString
B.unsafeCreate Int
len ((Ptr Word8 -> IO ()) -> ByteString)
-> (Poke -> Ptr Word8 -> IO ()) -> Poke -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Poke -> Ptr Word8 -> IO ()
wrapPoke
{-# INLINE runPoke #-}

wrapPoke :: Poke -> Ptr Word8 -> IO ()
wrapPoke :: Poke -> Ptr Word8 -> IO ()
wrapPoke (Poke Poke#
p) (Ptr Addr#
addr#) =
    (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
st# -> case Poke#
p Addr#
addr# State# RealWorld
st# of (# State# RealWorld
l, Addr#
_r #) -> (# State# RealWorld
l, () #))
{-# INLINE wrapPoke #-}

-- | Instructions on how to perform a sized write.
--
-- The 'Poke' in 'writePoke' must write the _exact_ number of bytes specified in
-- 'writeSize'. Otherwise, your computer explodes.
data Write = Write
  { Write -> Int
writeSize :: {-# UNPACK #-} !Int
  , Write -> Poke
writePoke :: !Poke -- unpack unusable TODO is strict good or not here
  }

-- | Construct a 'Write'.
write :: Int -> Poke# -> Write
write :: Int -> Poke# -> Write
write Int
len Poke#
p = Int -> Poke -> Write
Write Int
len (Poke# -> Poke
Poke Poke#
p)
{-# INLINE write #-}

-- | Sequence the 'Poke's, sum the sizes.
instance Semigroup Write where
    -- TODO feels like this might be INLINE[1] or even INLINE[0]?
    {-# INLINE (<>) #-}
    Write Int
ll Poke
lp <> :: Write -> Write -> Write
<> Write Int
rl Poke
rp = Int -> Poke -> Write
Write (Int
ll Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rl) (Poke
lp Poke -> Poke -> Poke
forall a. Semigroup a => a -> a -> a
<> Poke
rp)

-- | The empty 'Write' is the empty 'Poke', which writes zero bytes.
instance Monoid Write where
    {-# INLINE mempty #-}
    mempty :: Write
mempty = Int -> Poke -> Write
Write Int
0 Poke
forall a. Monoid a => a
mempty

-- | Serialize and show the resulting ByteString.
instance Show Write where showsPrec :: Int -> Write -> ShowS
showsPrec Int
p = Int -> ByteString -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
p (ByteString -> ShowS) -> (Write -> ByteString) -> Write -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Write -> ByteString
runWrite

runWrite :: Write -> B.ByteString
runWrite :: Write -> ByteString
runWrite (Write Int
len Poke
p) = Int -> Poke -> ByteString
runPoke Int
len Poke
p
{-# INLINE runWrite #-}