    -- | Interface to @zlib@ and @gzip@ compression for 'Bytestring's and 'Builder's

    {-# LANGUAGE GeneralizedNewtypeDeriving #-}

    module System.IO.Streams.Zlib
     ( -- * ByteString decompression
       gunzip
     , decompress
       -- * ByteString compression
     , gzip
     , compress
       -- * Builder compression
     , gzipBuilder
     , compressBuilder
       -- * Compression level
     , CompressionLevel(..)
     , defaultCompressionLevel
     ) where

    ------------------------------------------------------------------------------
    import           Data.ByteString                  (ByteString)
    import qualified Data.ByteString                  as S
    import           Data.IORef                       (newIORef, readIORef, writeIORef)
    import           Prelude                          hiding (read)
    ------------------------------------------------------------------------------
    import           Codec.Zlib                       (Deflate, Inflate, Popper, WindowBits (..), feedDeflate, feedInflate, finishDeflate, finishInflate, flushDeflate, flushInflate, initDeflate, initInflate)
    import           Data.ByteString.Builder          (Builder, byteString)
    import           Data.ByteString.Builder.Extra    (defaultChunkSize, flush)
    import           Data.ByteString.Builder.Internal (newBuffer)
    ------------------------------------------------------------------------------
    import           System.IO.Streams.Builder        (unsafeBuilderStream)
    import           System.IO.Streams.Internal       (InputStream, OutputStream, makeInputStream, makeOutputStream, read, write)

    ------------------------------------------------------------------------------
    gzipBits :: WindowBits
    gzipBits :: WindowBits
    gzipBits = Int -> WindowBits
    WindowBits Int
    31

    ------------------------------------------------------------------------------
    compressBits :: WindowBits
    compressBits :: WindowBits
    compressBits = Int -> WindowBits
    WindowBits Int
    15

    ------------------------------------------------------------------------------
    -- | Decompress an 'InputStream' of strict 'ByteString's from the @gzip@ format
    gunzip :: InputStream ByteString -> IO (InputStream ByteString)
    gunzip :: InputStream ByteString -> IO (InputStream ByteString)
    gunzip InputStream ByteString
    input = WindowBits -> IO Inflate
    initInflate WindowBits
    gzipBits IO Inflate
    -> (Inflate -> IO (InputStream ByteString))
    -> IO (InputStream ByteString)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= InputStream ByteString -> Inflate -> IO (InputStream ByteString)
    inflate InputStream ByteString
    input

    ------------------------------------------------------------------------------
    -- | Decompress an 'InputStream' of strict 'ByteString's from the @zlib@ format
    decompress :: InputStream ByteString -> IO (InputStream ByteString)
    decompress :: InputStream ByteString -> IO (InputStream ByteString)
    decompress InputStream ByteString
    input = WindowBits -> IO Inflate
    initInflate WindowBits
    compressBits IO Inflate
    -> (Inflate -> IO (InputStream ByteString))
    -> IO (InputStream ByteString)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= InputStream ByteString -> Inflate -> IO (InputStream ByteString)
    inflate InputStream ByteString
    input

    ------------------------------------------------------------------------------
    -- Note: bytes pushed back to this input stream are not propagated back to the
    -- source InputStream.
    data IS = Input
            | Popper Popper
            | Done

    inflate :: InputStream ByteString -> Inflate -> IO (InputStream ByteString)
    inflate :: InputStream ByteString -> Inflate -> IO (InputStream ByteString)
    inflate InputStream ByteString
    input Inflate
    state = do
        ref <- IS -> IO (IORef IS)
    forall a. a -> IO (IORef a)
    newIORef IS
    Input
        makeInputStream $ stream ref

      where
        stream :: IORef IS -> IO (Maybe ByteString)
    stream IORef IS
    ref = IO (Maybe ByteString)
    go
          where
            go :: IO (Maybe ByteString)
    go = IORef IS -> IO IS
    forall a. IORef a -> IO a
    readIORef IORef IS
    ref IO IS -> (IS -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= \IS
    st ->
                 case IS
    st of
                   IS
    Input    -> InputStream ByteString -> IO (Maybe ByteString)
    forall a. InputStream a -> IO (Maybe a)
    read InputStream ByteString
    input IO (Maybe ByteString)
    -> (Maybe ByteString -> IO (Maybe ByteString))
    -> IO (Maybe ByteString)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= IO (Maybe ByteString)
    -> (ByteString -> IO (Maybe ByteString))
    -> Maybe ByteString
    -> IO (Maybe ByteString)
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe IO (Maybe ByteString)
    eof ByteString -> IO (Maybe ByteString)
    chunk
                   Popper IO (Maybe ByteString)
    p -> IO (Maybe ByteString) -> IO (Maybe ByteString)
    pop IO (Maybe ByteString)
    p
                   IS
    Done     -> Maybe ByteString -> IO (Maybe ByteString)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Maybe ByteString
    forall a. Maybe a
    Nothing

            eof :: IO (Maybe ByteString)
    eof = do
                x <- Inflate -> IO ByteString
    finishInflate Inflate
    state
                writeIORef ref Done
                if (not $ S.null x)
                  then return $! Just x
                  else return Nothing

            chunk :: ByteString -> IO (Maybe ByteString)
    chunk ByteString
    s =
                if ByteString -> Bool
    S.null ByteString
    s
                  then do
                      out <- Inflate -> IO ByteString
    flushInflate Inflate
    state
                      return $! Just out
                  else Inflate -> ByteString -> IO (IO (Maybe ByteString))
    feedInflate Inflate
    state ByteString
    s IO (IO (Maybe ByteString))
    -> (IO (Maybe ByteString) -> IO (Maybe ByteString))
    -> IO (Maybe ByteString)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= \IO (Maybe ByteString)
    popper -> do
                      IORef IS -> IS -> IO ()
    forall a. IORef a -> a -> IO ()
    writeIORef IORef IS
    ref (IS -> IO ()) -> IS -> IO ()
    forall a b. (a -> b) -> a -> b
    $ IO (Maybe ByteString) -> IS
    Popper IO (Maybe ByteString)
    popper
                      IO (Maybe ByteString) -> IO (Maybe ByteString)
    pop IO (Maybe ByteString)
    popper

            pop :: IO (Maybe ByteString) -> IO (Maybe ByteString)
    pop IO (Maybe ByteString)
    popper = IO (Maybe ByteString)
    popper IO (Maybe ByteString)
    -> (Maybe ByteString -> IO (Maybe ByteString))
    -> IO (Maybe ByteString)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= IO (Maybe ByteString)
    -> (ByteString -> IO (Maybe ByteString))
    -> Maybe ByteString
    -> IO (Maybe ByteString)
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe IO (Maybe ByteString)
    backToInput (Maybe ByteString -> IO (Maybe ByteString)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Maybe ByteString -> IO (Maybe ByteString))
    -> (ByteString -> Maybe ByteString)
    -> ByteString
    -> IO (Maybe ByteString)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ByteString -> Maybe ByteString
    forall a. a -> Maybe a
    Just)
            backToInput :: IO (Maybe ByteString)
    backToInput = IORef IS -> IS -> IO ()
    forall a. IORef a -> a -> IO ()
    writeIORef IORef IS
    ref IS
    Input IO () -> IO (Maybe ByteString) -> IO (Maybe ByteString)
    forall a b. IO a -> IO b -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> InputStream ByteString -> IO (Maybe ByteString)
    forall a. InputStream a -> IO (Maybe a)
    read InputStream ByteString
    input IO (Maybe ByteString)
    -> (Maybe ByteString -> IO (Maybe ByteString))
    -> IO (Maybe ByteString)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= IO (Maybe ByteString)
    -> (ByteString -> IO (Maybe ByteString))
    -> Maybe ByteString
    -> IO (Maybe ByteString)
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe IO (Maybe ByteString)
    eof ByteString -> IO (Maybe ByteString)
    chunk

    ------------------------------------------------------------------------------
    deflateBuilder :: OutputStream Builder
                   -> Deflate
                   -> IO (OutputStream Builder)
    deflateBuilder :: OutputStream Builder -> Deflate -> IO (OutputStream Builder)
    deflateBuilder OutputStream Builder
    stream Deflate
    state = do
        zippedStr <- (Maybe ByteString -> IO ()) -> IO (OutputStream ByteString)
    forall a. (Maybe a -> IO ()) -> IO (OutputStream a)
    makeOutputStream Maybe ByteString -> IO ()
    bytestringStream IO (OutputStream ByteString)
    -> (OutputStream ByteString -> IO (OutputStream ByteString))
    -> IO (OutputStream ByteString)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>=
                     \OutputStream ByteString
    x -> OutputStream ByteString -> Deflate -> IO (OutputStream ByteString)
    deflate OutputStream ByteString
    x Deflate
    state

        -- we can use unsafeBuilderStream here because zlib is going to consume the
        -- stream
        unsafeBuilderStream (newBuffer defaultChunkSize) zippedStr

      where
        bytestringStream :: Maybe ByteString -> IO ()
    bytestringStream Maybe ByteString
    x = Maybe Builder -> OutputStream Builder -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    write ((ByteString -> Builder) -> Maybe ByteString -> Maybe Builder
    forall a b. (a -> b) -> Maybe a -> Maybe b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ByteString -> Builder
    cvt Maybe ByteString
    x) OutputStream Builder
    stream

        cvt :: ByteString -> Builder
    cvt ByteString
    s | ByteString -> Bool
    S.null ByteString
    s  = Builder
    flush
              | Bool
    otherwise = ByteString -> Builder
    byteString ByteString
    s

    ------------------------------------------------------------------------------
    -- | Convert an 'OutputStream' that consumes compressed 'Builder's into an
    -- 'OutputStream' that consumes uncompressed 'Builder's in the @gzip@ format
    gzipBuilder :: CompressionLevel
                -> OutputStream Builder
                -> IO (OutputStream Builder)
    gzipBuilder :: CompressionLevel
    -> OutputStream Builder -> IO (OutputStream Builder)
    gzipBuilder CompressionLevel
    level OutputStream Builder
    output =
        Int -> WindowBits -> IO Deflate
    initDeflate (CompressionLevel -> Int
    clamp CompressionLevel
    level) WindowBits
    gzipBits IO Deflate
    -> (Deflate -> IO (OutputStream Builder))
    -> IO (OutputStream Builder)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= OutputStream Builder -> Deflate -> IO (OutputStream Builder)
    deflateBuilder OutputStream Builder
    output

    ------------------------------------------------------------------------------
    -- | Convert an 'OutputStream' that consumes compressed 'Builder's into an
    -- 'OutputStream' that consumes uncompressed 'Builder's in the @zlib@ format
    compressBuilder :: CompressionLevel
                    -> OutputStream Builder
                    -> IO (OutputStream Builder)
    compressBuilder :: CompressionLevel
    -> OutputStream Builder -> IO (OutputStream Builder)
    compressBuilder CompressionLevel
    level OutputStream Builder
    output =
        Int -> WindowBits -> IO Deflate
    initDeflate (CompressionLevel -> Int
    clamp CompressionLevel
    level) WindowBits
    compressBits IO Deflate
    -> (Deflate -> IO (OutputStream Builder))
    -> IO (OutputStream Builder)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= OutputStream Builder -> Deflate -> IO (OutputStream Builder)
    deflateBuilder OutputStream Builder
    output

    ------------------------------------------------------------------------------
    deflate :: OutputStream ByteString
            -> Deflate
            -> IO (OutputStream ByteString)
    deflate :: OutputStream ByteString -> Deflate -> IO (OutputStream ByteString)
    deflate OutputStream ByteString
    output Deflate
    state = (Maybe ByteString -> IO ()) -> IO (OutputStream ByteString)
    forall a. (Maybe a -> IO ()) -> IO (OutputStream a)
    makeOutputStream Maybe ByteString -> IO ()
    stream
      where
        stream :: Maybe ByteString -> IO ()
    stream Maybe ByteString
    Nothing = IO (Maybe ByteString) -> IO ()
    popAll (Deflate -> IO (Maybe ByteString)
    finishDeflate Deflate
    state) IO () -> IO () -> IO ()
    forall a b. IO a -> IO b -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Maybe ByteString -> OutputStream ByteString -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    write Maybe ByteString
    forall a. Maybe a
    Nothing OutputStream ByteString
    output

        stream (Just ByteString
    s) = do
            -- Empty string means flush
            if ByteString -> Bool
    S.null ByteString
    s
              then do
                  IO (Maybe ByteString) -> IO ()
    popAll (Deflate -> IO (Maybe ByteString)
    flushDeflate Deflate
    state)
                  Maybe ByteString -> OutputStream ByteString -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    write (ByteString -> Maybe ByteString
    forall a. a -> Maybe a
    Just ByteString
    S.empty) OutputStream ByteString
    output

              else Deflate -> ByteString -> IO (IO (Maybe ByteString))
    feedDeflate Deflate
    state ByteString
    s IO (IO (Maybe ByteString))
    -> (IO (Maybe ByteString) -> IO ()) -> IO ()
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= IO (Maybe ByteString) -> IO ()
    popAll

        popAll :: IO (Maybe ByteString) -> IO ()
    popAll IO (Maybe ByteString)
    popper = IO ()
    go
          where
            go :: IO ()
    go = IO (Maybe ByteString)
    popper IO (Maybe ByteString) -> (Maybe ByteString -> IO ()) -> IO ()
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= IO () -> (ByteString -> IO ()) -> Maybe ByteString -> IO ()
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe (() -> IO ()
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (() -> IO ()) -> () -> IO ()
    forall a b. (a -> b) -> a -> b
    $! ()) (\ByteString
    s -> Maybe ByteString -> OutputStream ByteString -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    write (ByteString -> Maybe ByteString
    forall a. a -> Maybe a
    Just ByteString
    s) OutputStream ByteString
    output IO () -> IO () -> IO ()
    forall a b. IO a -> IO b -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> IO ()
    go)

    ------------------------------------------------------------------------------
    -- | Parameter that defines the tradeoff between speed and compression ratio
    newtype CompressionLevel = CompressionLevel Int
      deriving (ReadPrec [CompressionLevel]
    ReadPrec CompressionLevel
    Int -> ReadS CompressionLevel
    ReadS [CompressionLevel]
    (Int -> ReadS CompressionLevel)
    -> ReadS [CompressionLevel]
    -> ReadPrec CompressionLevel
    -> ReadPrec [CompressionLevel]
    -> Read CompressionLevel
    forall a.
    (Int -> ReadS a)
    -> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
    $creadsPrec :: Int -> ReadS CompressionLevel
    readsPrec :: Int -> ReadS CompressionLevel
    $creadList :: ReadS [CompressionLevel]
    readList :: ReadS [CompressionLevel]
    $creadPrec :: ReadPrec CompressionLevel
    readPrec :: ReadPrec CompressionLevel
    $creadListPrec :: ReadPrec [CompressionLevel]
    readListPrec :: ReadPrec [CompressionLevel]
    Read, CompressionLevel -> CompressionLevel -> Bool
    (CompressionLevel -> CompressionLevel -> Bool)
    -> (CompressionLevel -> CompressionLevel -> Bool)
    -> Eq CompressionLevel
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: CompressionLevel -> CompressionLevel -> Bool
    == :: CompressionLevel -> CompressionLevel -> Bool
    $c/= :: CompressionLevel -> CompressionLevel -> Bool
    /= :: CompressionLevel -> CompressionLevel -> Bool
    Eq, Int -> CompressionLevel -> ShowS
    [CompressionLevel] -> ShowS
    CompressionLevel -> String
    (Int -> CompressionLevel -> ShowS)
    -> (CompressionLevel -> String)
    -> ([CompressionLevel] -> ShowS)
    -> Show CompressionLevel
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> CompressionLevel -> ShowS
    showsPrec :: Int -> CompressionLevel -> ShowS
    $cshow :: CompressionLevel -> String
    show :: CompressionLevel -> String
    $cshowList :: [CompressionLevel] -> ShowS
    showList :: [CompressionLevel] -> ShowS
    Show, Integer -> CompressionLevel
    CompressionLevel -> CompressionLevel
    CompressionLevel -> CompressionLevel -> CompressionLevel
    (CompressionLevel -> CompressionLevel -> CompressionLevel)
    -> (CompressionLevel -> CompressionLevel -> CompressionLevel)
    -> (CompressionLevel -> CompressionLevel -> CompressionLevel)
    -> (CompressionLevel -> CompressionLevel)
    -> (CompressionLevel -> CompressionLevel)
    -> (CompressionLevel -> CompressionLevel)
    -> (Integer -> CompressionLevel)
    -> Num CompressionLevel
    forall a.
    (a -> a -> a)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> (a -> a)
    -> (a -> a)
    -> (a -> a)
    -> (Integer -> a)
    -> Num a
    $c+ :: CompressionLevel -> CompressionLevel -> CompressionLevel
    + :: CompressionLevel -> CompressionLevel -> CompressionLevel
    $c- :: CompressionLevel -> CompressionLevel -> CompressionLevel
    - :: CompressionLevel -> CompressionLevel -> CompressionLevel
    $c* :: CompressionLevel -> CompressionLevel -> CompressionLevel
    * :: CompressionLevel -> CompressionLevel -> CompressionLevel
    $cnegate :: CompressionLevel -> CompressionLevel
    negate :: CompressionLevel -> CompressionLevel
    $cabs :: CompressionLevel -> CompressionLevel
    abs :: CompressionLevel -> CompressionLevel
    $csignum :: CompressionLevel -> CompressionLevel
    signum :: CompressionLevel -> CompressionLevel
    $cfromInteger :: Integer -> CompressionLevel
    fromInteger :: Integer -> CompressionLevel
    Num)

    ------------------------------------------------------------------------------
    -- | A compression level that balances speed with compression ratio
    defaultCompressionLevel :: CompressionLevel
    defaultCompressionLevel :: CompressionLevel
    defaultCompressionLevel = Int -> CompressionLevel
    CompressionLevel Int
    5

    ------------------------------------------------------------------------------
    clamp :: CompressionLevel -> Int
    clamp :: CompressionLevel -> Int
    clamp (CompressionLevel Int
    x) = Int -> Int -> Int
    forall a. Ord a => a -> a -> a
    min Int
    9 (Int -> Int -> Int
    forall a. Ord a => a -> a -> a
    max Int
    x Int
    0)

    ------------------------------------------------------------------------------
    -- | Convert an 'OutputStream' that consumes compressed 'ByteString's into an
    -- 'OutputStream' that consumes uncompressed 'ByteString's in the @gzip@ format
    gzip :: CompressionLevel
         -> OutputStream ByteString
         -> IO (OutputStream ByteString)
    gzip :: CompressionLevel
    -> OutputStream ByteString -> IO (OutputStream ByteString)
    gzip CompressionLevel
    level OutputStream ByteString
    output = Int -> WindowBits -> IO Deflate
    initDeflate (CompressionLevel -> Int
    clamp CompressionLevel
    level) WindowBits
    gzipBits IO Deflate
    -> (Deflate -> IO (OutputStream ByteString))
    -> IO (OutputStream ByteString)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= OutputStream ByteString -> Deflate -> IO (OutputStream ByteString)
    deflate OutputStream ByteString
    output

    ------------------------------------------------------------------------------
    -- | Convert an 'OutputStream' that consumes compressed 'ByteString's into an
    -- 'OutputStream' that consumes uncompressed 'ByteString's in the @zlib@ format
    compress :: CompressionLevel
             -> OutputStream ByteString
             -> IO (OutputStream ByteString)
    compress :: CompressionLevel
    -> OutputStream ByteString -> IO (OutputStream ByteString)
    compress CompressionLevel
    level OutputStream ByteString
    output = Int -> WindowBits -> IO Deflate
    initDeflate (CompressionLevel -> Int
    clamp CompressionLevel
    level) WindowBits
    compressBits IO Deflate
    -> (Deflate -> IO (OutputStream ByteString))
    -> IO (OutputStream ByteString)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>=
                            OutputStream ByteString -> Deflate -> IO (OutputStream ByteString)
    deflate OutputStream ByteString
    output