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

module Streamly.External.ByteString
  ( toArray
  , fromArray

  , read
  , writeN
  , write
  )
where

import Control.Monad.IO.Class (MonadIO)
import Data.Word (Word8)
#if MIN_VERSION_base(4, 10, 0)
import Foreign.ForeignPtr (plusForeignPtr)
#else
import Foreign.ForeignPtr.Compat (plusForeignPtr)
#endif
import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
import GHC.Ptr (minusPtr, plusPtr)
import Streamly.Data.Unfold (Unfold, lmap)
import Streamly.Data.Fold (Fold)

-- Internal imports
import Data.ByteString.Internal (ByteString(..))
import Streamly.Internal.Data.Array.Foreign.Type (Array(..))

import qualified Streamly.Data.Array.Foreign as A

import Prelude hiding (read)

-- | Convert a 'ByteString' to an array of 'Word8'. This function unwraps the
-- 'ByteString' and wraps it with 'Array' constructors and hence the operation
-- is performed in constant time.
{-# INLINE toArray #-}
toArray :: ByteString -> Array Word8
toArray :: ByteString -> Array Word8
toArray (PS ForeignPtr Word8
fp Int
off Int
len) = ForeignPtr Word8 -> Ptr Word8 -> Array Word8
forall a. ForeignPtr a -> Ptr a -> Array a
Array ForeignPtr Word8
forall b. ForeignPtr b
nfp Ptr Word8
forall b. Ptr b
endPtr
  where
    nfp :: ForeignPtr b
nfp = ForeignPtr Word8
fp ForeignPtr Word8 -> Int -> ForeignPtr b
forall a b. ForeignPtr a -> Int -> ForeignPtr b
`plusForeignPtr` Int
off
    endPtr :: Ptr b
endPtr = ForeignPtr Any -> Ptr Any
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Any
forall b. ForeignPtr b
nfp Ptr Any -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len

-- | 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 {Ptr Word8
ForeignPtr Word8
aStart :: forall a. Array a -> ForeignPtr a
aEnd :: forall a. Array a -> Ptr a
aEnd :: Ptr Word8
aStart :: ForeignPtr Word8
..}
    | 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 -> Int -> ByteString
PS ForeignPtr Word8
aStart Int
0 Int
aLen
  where
    aStartPtr :: Ptr Word8
aStartPtr = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
aStart
    aLen :: Int
aLen = Ptr Word8
aEnd Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
aStartPtr

-- | Unfold a strict ByteString to a stream of Word8.
{-# INLINE read #-}
read :: Monad m => Unfold m ByteString Word8
read :: Unfold m ByteString Word8
read = (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, Storable a) =>
Unfold m (Array a) a
A.read

-- | 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, Storable a) =>
Int -> Fold m a (Array a)
A.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, Storable a) =>
Fold m a (Array a)
A.write