{-# LANGUAGE FlexibleContexts #-}

{-# LANGUAGE RankNTypes #-}

-- | Streaming compression and decompression using conduits.

--

-- Parts of this code were taken from zlib-enum and adapted for conduits.

module Data.Conduit.Zlib (

    -- * Conduits

    compress, decompress, gzip, ungzip,

    -- * Flushing

    compressFlush, decompressFlush,

    -- * Decompression combinators

    multiple,

    -- * Re-exported from zlib-bindings

    WindowBits (..), defaultWindowBits

) where

import Data.Streaming.Zlib

import Data.Conduit

import Data.ByteString (ByteString)

import qualified Data.ByteString as S

import Control.Monad (unless, liftM)

import Control.Monad.Trans.Class (lift, MonadTrans)

import Control.Monad.Primitive (PrimMonad, unsafePrimToPrim)

import Control.Monad.Base (MonadBase, liftBase)

import Control.Monad.Trans.Resource (MonadThrow, monadThrow)

import Data.Function (fix)

-- | Gzip compression with default parameters.

gzip :: (MonadThrow m, MonadBase base m, PrimMonad base) => Conduit ByteString m ByteString

gzip = compress 1 (WindowBits 31)

-- | Gzip decompression with default parameters.

ungzip :: (MonadBase base m, PrimMonad base, MonadThrow m) => Conduit ByteString m ByteString

ungzip = decompress (WindowBits 31)

unsafeLiftIO :: (MonadBase base m, PrimMonad base, MonadThrow m) => IO a -> m a

unsafeLiftIO = liftBase . unsafePrimToPrim

-- |

-- Decompress (inflate) a stream of 'ByteString's. For example:

--

-- >    sourceFile "test.z" $= decompress defaultWindowBits $$ sinkFile "test"

decompress

    :: (MonadBase base m, PrimMonad base, MonadThrow m)

    => WindowBits -- ^ Zlib parameter (see the zlib-bindings package as well as the zlib C library)

    -> Conduit ByteString m ByteString

decompress =

    helperDecompress (liftM (fmap Chunk) await) yield' leftover

  where

    yield' Flush = return ()

    yield' (Chunk bs) = yield bs

-- | Same as 'decompress', but allows you to explicitly flush the stream.

decompressFlush

    :: (MonadBase base m, PrimMonad base, MonadThrow m)

    => WindowBits -- ^ Zlib parameter (see the zlib-bindings package as well as the zlib C library)

    -> Conduit (Flush ByteString) m (Flush ByteString)

decompressFlush = helperDecompress await yield (leftover . Chunk)

helperDecompress :: (Monad (t m), MonadBase base m, PrimMonad base, MonadThrow m, MonadTrans t)

                 => t m (Maybe (Flush ByteString))

                 -> (Flush ByteString -> t m ())

                 -> (ByteString -> t m ())

                 -> WindowBits

                 -> t m ()

helperDecompress await' yield' leftover' config = do

    -- Initialize the stateful inflater, which will be used below

    -- This inflater is never exposed outside of this function

    inf <- lift $ unsafeLiftIO $ initInflate config

    -- Some helper functions used by the main feeder loop below

    let -- Flush any remaining inflated bytes downstream

        flush = do

            chunk <- lift $ unsafeLiftIO $ flushInflate inf

            unless (S.null chunk) $ yield' $ Chunk chunk

        -- Get any input which is unused by the inflater

        getUnused = lift $ unsafeLiftIO $ getUnusedInflate inf

        -- If there is any unused data, return it as leftovers to the stream

        unused = do

            rem' <- getUnused

            unless (S.null rem') $ leftover' rem'

    -- Main loop: feed data from upstream into the inflater

    fix $ \feeder -> do

        mnext <- await'

        case mnext of

            -- No more data is available from upstream

            Nothing -> do

                -- Flush any remaining uncompressed data

                flush

                -- Return the rest of the unconsumed data as leftovers

                unused

            -- Another chunk of compressed data arrived

            Just (Chunk x) -> do

                -- Feed the compressed data into the inflater, returning a

                -- "popper" which will return chunks of decompressed data

                popper <- lift $ unsafeLiftIO $ feedInflate inf x

                -- Loop over the popper grabbing decompressed chunks and

                -- yielding them downstream

                fix $ \pop -> do

                    mbs <- lift $ unsafeLiftIO popper

                    case mbs of

                        -- No more data from this popper

                        PRDone -> do

                            rem' <- getUnused

                            if S.null rem'

                                -- No data was unused by the inflater, so let's

                                -- fill it up again and get more data out of it

                                then feeder

                                -- In this case, there is some unconsumed data,

                                -- meaning the compressed stream is complete.

                                -- At this point, we need to stop feeding,

                                -- return the unconsumed data as leftovers, and

                                -- flush any remaining content (which should be

                                -- nothing)

                                else do

                                    flush

                                    leftover' rem'

                        -- Another chunk available, yield it downstream and

                        -- loop again

                        PRNext bs -> do

                            yield' (Chunk bs)

                            pop

                        -- An error occurred inside zlib, throw it

                        PRError e -> lift $ monadThrow e

            -- We've been asked to flush the stream

            Just Flush -> do

                -- Get any uncompressed data waiting for us

                flush

                -- Put a Flush in the stream

                yield' Flush

                -- Feed in more data

                feeder

-- |

-- Compress (deflate) a stream of 'ByteString's. The 'WindowBits' also control

-- the format (zlib vs. gzip).

compress

    :: (MonadBase base m, PrimMonad base, MonadThrow m)

    => Int         -- ^ Compression level

    -> WindowBits  -- ^ Zlib parameter (see the zlib-bindings package as well as the zlib C library)

    -> Conduit ByteString m ByteString

compress =

    helperCompress (liftM (fmap Chunk) await) yield'

  where

    yield' Flush = return ()

    yield' (Chunk bs) = yield bs

-- | Same as 'compress', but allows you to explicitly flush the stream.

compressFlush

    :: (MonadBase base m, PrimMonad base, MonadThrow m)

    => Int         -- ^ Compression level

    -> WindowBits  -- ^ Zlib parameter (see the zlib-bindings package as well as the zlib C library)

    -> Conduit (Flush ByteString) m (Flush ByteString)

compressFlush = helperCompress await yield

helperCompress :: (Monad (t m), MonadBase base m, PrimMonad base, MonadThrow m, MonadTrans t)

               => t m (Maybe (Flush ByteString))

               -> (Flush ByteString -> t m ())

               -> Int

               -> WindowBits

               -> t m ()

helperCompress await' yield' level config =

    await' >>= maybe (return ()) start

  where

    start input = do

        def <- lift $ unsafeLiftIO $ initDeflate level config

        push def input

    continue def = await' >>= maybe (close def) (push def)

    goPopper popper = do

        mbs <- lift $ unsafeLiftIO popper

        case mbs of

            PRDone -> return ()

            PRNext bs -> yield' (Chunk bs) >> goPopper popper

            PRError e -> lift $ monadThrow e

    push def (Chunk x) = do

        popper <- lift $ unsafeLiftIO $ feedDeflate def x

        goPopper popper

        continue def

    push def Flush = do

        mchunk <- lift $ unsafeLiftIO $ flushDeflate def

        case mchunk of

            PRDone -> return ()

            PRNext x -> yield' $ Chunk x

            PRError e -> lift $ monadThrow e

        yield' Flush

        continue def

    close def = do

        mchunk <- lift $ unsafeLiftIO $ finishDeflate def

        case mchunk of

            PRDone -> return ()

            PRNext chunk -> yield' (Chunk chunk) >> close def

            PRError e -> lift $ monadThrow e

-- | The standard 'decompress' and 'ungzip' functions will only decompress a

-- single compressed entity from the stream. This combinator will exhaust the

-- stream completely of all individual compressed entities. This is useful for

-- cases where you have a concatenated archive, e.g. @cat file1.gz file2.gz >

-- combined.gz@.

--

-- Usage:

--

-- > sourceFile "combined.gz" $$ multiple ungzip =$ consume

--

-- This combinator will not fail on an empty stream. If you want to ensure that

-- at least one compressed entity in the stream exists, consider a usage such

-- as:

--

-- > sourceFile "combined.gz" $$ (ungzip >> multiple ungzip) =$ consume

--

-- @since 1.1.10

multiple :: Monad m

         => Conduit ByteString m a

         -> Conduit ByteString m a

multiple inner =

    loop

  where

    loop = do

        mbs <- await

        case mbs of

            Nothing -> return ()

            Just bs

                | S.null bs -> loop

                | otherwise -> do

                    leftover bs

                    inner

                    loop