{-# LANGUAGE BangPatterns #-}

module Parser.Lathe.Internal.ByteString
  ( toStrictLen
  , toShortLen
  ) where

import           Control.Monad.ST
import           Control.Monad.ST.Unsafe
import           Data.ByteString.Internal as B (ByteString (..), unsafeCreate)
import qualified Data.ByteString.Unsafe as B (unsafeUseAsCStringLen)
import           Data.ByteString.Lazy.Internal as L (ByteString (..))
import           Data.ByteString.Short.Internal (ShortByteString (..))
import           Data.Primitive.ByteArray
import           Foreign.Marshal.Utils (copyBytes)
import           Foreign.Ptr



-- | Convert a lazy ByteString of known size to a strict one.
toStrictLen :: Int -> L.ByteString -> B.ByteString
toStrictLen :: Int -> ByteString -> ByteString
toStrictLen Int
total ByteString
xs =
  Int -> (Ptr Word8 -> IO ()) -> ByteString
B.unsafeCreate Int
total ((Ptr Word8 -> IO ()) -> ByteString)
-> (Ptr Word8 -> IO ()) -> ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dest0 ->

    let go :: Ptr b -> ByteString -> IO ()
go !Ptr b
dest ByteString
bs =
          case ByteString
bs of
            L.Chunk ByteString
b ByteString
cs -> do
              Int
len <- ByteString -> (CStringLen -> IO Int) -> IO Int
forall a. ByteString -> (CStringLen -> IO a) -> IO a
B.unsafeUseAsCStringLen ByteString
b ((CStringLen -> IO Int) -> IO Int)
-> (CStringLen -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
src, Int
len) -> do
                       Ptr b -> Ptr b -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr b
dest (Ptr CChar -> Ptr b
forall a b. Ptr a -> Ptr b
castPtr Ptr CChar
src) Int
len
                       Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
len

              Ptr b -> ByteString -> IO ()
go (Ptr b -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr b
dest Int
len) ByteString
cs

            ByteString
L.Empty      -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    in Ptr Word8 -> ByteString -> IO ()
forall {b}. Ptr b -> ByteString -> IO ()
go Ptr Word8
dest0 ByteString
xs



-- | Convert a lazy ByteString of known size to a short one.
toShortLen :: Int -> L.ByteString -> ShortByteString
toShortLen :: Int -> ByteString -> ShortByteString
toShortLen Int
total ByteString
xs =
  (forall s. ST s ShortByteString) -> ShortByteString
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s ShortByteString) -> ShortByteString)
-> (forall s. ST s ShortByteString) -> ShortByteString
forall a b. (a -> b) -> a -> b
$ do
    MutableByteArray s
marr <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
total

    let go :: Int -> ByteString -> ST s ()
go !Int
n ByteString
bs =
          case ByteString
bs of
            L.Chunk ByteString
b ByteString
cs -> do
              Int
len <- IO Int -> ST s Int
forall a s. IO a -> ST s a
unsafeIOToST (IO Int -> ST s Int)
-> ((CStringLen -> IO Int) -> IO Int)
-> (CStringLen -> IO Int)
-> ST s Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                       ByteString -> (CStringLen -> IO Int) -> IO Int
forall a. ByteString -> (CStringLen -> IO a) -> IO a
B.unsafeUseAsCStringLen ByteString
b ((CStringLen -> IO Int) -> ST s Int)
-> (CStringLen -> IO Int) -> ST s Int
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
src, Int
len) -> do
                         ST s () -> IO ()
forall s a. ST s a -> IO a
unsafeSTToIO (ST s () -> IO ()) -> ST s () -> IO ()
forall a b. (a -> b) -> a -> b
$ MutableByteArray (PrimState (ST s))
-> Int -> Ptr CChar -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutableByteArray (PrimState m) -> Int -> Ptr a -> Int -> m ()
copyPtrToMutableByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
marr Int
n Ptr CChar
src Int
len
                         Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
len

              Int -> ByteString -> ST s ()
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) ByteString
cs

            ByteString
L.Empty      -> () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    Int -> ByteString -> ST s ()
forall {s}. Int -> ByteString -> ST s ()
go Int
0 ByteString
xs
    (\(ByteArray ByteArray#
arr) -> ByteArray# -> ShortByteString
SBS ByteArray#
arr) (ByteArray -> ShortByteString)
-> ST s ByteArray -> ST s ShortByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
marr