-- | 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 (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 (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
    IORef IS
ref <- IS -> IO (IORef IS)
forall a. a -> IO (IORef a)
newIORef IS
Input
    IO (Maybe ByteString) -> IO (InputStream ByteString)
forall a. IO (Maybe a) -> IO (InputStream a)
makeInputStream (IO (Maybe ByteString) -> IO (InputStream ByteString))
-> IO (Maybe ByteString) -> IO (InputStream ByteString)
forall a b. (a -> b) -> a -> b
$ IORef IS -> IO (Maybe ByteString)
stream IORef IS
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 (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 (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 (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing

        eof :: IO (Maybe ByteString)
eof = do
            ByteString
x <- Inflate -> IO ByteString
finishInflate Inflate
state
            IORef IS -> IS -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef IS
ref IS
Done
            if (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ByteString -> Bool
S.null ByteString
x)
              then Maybe ByteString -> IO (Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ByteString -> IO (Maybe ByteString))
-> Maybe ByteString -> IO (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$! ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
x
              else Maybe ByteString -> IO (Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing

        chunk :: ByteString -> IO (Maybe ByteString)
chunk ByteString
s =
            if ByteString -> Bool
S.null ByteString
s
              then do
                  ByteString
out <- Inflate -> IO ByteString
flushInflate Inflate
state
                  Maybe ByteString -> IO (Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ByteString -> IO (Maybe ByteString))
-> Maybe ByteString -> IO (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$! ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
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 (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 (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 (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 (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 (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
    OutputStream ByteString
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 (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
    IO Buffer -> OutputStream ByteString -> IO (OutputStream Builder)
unsafeBuilderStream (Int -> IO Buffer
newBuffer Int
defaultChunkSize) OutputStream ByteString
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 (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 (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 (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 (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 (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 (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 (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 (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
readListPrec :: ReadPrec [CompressionLevel]
$creadListPrec :: ReadPrec [CompressionLevel]
readPrec :: ReadPrec CompressionLevel
$creadPrec :: ReadPrec CompressionLevel
readList :: ReadS [CompressionLevel]
$creadList :: ReadS [CompressionLevel]
readsPrec :: Int -> ReadS CompressionLevel
$creadsPrec :: Int -> ReadS CompressionLevel
Read, CompressionLevel -> CompressionLevel -> Bool
(CompressionLevel -> CompressionLevel -> Bool)
-> (CompressionLevel -> CompressionLevel -> Bool)
-> Eq CompressionLevel
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CompressionLevel -> CompressionLevel -> Bool
$c/= :: CompressionLevel -> CompressionLevel -> Bool
== :: CompressionLevel -> CompressionLevel -> Bool
$c== :: 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
showList :: [CompressionLevel] -> ShowS
$cshowList :: [CompressionLevel] -> ShowS
show :: CompressionLevel -> String
$cshow :: CompressionLevel -> String
showsPrec :: Int -> CompressionLevel -> ShowS
$cshowsPrec :: Int -> 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
fromInteger :: Integer -> CompressionLevel
$cfromInteger :: Integer -> CompressionLevel
signum :: CompressionLevel -> CompressionLevel
$csignum :: CompressionLevel -> CompressionLevel
abs :: CompressionLevel -> CompressionLevel
$cabs :: CompressionLevel -> CompressionLevel
negate :: CompressionLevel -> CompressionLevel
$cnegate :: CompressionLevel -> CompressionLevel
* :: CompressionLevel -> CompressionLevel -> CompressionLevel
$c* :: CompressionLevel -> CompressionLevel -> CompressionLevel
- :: CompressionLevel -> CompressionLevel -> CompressionLevel
$c- :: CompressionLevel -> CompressionLevel -> CompressionLevel
+ :: CompressionLevel -> CompressionLevel -> CompressionLevel
$c+ :: CompressionLevel -> CompressionLevel -> 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 (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 (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
                        OutputStream ByteString -> Deflate -> IO (OutputStream ByteString)
deflate OutputStream ByteString
output