{-# 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# (unsafeCoerce# :: forall a b. a -> b
unsafeCoerce# MutableByteArray# RealWorld
marr#)

-- | Helper function that creates a ForeignPtr
{-# INLINE makeForeignPtr #-}
makeForeignPtr :: MutableByteArray -> Int -> ForeignPtr a
makeForeignPtr :: forall a. MutableByteArray -> Int -> ForeignPtr a
makeForeignPtr (MutableByteArray MutableByteArray# RealWorld
marr#) (I# Int#
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. 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 forall a. MutableByteArray -> Int -> Int -> Array a
Array (MutableByteArray# RealWorld -> MutableByteArray
MutableByteArray MutableByteArray# RealWorld
marr#) Int
off (Int
off forall a. Num a => a -> a -> a
+ Int
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
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 forall a. Eq a => a -> a -> Bool
== Int
0 = forall a. Monoid a => a
mempty
    | Bool
otherwise = ForeignPtr Word8 -> Int -> ByteString
BS (forall a. MutableByteArray -> Int -> ForeignPtr a
makeForeignPtr MutableByteArray
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