{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Streamly.External.ByteString.Lazy
  ( readChunks
  , read

  , toChunks
  , fromChunks
  , fromChunksIO
  )
where

import Data.Word (Word8)
import Streamly.Data.Unfold (many)
import Streamly.Data.Array.Foreign (Array)
import System.IO.Unsafe (unsafeInterleaveIO)

-- Internal imports
import Data.ByteString.Lazy.Internal (ByteString(..), chunk)
import Streamly.Internal.Data.Stream.StreamD.Type (Step(..))
import Streamly.Internal.Data.Unfold.Type (Unfold(..))

import qualified Streamly.External.ByteString as Strict
import qualified Streamly.Data.Array.Foreign as A
import qualified Streamly.Prelude as S

import Prelude hiding (concat, read)

-- | Unfold a lazy ByteString to a stream of 'Array' 'Words'.
{-# INLINE  readChunks #-}
readChunks :: Monad m => Unfold m ByteString (Array Word8)
readChunks :: Unfold m ByteString (Array Word8)
readChunks = (ByteString -> m (Step ByteString (Array Word8)))
-> (ByteString -> m ByteString)
-> Unfold m ByteString (Array Word8)
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold ByteString -> m (Step ByteString (Array Word8))
forall (m :: * -> *).
Monad m =>
ByteString -> m (Step ByteString (Array Word8))
step ByteString -> m ByteString
forall a. a -> m a
seed
  where
    seed :: a -> m a
seed = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
    step :: ByteString -> m (Step ByteString (Array Word8))
step (Chunk ByteString
bs ByteString
bl) = Step ByteString (Array Word8) -> m (Step ByteString (Array Word8))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step ByteString (Array Word8)
 -> m (Step ByteString (Array Word8)))
-> Step ByteString (Array Word8)
-> m (Step ByteString (Array Word8))
forall a b. (a -> b) -> a -> b
$ Array Word8 -> ByteString -> Step ByteString (Array Word8)
forall s a. a -> s -> Step s a
Yield (ByteString -> Array Word8
Strict.toArray ByteString
bs) ByteString
bl
    step ByteString
Empty = Step ByteString (Array Word8) -> m (Step ByteString (Array Word8))
forall (m :: * -> *) a. Monad m => a -> m a
return Step ByteString (Array Word8)
forall s a. Step s a
Stop

-- | Unfold a lazy ByteString to a stream of Word8
{-# INLINE read #-}
read :: Monad m => Unfold m ByteString Word8
read :: Unfold m ByteString Word8
read = Unfold m ByteString (Array Word8)
-> Unfold m (Array Word8) Word8 -> Unfold m ByteString Word8
forall (m :: * -> *) a b c.
Monad m =>
Unfold m a b -> Unfold m b c -> Unfold m a c
many Unfold m ByteString (Array Word8)
forall (m :: * -> *). Monad m => Unfold m ByteString (Array Word8)
readChunks Unfold m (Array Word8) Word8
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Unfold m (Array a) a
A.read

-- | Convert a lazy 'ByteString' to a serial stream of 'Array' 'Word8'.
{-# INLINE toChunks #-}
toChunks :: Monad m => ByteString -> S.SerialT m (Array Word8)
toChunks :: ByteString -> SerialT m (Array Word8)
toChunks = Unfold m ByteString (Array Word8)
-> ByteString -> SerialT m (Array Word8)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a b.
(IsStream t, Monad m) =>
Unfold m a b -> a -> t m b
S.unfold Unfold m ByteString (Array Word8)
forall (m :: * -> *). Monad m => Unfold m ByteString (Array Word8)
readChunks

{-
newtype LazyIO a = LazyIO { runLazy :: IO a } deriving (Functor, Applicative)

liftToLazy :: IO a -> LazyIO a
liftToLazy = LazyIO

instance Monad LazyIO where
    return = pure
    LazyIO a >>= f = LazyIO (unsafeInterleaveIO a >>= unsafeInterleaveIO . runLazy . f)
-}

-- | Convert a serial stream of 'Array' 'Word8' to a lazy 'ByteString'.
--
-- IMPORTANT NOTE: This function is lazy only for lazy monads
-- (e.g. Identity). For strict monads (e.g. /IO/) it consumes the entire input
-- before generating the output. For /IO/ monad please use fromChunksIO
-- instead.
--
-- For strict monads like /IO/ you could create a newtype wrapper to make the
-- monad bind operation lazy and lift the stream to that type using hoist, then
-- you can use this function to generate the bytestring lazily. For example you
-- can wrap the /IO/ type to make the bind lazy like this:
--
-- @
-- newtype LazyIO a = LazyIO { runLazy :: IO a } deriving (Functor, Applicative)
--
-- liftToLazy :: IO a -> LazyIO a
-- liftToLazy = LazyIO
--
-- instance Monad LazyIO where
--   return = pure
--   LazyIO a >>= f = LazyIO (unsafeInterleaveIO a >>= unsafeInterleaveIO . runLazy . f)
-- @
--
-- /fromChunks/ can then be used as,
-- @
-- {-# INLINE fromChunksIO #-}
-- fromChunksIO :: SerialT IO (Array Word8) -> IO ByteString
-- fromChunksIO str = runLazy (fromChunks (S.hoist liftToLazy str))
-- @
{-# INLINE fromChunks #-}
fromChunks :: Monad m => S.SerialT m (Array Word8) -> m ByteString
fromChunks :: SerialT m (Array Word8) -> m ByteString
fromChunks = (ByteString -> ByteString -> ByteString)
-> ByteString -> SerialT m ByteString -> m ByteString
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> SerialT m a -> m b
S.foldr ByteString -> ByteString -> ByteString
chunk ByteString
Empty (SerialT m ByteString -> m ByteString)
-> (SerialT m (Array Word8) -> SerialT m ByteString)
-> SerialT m (Array Word8)
-> m ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array Word8 -> ByteString)
-> SerialT m (Array Word8) -> SerialT m ByteString
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a b.
(IsStream t, Monad m) =>
(a -> b) -> t m a -> t m b
S.map Array Word8 -> ByteString
Strict.fromArray

-- | Convert a serial stream of 'Array' 'Word8' to a lazy 'ByteString' in the
-- /IO/ monad.
{-# INLINE fromChunksIO #-}
fromChunksIO :: S.SerialT IO (Array Word8) -> IO ByteString
fromChunksIO :: SerialT IO (Array Word8) -> IO ByteString
fromChunksIO =
-- Although the /IO/ monad is strict in nature we emulate laziness using
-- 'unsafeInterleaveIO'.
    (ByteString -> IO ByteString -> IO ByteString)
-> IO ByteString -> SerialT IO ByteString -> IO ByteString
forall (m :: * -> *) a b.
Monad m =>
(a -> m b -> m b) -> m b -> SerialT m a -> m b
S.foldrM
        (\ByteString
x IO ByteString
b -> IO ByteString -> IO ByteString
forall a. IO a -> IO a
unsafeInterleaveIO IO ByteString
b IO ByteString -> (ByteString -> IO ByteString) -> IO ByteString
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString)
-> (ByteString -> ByteString) -> ByteString -> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString -> ByteString
chunk ByteString
x)
        (ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
Empty) (SerialT IO ByteString -> IO ByteString)
-> (SerialT IO (Array Word8) -> SerialT IO ByteString)
-> SerialT IO (Array Word8)
-> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    (Array Word8 -> ByteString)
-> SerialT IO (Array Word8) -> SerialT IO ByteString
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a b.
(IsStream t, Monad m) =>
(a -> b) -> t m a -> t m b
S.map Array Word8 -> ByteString
Strict.fromArray