{-# 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.Data.Array.Type (Array(..))
import Streamly.Internal.Data.Unboxed (MutableByteArray(..))
import Streamly.Internal.System.IO (unsafeInlineIO)

import qualified Streamly.Data.Array as Array
import qualified Streamly.Internal.Data.Unfold as Unfold (fold, mkUnfoldrM)
import qualified Streamly.Internal.Data.Unboxed as Unboxed (nil)
import qualified Streamly.Internal.Data.Stream.StreamD as StreamD (Step(Yield))

import Prelude hiding (read)

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

-- | Helper function that creates a ForeignPtr
{-# INLINE makeForeignPtr #-}
makeForeignPtr :: MutableByteArray -> Int -> ForeignPtr a
makeForeignPtr :: MutableByteArray -> Int -> ForeignPtr a
makeForeignPtr (MutableByteArray MutableByteArray# RealWorld
marr#) (I# Int#
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 (BS (ForeignPtr Addr#
addr# ForeignPtrContents
_) Int
_)
    | 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 = MutableByteArray -> Int -> Int -> Array Word8
forall a. MutableByteArray -> Int -> Int -> Array a
Array MutableByteArray
Unboxed.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 MutableByteArray -> Int -> Int -> Array Word8
forall a. MutableByteArray -> Int -> Int -> Array a
Array (MutableByteArray# RealWorld -> MutableByteArray
MutableByteArray MutableByteArray# RealWorld
marr#) Int
off (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)
toArray (BS ForeignPtr Word8
fptr Int
len) =
    IO (Array Word8) -> Array Word8
forall a. IO a -> a
unsafeInlineIO
        (IO (Array Word8) -> Array Word8)
-> IO (Array Word8) -> Array Word8
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8
-> (Ptr Word8 -> IO (Array Word8)) -> IO (Array Word8)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fptr ((Ptr Word8 -> IO (Array Word8)) -> IO (Array Word8))
-> (Ptr Word8 -> IO (Array Word8)) -> IO (Array Word8)
forall a b. (a -> b) -> a -> b
$ Fold IO Word8 (Array Word8)
-> Unfold IO (Ptr Word8) Word8 -> Ptr Word8 -> IO (Array Word8)
forall (m :: * -> *) b c a.
Monad m =>
Fold m b c -> Unfold m a b -> a -> m c
Unfold.fold (Int -> Fold IO Word8 (Array Word8)
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 =
        (Ptr Word8 -> IO (Step (Ptr Word8) Word8))
-> Unfold IO (Ptr Word8) Word8
forall (m :: * -> *) a b.
Applicative m =>
(a -> m (Step a b)) -> Unfold m a b
Unfold.mkUnfoldrM
            (\Ptr Word8
ptr -> (Word8 -> Ptr Word8 -> Step (Ptr Word8) Word8)
-> Ptr Word8 -> Word8 -> Step (Ptr Word8) Word8
forall a b c. (a -> b -> c) -> b -> a -> c
flip Word8 -> Ptr Word8 -> Step (Ptr Word8) Word8
forall s a. a -> s -> Step s a
StreamD.Yield (Ptr Word8
ptr Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
1) (Word8 -> Step (Ptr Word8) Word8)
-> IO Word8 -> IO (Step (Ptr Word8) Word8)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr Word8 -> IO Word8
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
MutableByteArray
arrContents :: forall a. Array a -> MutableByteArray
arrStart :: forall a. Array a -> Int
arrEnd :: forall a. Array a -> Int
arrEnd :: Int
arrStart :: Int
arrContents :: MutableByteArray
..}
    | Int
aLen Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = ByteString
forall a. Monoid a => a
mempty
    | Bool
otherwise = ForeignPtr Word8 -> Int -> ByteString
BS (MutableByteArray -> Int -> ForeignPtr Word8
forall a. MutableByteArray -> Int -> ForeignPtr a
makeForeignPtr MutableByteArray
arrContents Int
arrStart) Int
aLen
  where
    aLen :: Int
aLen = Int
arrEnd Int -> Int -> Int
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 :: 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 :: 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 :: * -> *) 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 :: 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 :: * -> *) 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 :: Unfold m ByteString Word8
read = Unfold m ByteString Word8
forall (m :: * -> *). Monad m => Unfold m ByteString Word8
reader