{-# LANGUAGE CPP             #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE MagicHash #-}

module Streamly.External.ByteString
  ( toArray
  , fromArray

  , reader

  , writeN
  , write

  -- Deprecated
  , read
  )
where

import Control.Monad.IO.Class (MonadIO)
import Data.Word (Word8)
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Storable (peek)
import GHC.Exts
    ( Addr#
    , MutableByteArray#
    , RealWorld
    , byteArrayContents#
    , minusAddr#
    , plusAddr#
    , unsafeCoerce#
    )
import GHC.ForeignPtr (ForeignPtr(..), ForeignPtrContents(..))
import GHC.Int (Int(..))
import GHC.Ptr (Ptr(..), nullPtr, plusPtr)
import Streamly.Data.Fold (Fold)
import Streamly.Data.Unfold (Unfold, lmap)

-- Internal imports
import Data.ByteString.Internal (ByteString(..))
import Streamly.Internal.System.IO (unsafeInlineIO)

import qualified Streamly.Data.Array as Array
import qualified Streamly.Internal.Data.Unfold as Unfold (fold, mkUnfoldrM)

#if !(MIN_VERSION_bytestring(0,11,0))
import GHC.ForeignPtr (plusForeignPtr)
#endif

#if MIN_VERSION_streamly_core(0,2,0)
import Streamly.Internal.Data.Array (Array(..))
import Streamly.Internal.Data.MutByteArray (MutByteArray(..))
import qualified Streamly.Internal.Data.Array as Array
import qualified Streamly.Internal.Data.MutByteArray as MutBA
import qualified Streamly.Internal.Data.Stream as StreamD (Step(Yield))
#else
import Streamly.Internal.Data.Array.Type (Array(..))
import Streamly.Internal.Data.Unboxed (MutableByteArray(..))
import qualified Streamly.Internal.Data.Unboxed as MutBA
import qualified Streamly.Internal.Data.Stream.StreamD as StreamD (Step(Yield))
#endif

import Prelude hiding (read)

#if MIN_VERSION_streamly_core(0,2,0)
#define MUT_BYTE_ARRAY MutByteArray
#else
#define MUT_BYTE_ARRAY MutableByteArray
#endif

#if MIN_VERSION_streamly_core(0,2,2)
#define NIL MutBA.empty
#else
#define NIL MutBA.nil
#endif

#if MIN_VERSION_bytestring(0,11,0)
#define CONSTRUCTOR(a, b, c) BS a c
#define WHEN_0_10_12(x)
#else
#define CONSTRUCTOR(a, b, c) PS a b c
#define WHEN_0_10_12(x) x
#endif


{-# INLINE ensurePinned #-}
ensurePinned :: Array a -> IO (Array a)
{-# INLINE pinnedCreateOf #-}
pinnedCreateOf :: MonadIO m => Int -> Fold m Word8 (Array Word8)
{-# INLINE pinnedCreate #-}
pinnedCreate :: MonadIO m => Fold m Word8 (Array Word8)

#if MIN_VERSION_streamly_core(0,2,2)
ensurePinned :: forall a. Array a -> IO (Array a)
ensurePinned = Array a -> IO (Array a)
forall a. Array a -> IO (Array a)
Array.pin
pinnedCreateOf :: forall (m :: * -> *).
MonadIO m =>
Int -> Fold m Word8 (Array Word8)
pinnedCreateOf = Int -> Fold m Word8 (Array Word8)
forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (Array a)
Array.pinnedCreateOf
pinnedCreate :: forall (m :: * -> *). MonadIO m => Fold m Word8 (Array Word8)
pinnedCreate = Fold m Word8 (Array Word8)
forall (m :: * -> *) a. (MonadIO m, Unbox a) => Fold m a (Array a)
Array.pinnedCreate
#elif MIN_VERSION_streamly_core(0,2,0)
ensurePinned = Array.pin
pinnedCreateOf = Array.pinnedWriteN
pinnedCreate = Array.pinnedWrite
#else
ensurePinned = pure
pinnedCreateOf = Array.writeN
pinnedCreate = Array.write
#endif

{-# INLINE mutableByteArrayContents# #-}
mutableByteArrayContents# :: MutableByteArray# RealWorld -> Addr#
mutableByteArrayContents# :: MutableByteArray# RealWorld -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
marr# = ByteArray# -> Addr#
byteArrayContents# (MutableByteArray# RealWorld -> ByteArray#
forall a b. a -> b
unsafeCoerce# MutableByteArray# RealWorld
marr#)

-- | Helper function that creates a ForeignPtr
{-# INLINE makeForeignPtr #-}
makeForeignPtr :: MUT_BYTE_ARRAY -> Int -> ForeignPtr a
makeForeignPtr :: forall a. MutByteArray -> Int -> ForeignPtr a
makeForeignPtr (MUT_BYTE_ARRAY marr#) (I# off#) =
    Addr# -> ForeignPtrContents -> ForeignPtr a
forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr
        (MutableByteArray# RealWorld -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
marr# Addr# -> Int# -> Addr#
`plusAddr#` Int#
off#)
        (MutableByteArray# RealWorld -> ForeignPtrContents
PlainPtr MutableByteArray# RealWorld
marr#)

-- | Convert a 'ByteString' to an array of 'Word8'. It can be done in constant
-- time only for GHC allocated memory. For foreign allocator allocated memory
-- there is a copy involved.
{-# INLINE toArray #-}
toArray :: ByteString -> Array Word8
toArray :: ByteString -> Array Word8
toArray (CONSTRUCTOR((ForeignPtr addr# _), _, _))
    | Addr# -> Ptr Any
forall a. Addr# -> Ptr a
Ptr Addr#
addr# Ptr Any -> Ptr Any -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Any
forall a. Ptr a
nullPtr = MutByteArray -> Int -> Int -> Array Word8
forall a. MutByteArray -> Int -> Int -> Array a
Array NIL 0 0
toArray (CONSTRUCTOR((ForeignPtr addr# (PlainPtr marr#)), off0, len)) =
    let off :: Int
off = Int# -> Int
I# (Addr#
addr# Addr# -> Addr# -> Int#
`minusAddr#` MutableByteArray# RealWorld -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
marr#)
                  WHEN_0_10_12(+ off0)
     in MutByteArray -> Int -> Int -> Array Word8
forall a. MutByteArray -> Int -> Int -> Array a
Array (MUT_BYTE_ARRAY marr#) off (off + len)
toArray (CONSTRUCTOR(fptr, off, len)) =
    unsafeInlineIO
        $ withForeignPtr (fptr WHEN_0_10_12(`plusForeignPtr` off))
        $ Unfold.fold (Array.writeN len) generator

    where

    generator =
        Unfold.mkUnfoldrM
            (\ptr -> flip StreamD.Yield (ptr `plusPtr` 1) <$> peek ptr)

-- | Convert an array of 'Word8' to a 'ByteString'.
--
-- Please ensure that the array is pinned when using this function.

-- If the array is pinned, the operation is performed in constant time. Whereas
-- for an unpinned array a copy is involved to pin it.
--
{-# INLINE fromArray #-}
fromArray :: Array Word8 -> ByteString
fromArray :: Array Word8 -> ByteString
fromArray Array Word8
arr
    | Int
aLen Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = ByteString
forall a. Monoid a => a
mempty
    | Bool
otherwise = IO ByteString -> ByteString
forall a. IO a -> a
unsafeInlineIO (IO ByteString -> ByteString) -> IO ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ do
        Array{Int
MutByteArray
arrContents :: MutByteArray
arrStart :: Int
arrEnd :: Int
arrContents :: forall a. Array a -> MutByteArray
arrStart :: forall a. Array a -> Int
arrEnd :: forall a. Array a -> Int
..} <- Array Word8 -> IO (Array Word8)
forall a. Array a -> IO (Array a)
ensurePinned Array Word8
arr
        ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ CONSTRUCTOR((makeForeignPtr arrContents arrStart), 0, aLen)

    where

    aLen :: Int
aLen = Array Word8 -> Int
forall a. Array a -> Int
arrEnd Array Word8
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Array Word8 -> Int
forall a. Array a -> Int
arrStart Array Word8
arr

-- | Unfold a strict ByteString to a stream of Word8.
{-# INLINE reader #-}
reader :: Monad m => Unfold m ByteString Word8
reader :: forall (m :: * -> *). Monad m => Unfold m ByteString Word8
reader = (ByteString -> Array Word8)
-> Unfold m (Array Word8) Word8 -> Unfold m ByteString Word8
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap ByteString -> Array Word8
toArray Unfold m (Array Word8) Word8
forall (m :: * -> *) a. (Monad m, Unbox a) => Unfold m (Array a) a
Array.reader

-- | Fold a stream of Word8 to a strict ByteString of given size in bytes.
{-# INLINE writeN #-}
writeN :: MonadIO m => Int -> Fold m Word8 ByteString
writeN :: forall (m :: * -> *). MonadIO m => Int -> Fold m Word8 ByteString
writeN Int
i = Array Word8 -> ByteString
fromArray (Array Word8 -> ByteString)
-> Fold m Word8 (Array Word8) -> Fold m Word8 ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Fold m Word8 (Array Word8)
forall (m :: * -> *).
MonadIO m =>
Int -> Fold m Word8 (Array Word8)
pinnedCreateOf Int
i

-- | Fold a stream of Word8 to a strict ByteString of appropriate size.
{-# INLINE write #-}
write :: MonadIO m => Fold m Word8 ByteString
write :: forall (m :: * -> *). MonadIO m => Fold m Word8 ByteString
write = Array Word8 -> ByteString
fromArray (Array Word8 -> ByteString)
-> Fold m Word8 (Array Word8) -> Fold m Word8 ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Fold m Word8 (Array Word8)
forall (m :: * -> *). MonadIO m => Fold m Word8 (Array Word8)
pinnedCreate

--------------------------------------------------------------------------------
-- Deprecated
--------------------------------------------------------------------------------

{-# DEPRECATED read "Please use reader instead." #-}
{-# INLINE read #-}
read :: Monad m => Unfold m ByteString Word8
read :: forall (m :: * -> *). Monad m => Unfold m ByteString Word8
read = Unfold m ByteString Word8
forall (m :: * -> *). Monad m => Unfold m ByteString Word8
reader