{-# 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_streamly_core(0,2,0)
import Streamly.Internal.Data.Array (Array(..))
import Streamly.Internal.Data.MutByteArray (MutByteArray(..))
import qualified Streamly.Internal.Data.MutByteArray as MutBA (nil)
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 (nil)
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

{-# INLINE mutableByteArrayContents# #-}
mutableByteArrayContents# :: MutableByteArray# RealWorld -> Addr#
mutableByteArrayContents# :: MutableByteArray# RealWorld -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
marr# = ByteArray# -> Addr#
byteArrayContents# (unsafeCoerce# :: 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#) =
    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 (BS (ForeignPtr Addr#
addr# ForeignPtrContents
_) Int
_)
    | forall a. Addr# -> Ptr a
Ptr Addr#
addr# forall a. Eq a => a -> a -> Bool
== forall a. Ptr a
nullPtr = forall a. MutByteArray -> Int -> Int -> Array a
Array MutByteArray
MutBA.nil Int
0 Int
0
toArray (BS (ForeignPtr Addr#
addr# (PlainPtr MutableByteArray# RealWorld
marr#)) Int
len) =
    let off :: Int
off = Int# -> Int
I# (Addr#
addr# Addr# -> Addr# -> Int#
`minusAddr#` MutableByteArray# RealWorld -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
marr#)
     in forall a. MutByteArray -> Int -> Int -> Array a
Array (MUT_BYTE_ARRAY marr#) off (off + len)
toArray (BS ForeignPtr Word8
fptr Int
len) =
    forall a. IO a -> a
unsafeInlineIO
        forall a b. (a -> b) -> a -> b
$ forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fptr forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) b c a.
Monad m =>
Fold m b c -> Unfold m a b -> a -> m c
Unfold.fold (forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (Array a)
Array.writeN Int
len) Unfold IO (Ptr Word8) Word8
generator

    where

    generator :: Unfold IO (Ptr Word8) Word8
generator =
        forall (m :: * -> *) a b.
Applicative m =>
(a -> m (Step a b)) -> Unfold m a b
Unfold.mkUnfoldrM
            (\Ptr Word8
ptr -> forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. a -> s -> Step s a
StreamD.Yield (Ptr Word8
ptr forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
1) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Storable a => Ptr a -> IO a
peek Ptr Word8
ptr)

-- | Convert an array of 'Word8' to a 'ByteString'. This function unwraps the
-- 'Array' and wraps it with 'ByteString' constructors and hence the operation
-- is performed in constant time.
{-# INLINE fromArray #-}
fromArray :: Array Word8 -> ByteString
fromArray :: Array Word8 -> ByteString
fromArray (Array {Int
MutByteArray
arrContents :: forall a. Array a -> MutByteArray
arrStart :: forall a. Array a -> Int
arrEnd :: forall a. Array a -> Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutByteArray
..})
    | Int
aLen forall a. Eq a => a -> a -> Bool
== Int
0 = forall a. Monoid a => a
mempty
    | Bool
otherwise = ForeignPtr Word8 -> Int -> ByteString
BS (forall a. MutByteArray -> Int -> ForeignPtr a
makeForeignPtr MutByteArray
arrContents Int
arrStart) Int
aLen

    where

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

-- | 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 = forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap ByteString -> Array Word8
toArray 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
(MonadIO m, Unbox a) =>
Int -> Fold m a (Array a)
Array.writeN 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. (MonadIO m, Unbox a) => Fold m a (Array a)
Array.write

--------------------------------------------------------------------------------
-- 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 = forall (m :: * -> *). Monad m => Unfold m ByteString Word8
reader