module Codec.Compression.Zlib.Internal (
compress,
decompress,
CompressStream(..),
compressST,
compressIO,
foldCompressStream,
foldCompressStreamWithInput,
DecompressStream(..),
DecompressError(..),
decompressST,
decompressIO,
foldDecompressStream,
foldDecompressStreamWithInput,
CompressParams(..),
defaultCompressParams,
DecompressParams(..),
defaultDecompressParams,
Stream.Format(..),
Stream.gzipFormat,
Stream.zlibFormat,
Stream.rawFormat,
Stream.gzipOrZlibFormat,
Stream.CompressionLevel(..),
Stream.defaultCompression,
Stream.noCompression,
Stream.bestSpeed,
Stream.bestCompression,
Stream.compressionLevel,
Stream.Method(..),
Stream.deflateMethod,
Stream.WindowBits(..),
Stream.defaultWindowBits,
Stream.windowBits,
Stream.MemoryLevel(..),
Stream.defaultMemoryLevel,
Stream.minMemoryLevel,
Stream.maxMemoryLevel,
Stream.memoryLevel,
Stream.CompressionStrategy(..),
Stream.defaultStrategy,
Stream.filteredStrategy,
Stream.huffmanOnlyStrategy,
) where
import Prelude hiding (length)
import Control.Monad (when)
import Control.Exception (Exception, throw, assert)
import Control.Monad.ST.Lazy hiding (stToIO)
import Control.Monad.ST.Strict (stToIO)
import Data.Typeable (Typeable)
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Lazy.Internal as L
import qualified Data.ByteString as S
import qualified Data.ByteString.Internal as S
import qualified Codec.Compression.Zlib.Stream as Stream
import Codec.Compression.Zlib.Stream (Stream)
data CompressParams = CompressParams {
compressLevel :: !Stream.CompressionLevel,
compressMethod :: !Stream.Method,
compressWindowBits :: !Stream.WindowBits,
compressMemoryLevel :: !Stream.MemoryLevel,
compressStrategy :: !Stream.CompressionStrategy,
compressBufferSize :: !Int,
compressDictionary :: Maybe S.ByteString
}
data DecompressParams = DecompressParams {
decompressWindowBits :: !Stream.WindowBits,
decompressBufferSize :: !Int,
decompressDictionary :: Maybe S.ByteString
}
defaultCompressParams :: CompressParams
defaultCompressParams = CompressParams {
compressLevel = Stream.defaultCompression,
compressMethod = Stream.deflateMethod,
compressWindowBits = Stream.defaultWindowBits,
compressMemoryLevel = Stream.defaultMemoryLevel,
compressStrategy = Stream.defaultStrategy,
compressBufferSize = defaultCompressBufferSize,
compressDictionary = Nothing
}
defaultDecompressParams :: DecompressParams
defaultDecompressParams = DecompressParams {
decompressWindowBits = Stream.defaultWindowBits,
decompressBufferSize = defaultDecompressBufferSize,
decompressDictionary = Nothing
}
defaultCompressBufferSize, defaultDecompressBufferSize :: Int
defaultCompressBufferSize = 16 * 1024 L.chunkOverhead
defaultDecompressBufferSize = 32 * 1024 L.chunkOverhead
data DecompressStream m
= DecompressInputRequired (S.ByteString -> m (DecompressStream m))
| DecompressOutputAvailable S.ByteString (m (DecompressStream m))
| DecompressStreamEnd S.ByteString
| DecompressStreamError DecompressError
data DecompressError =
TruncatedInput
| DictionaryRequired
| DictionaryMismatch
| DataFormatError String
deriving (Eq, Typeable)
instance Show DecompressError where
show TruncatedInput = modprefix "premature end of compressed data stream"
show DictionaryRequired = modprefix "compressed data stream requires custom dictionary"
show DictionaryMismatch = modprefix "given dictionary does not match the expected one"
show (DataFormatError detail) = modprefix ("compressed data stream format error (" ++ detail ++ ")")
modprefix :: ShowS
modprefix = ("Codec.Compression.Zlib: " ++)
instance Exception DecompressError
foldDecompressStream :: Monad m
=> ((S.ByteString -> m a) -> m a)
-> (S.ByteString -> m a -> m a)
-> (S.ByteString -> m a)
-> (DecompressError -> m a)
-> DecompressStream m -> m a
foldDecompressStream input output end err = fold
where
fold (DecompressInputRequired next) =
input (\x -> next x >>= fold)
fold (DecompressOutputAvailable outchunk next) =
output outchunk (next >>= fold)
fold (DecompressStreamEnd inchunk) = end inchunk
fold (DecompressStreamError derr) = err derr
foldDecompressStreamWithInput :: (S.ByteString -> a -> a)
-> (L.ByteString -> a)
-> (DecompressError -> a)
-> (forall s. DecompressStream (ST s))
-> L.ByteString
-> a
foldDecompressStreamWithInput chunk end err = \s lbs ->
runST (fold s (L.toChunks lbs))
where
fold (DecompressInputRequired next) [] =
next S.empty >>= \strm -> fold strm []
fold (DecompressInputRequired next) (inchunk:inchunks) =
next inchunk >>= \s -> fold s inchunks
fold (DecompressOutputAvailable outchunk next) inchunks = do
r <- next >>= \s -> fold s inchunks
return $ chunk outchunk r
fold (DecompressStreamEnd inchunk) inchunks =
return $ end (L.fromChunks (inchunk:inchunks))
fold (DecompressStreamError derr) _ =
return $ err derr
data CompressStream m
= CompressInputRequired (S.ByteString -> m (CompressStream m))
| CompressOutputAvailable S.ByteString (m (CompressStream m))
| CompressStreamEnd
foldCompressStream :: Monad m
=> ((S.ByteString -> m a) -> m a)
-> (S.ByteString -> m a -> m a)
-> m a
-> CompressStream m -> m a
foldCompressStream input output end = fold
where
fold (CompressInputRequired next) =
input (\x -> next x >>= fold)
fold (CompressOutputAvailable outchunk next) =
output outchunk (next >>= fold)
fold CompressStreamEnd =
end
foldCompressStreamWithInput :: (S.ByteString -> a -> a)
-> a
-> (forall s. CompressStream (ST s))
-> L.ByteString
-> a
foldCompressStreamWithInput chunk end = \s lbs ->
runST (fold s (L.toChunks lbs))
where
fold (CompressInputRequired next) [] =
next S.empty >>= \strm -> fold strm []
fold (CompressInputRequired next) (inchunk:inchunks) =
next inchunk >>= \s -> fold s inchunks
fold (CompressOutputAvailable outchunk next) inchunks = do
r <- next >>= \s -> fold s inchunks
return $ chunk outchunk r
fold CompressStreamEnd _inchunks =
return end
compress :: Stream.Format -> CompressParams -> L.ByteString -> L.ByteString
compressST :: Stream.Format -> CompressParams -> CompressStream (ST s)
compressIO :: Stream.Format -> CompressParams -> CompressStream IO
compress format params = compressStreamToLBS (compressStream format params)
compressST format params = compressStreamToST (compressStream format params)
compressIO format params = compressStreamToIO (compressStream format params)
compressStream :: Stream.Format -> CompressParams -> CompressStream Stream
compressStream format (CompressParams compLevel method bits memLevel
strategy initChunkSize mdict) =
CompressInputRequired $ \chunk -> do
Stream.deflateInit format compLevel method bits memLevel strategy
setDictionary mdict
case chunk of
_ | S.null chunk ->
fillBuffers 20
S.PS inFPtr offset length -> do
Stream.pushInputBuffer inFPtr offset length
fillBuffers initChunkSize
where
fillBuffers :: Int -> Stream (CompressStream Stream)
fillBuffers outChunkSize = do
#ifdef DEBUG
Stream.consistencyCheck
#endif
inputBufferEmpty <- Stream.inputBufferEmpty
outputBufferFull <- Stream.outputBufferFull
assert (inputBufferEmpty || outputBufferFull) $ return ()
when outputBufferFull $ do
outFPtr <- Stream.unsafeLiftIO (S.mallocByteString outChunkSize)
Stream.pushOutputBuffer outFPtr 0 outChunkSize
if inputBufferEmpty
then return $ CompressInputRequired $ \chunk ->
case chunk of
_ | S.null chunk -> drainBuffers True
S.PS inFPtr offset length -> do
Stream.pushInputBuffer inFPtr offset length
drainBuffers False
else drainBuffers False
drainBuffers :: Bool -> Stream (CompressStream Stream)
drainBuffers lastChunk = do
inputBufferEmpty' <- Stream.inputBufferEmpty
outputBufferFull' <- Stream.outputBufferFull
assert(not outputBufferFull'
&& (lastChunk || not inputBufferEmpty')) $ return ()
let flush = if lastChunk then Stream.Finish else Stream.NoFlush
status <- Stream.deflate flush
case status of
Stream.Ok -> do
outputBufferFull <- Stream.outputBufferFull
if outputBufferFull
then do (outFPtr, offset, length) <- Stream.popOutputBuffer
let chunk = S.PS outFPtr offset length
return $ CompressOutputAvailable chunk $ do
fillBuffers defaultCompressBufferSize
else do fillBuffers defaultCompressBufferSize
Stream.StreamEnd -> do
inputBufferEmpty <- Stream.inputBufferEmpty
assert inputBufferEmpty $ return ()
outputBufferBytesAvailable <- Stream.outputBufferBytesAvailable
if outputBufferBytesAvailable > 0
then do (outFPtr, offset, length) <- Stream.popOutputBuffer
let chunk = S.PS outFPtr offset length
Stream.finalise
return $ CompressOutputAvailable chunk (return CompressStreamEnd)
else do Stream.finalise
return CompressStreamEnd
Stream.Error code msg -> case code of
Stream.BufferError -> fail "BufferError should be impossible!"
Stream.NeedDict _ -> fail "NeedDict is impossible!"
_ -> fail msg
setDictionary :: Maybe S.ByteString -> Stream ()
setDictionary (Just dict)
| Stream.formatSupportsDictionary format = do
status <- Stream.deflateSetDictionary dict
case status of
Stream.Ok -> return ()
Stream.Error _ msg -> fail msg
_ -> fail "error when setting deflate dictionary"
setDictionary _ = return ()
decompress :: Stream.Format -> DecompressParams -> L.ByteString -> L.ByteString
decompressST :: Stream.Format -> DecompressParams -> DecompressStream (ST s)
decompressIO :: Stream.Format -> DecompressParams -> DecompressStream IO
decompress format params = decompressStreamToLBS (decompressStream format params)
decompressST format params = decompressStreamToST (decompressStream format params)
decompressIO format params = decompressStreamToIO (decompressStream format params)
decompressStream :: Stream.Format -> DecompressParams -> DecompressStream Stream
decompressStream format (DecompressParams bits initChunkSize mdict) =
DecompressInputRequired $ \chunk -> do
Stream.inflateInit format bits
case chunk of
_ | S.null chunk ->
fillBuffers 4
S.PS inFPtr offset length -> do
Stream.pushInputBuffer inFPtr offset length
fillBuffers initChunkSize
where
fillBuffers :: Int
-> Stream (DecompressStream Stream)
fillBuffers outChunkSize = do
#ifdef DEBUG
Stream.consistencyCheck
#endif
inputBufferEmpty <- Stream.inputBufferEmpty
outputBufferFull <- Stream.outputBufferFull
assert (inputBufferEmpty || outputBufferFull) $ return ()
when outputBufferFull $ do
outFPtr <- Stream.unsafeLiftIO (S.mallocByteString outChunkSize)
Stream.pushOutputBuffer outFPtr 0 outChunkSize
if inputBufferEmpty
then return $ DecompressInputRequired $ \chunk ->
case chunk of
_ | S.null chunk -> drainBuffers True
S.PS inFPtr offset length -> do
Stream.pushInputBuffer inFPtr offset length
drainBuffers False
else drainBuffers False
drainBuffers :: Bool -> Stream (DecompressStream Stream)
drainBuffers lastChunk = do
inputBufferEmpty' <- Stream.inputBufferEmpty
outputBufferFull' <- Stream.outputBufferFull
assert(not outputBufferFull'
&& (lastChunk || not inputBufferEmpty')) $ return ()
status <- Stream.inflate Stream.NoFlush
case status of
Stream.Ok -> do
outputBufferFull <- Stream.outputBufferFull
if outputBufferFull
then do (outFPtr, offset, length) <- Stream.popOutputBuffer
let chunk = S.PS outFPtr offset length
return $ DecompressOutputAvailable chunk $ do
fillBuffers defaultDecompressBufferSize
else do fillBuffers defaultDecompressBufferSize
Stream.StreamEnd -> do
inputBufferEmpty <- Stream.inputBufferEmpty
if inputBufferEmpty
then do finish (DecompressStreamEnd S.empty)
else do (inFPtr, offset, length) <- Stream.remainingInputBuffer
let inchunk = S.PS inFPtr offset length
finish (DecompressStreamEnd inchunk)
Stream.Error code msg -> case code of
Stream.BufferError -> finish (DecompressStreamError TruncatedInput)
Stream.NeedDict adler -> do
err <- setDictionary adler mdict
case err of
Just streamErr -> finish streamErr
Nothing -> drainBuffers lastChunk
Stream.DataError -> finish (DecompressStreamError (DataFormatError msg))
_ -> fail msg
finish end = do
outputBufferBytesAvailable <- Stream.outputBufferBytesAvailable
if outputBufferBytesAvailable > 0
then do (outFPtr, offset, length) <- Stream.popOutputBuffer
Stream.finalise
return (DecompressOutputAvailable (S.PS outFPtr offset length) (return end))
else do Stream.finalise
return end
setDictionary :: Stream.DictionaryHash -> Maybe S.ByteString
-> Stream (Maybe (DecompressStream Stream))
setDictionary _adler Nothing =
return $ Just (DecompressStreamError DictionaryRequired)
setDictionary _adler (Just dict) = do
status <- Stream.inflateSetDictionary dict
case status of
Stream.Ok -> return Nothing
Stream.Error Stream.DataError _ ->
return $ Just (DecompressStreamError DictionaryMismatch)
_ -> fail "error when setting inflate dictionary"
compressStreamToLBS :: CompressStream Stream -> L.ByteString -> L.ByteString
compressStreamToLBS = \strm inchunks ->
runST (do zstate <- strictToLazyST $ Stream.mkState
go strm zstate inchunks)
where
go :: CompressStream Stream -> Stream.State s
-> L.ByteString -> ST s L.ByteString
go (CompressInputRequired next) zstate L.Empty = do
(strm', zstate') <- strictToLazyST $ Stream.runStream (next S.empty) zstate
go strm' zstate' L.Empty
go (CompressInputRequired next) zstate (L.Chunk inchunk inchunks') = do
(strm', zstate') <- strictToLazyST $ Stream.runStream (next inchunk) zstate
go strm' zstate' inchunks'
go (CompressOutputAvailable outchunk next) zstate inchunks = do
(strm', zstate') <- strictToLazyST $ Stream.runStream next zstate
outchunks <- go strm' zstate' inchunks
return (L.Chunk outchunk outchunks)
go CompressStreamEnd _ _ = return L.Empty
compressStreamToIO :: CompressStream Stream -> CompressStream IO
compressStreamToIO =
\(CompressInputRequired next) ->
CompressInputRequired $ \chunk -> do
zstate <- stToIO Stream.mkState
(strm', zstate') <- stToIO $ Stream.runStream (next chunk) zstate
return (go strm' zstate')
where
go :: CompressStream Stream -> Stream.State RealWorld -> CompressStream IO
go (CompressInputRequired next) zstate =
CompressInputRequired $ \chunk -> do
(strm', zstate') <- stToIO $ Stream.runStream (next chunk) zstate
return (go strm' zstate')
go (CompressOutputAvailable chunk next) zstate =
CompressOutputAvailable chunk $ do
(strm', zstate') <- stToIO $ Stream.runStream next zstate
return (go strm' zstate')
go CompressStreamEnd _ = CompressStreamEnd
compressStreamToST :: CompressStream Stream -> CompressStream (ST s)
compressStreamToST =
\(CompressInputRequired next) ->
CompressInputRequired $ \chunk -> do
zstate <- strictToLazyST $ Stream.mkState
(strm', zstate') <- strictToLazyST $ Stream.runStream (next chunk) zstate
return (go strm' zstate')
where
go :: CompressStream Stream -> Stream.State s -> CompressStream (ST s)
go (CompressInputRequired next) zstate =
CompressInputRequired $ \chunk -> do
(strm', zstate') <- strictToLazyST $ Stream.runStream (next chunk) zstate
return (go strm' zstate')
go (CompressOutputAvailable chunk next) zstate =
CompressOutputAvailable chunk $ do
(strm', zstate') <- strictToLazyST $ Stream.runStream next zstate
return (go strm' zstate')
go CompressStreamEnd _ = CompressStreamEnd
decompressStreamToLBS :: DecompressStream Stream -> L.ByteString -> L.ByteString
decompressStreamToLBS = \strm inchunks ->
runST (do zstate <- strictToLazyST Stream.mkState
go strm zstate inchunks)
where
go :: DecompressStream Stream -> Stream.State s
-> L.ByteString -> ST s L.ByteString
go (DecompressInputRequired next) zstate L.Empty = do
(strm', zstate') <- strictToLazyST $ Stream.runStream (next S.empty) zstate
go strm' zstate' L.Empty
go (DecompressInputRequired next) zstate (L.Chunk inchunk inchunks') = do
(strm', zstate') <- strictToLazyST $ Stream.runStream (next inchunk) zstate
go strm' zstate' inchunks'
go (DecompressOutputAvailable outchunk next) zstate inchunks = do
(strm', zstate') <- strictToLazyST $ Stream.runStream next zstate
outchunks <- go strm' zstate' inchunks
return (L.Chunk outchunk outchunks)
go (DecompressStreamEnd _) _ !_inchunks = return L.Empty
go (DecompressStreamError err) _ _ = throw err
decompressStreamToIO :: DecompressStream Stream -> DecompressStream IO
decompressStreamToIO =
\(DecompressInputRequired next) ->
DecompressInputRequired $ \chunk -> do
zstate <- stToIO Stream.mkState
(strm', zstate') <- stToIO $ Stream.runStream (next chunk) zstate
return (go strm' zstate')
where
go :: DecompressStream Stream -> Stream.State RealWorld -> DecompressStream IO
go (DecompressInputRequired next) zstate =
DecompressInputRequired $ \chunk -> do
(strm', zstate') <- stToIO $ Stream.runStream (next chunk) zstate
return (go strm' zstate')
go (DecompressOutputAvailable chunk next) zstate =
DecompressOutputAvailable chunk $ do
(strm', zstate') <- stToIO $ Stream.runStream next zstate
return (go strm' zstate')
go (DecompressStreamEnd chunk) _ = DecompressStreamEnd chunk
go (DecompressStreamError err) _ = DecompressStreamError err
decompressStreamToST :: DecompressStream Stream -> DecompressStream (ST s)
decompressStreamToST =
\(DecompressInputRequired next) ->
DecompressInputRequired $ \chunk -> do
zstate <- strictToLazyST Stream.mkState
(strm', zstate') <- strictToLazyST $ Stream.runStream (next chunk) zstate
return (go strm' zstate')
where
go :: DecompressStream Stream -> Stream.State s -> DecompressStream (ST s)
go (DecompressInputRequired next) zstate =
DecompressInputRequired $ \chunk -> do
(strm', zstate') <- strictToLazyST $ Stream.runStream (next chunk) zstate
return (go strm' zstate')
go (DecompressOutputAvailable chunk next) zstate =
DecompressOutputAvailable chunk $ do
(strm', zstate') <- strictToLazyST $ Stream.runStream next zstate
return (go strm' zstate')
go (DecompressStreamEnd chunk) _ = DecompressStreamEnd chunk
go (DecompressStreamError err) _ = DecompressStreamError err