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

module Streamly.External.ByteString
  ( toArray
  , fromArray

  , read
  , writeN
  , write
  )
where

import Control.Monad.IO.Class (MonadIO)
import Data.Word (Word8)
import GHC.ForeignPtr (ForeignPtr(..))
import GHC.Ptr (Ptr(..), minusPtr, nullPtr, 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 Streamly.Internal.Data.Array.Foreign.Mut.Type
    (ArrayContents, arrayToFptrContents, fptrToArrayContents, nilArrayContents)

import qualified Streamly.Data.Array.Foreign as A

import Prelude hiding (read)

-- | Helper function that creates a ForeignPtr
makeForeignPtr :: ArrayContents -> Ptr a -> ForeignPtr a
makeForeignPtr :: ArrayContents -> Ptr a -> ForeignPtr a
makeForeignPtr ArrayContents
contents (Ptr Addr#
addr#) =
    Addr# -> ForeignPtrContents -> ForeignPtr a
forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr Addr#
addr# (ArrayContents -> ForeignPtrContents
arrayToFptrContents ArrayContents
contents)

-- | 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 Addr#
addr# ForeignPtrContents
_) Int
_ 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 = ArrayContents -> Ptr Word8 -> Ptr Word8 -> Array Word8
forall a. ArrayContents -> Ptr a -> Ptr a -> Array a
Array ArrayContents
nilArrayContents Ptr Word8
forall a. Ptr a
nullPtr Ptr Word8
forall a. Ptr a
nullPtr
toArray (PS (ForeignPtr Addr#
addr# ForeignPtrContents
fpcontents) Int
off Int
len) =
    ArrayContents -> Ptr Word8 -> Ptr Word8 -> Array Word8
forall a. ArrayContents -> Ptr a -> Ptr a -> Array a
Array (ForeignPtrContents -> ArrayContents
fptrToArrayContents ForeignPtrContents
fpcontents) Ptr Word8
forall a. Ptr a
startPtr Ptr Word8
forall a. Ptr a
endPtr
  where
    startPtr :: Ptr b
startPtr = Addr# -> Ptr Any
forall a. Addr# -> Ptr a
Ptr Addr#
addr# Ptr Any -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off
    endPtr :: Ptr b
endPtr = Ptr Any
forall a. Ptr a
startPtr 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
ArrayContents
arrContents :: forall a. Array a -> ArrayContents
arrStart :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aEnd :: Ptr Word8
arrStart :: Ptr Word8
arrContents :: ArrayContents
..}
    | 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 (ArrayContents -> Ptr Word8 -> ForeignPtr Word8
forall a. ArrayContents -> Ptr a -> ForeignPtr a
makeForeignPtr ArrayContents
arrContents Ptr Word8
arrStart) Int
0 Int
aLen
  where
    aLen :: Int
aLen = Ptr Word8
aEnd Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
arrStart

-- | 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