diff --git a/Codec/Compression/Zlib/Internal.hs b/Codec/Compression/Zlib/Internal.hs --- a/Codec/Compression/Zlib/Internal.hs +++ b/Codec/Compression/Zlib/Internal.hs @@ -84,7 +84,8 @@ compressWindowBits :: !Stream.WindowBits, compressMemoryLevel :: !Stream.MemoryLevel, compressStrategy :: !Stream.CompressionStrategy, - compressBufferSize :: !Int + compressBufferSize :: !Int, + compressDictionary :: Maybe S.ByteString } -- | The full set of parameters for decompression. The defaults are @@ -108,7 +109,8 @@ -- data DecompressParams = DecompressParams { decompressWindowBits :: !Stream.WindowBits, - decompressBufferSize :: !Int + decompressBufferSize :: !Int, + decompressDictionary :: Maybe S.ByteString } -- | The default set of parameters for compression. This is typically used with @@ -121,7 +123,8 @@ compressWindowBits = Stream.defaultWindowBits, compressMemoryLevel = Stream.defaultMemoryLevel, compressStrategy = Stream.defaultStrategy, - compressBufferSize = defaultCompressBufferSize + compressBufferSize = defaultCompressBufferSize, + compressDictionary = Nothing } -- | The default set of parameters for decompression. This is typically used with @@ -130,7 +133,8 @@ defaultDecompressParams :: DecompressParams defaultDecompressParams = DecompressParams { decompressWindowBits = Stream.defaultWindowBits, - decompressBufferSize = defaultDecompressBufferSize + decompressBufferSize = defaultDecompressBufferSize, + decompressDictionary = Nothing } -- | The default chunk sizes for the output of compression and decompression @@ -162,8 +166,8 @@ -- | It is possible to do zlib compression with a custom dictionary. This -- allows slightly higher compression ratios for short files. However such -- compressed streams require the same dictionary when decompressing. This - -- zlib binding does not currently support custom dictionaries. This error - -- is for when we encounter a compressed stream that needs a dictionary. + -- error is for when we encounter a compressed stream that needs a + -- dictionary, and it's not provided. | DictionaryRequired -- | If the compressed data stream is corrupted in any way then you will @@ -214,15 +218,17 @@ -> L.ByteString -> L.ByteString compress format - (CompressParams compLevel method bits memLevel strategy initChunkSize) + (CompressParams compLevel method bits memLevel strategy initChunkSize mdict) input = L.fromChunks $ Stream.run $ do Stream.deflateInit format compLevel method bits memLevel strategy + setDictionary mdict case L.toChunks input of [] -> fillBuffers 20 [] --gzip header is 20 bytes, others even smaller S.PS inFPtr offset length : chunks -> do Stream.pushInputBuffer inFPtr offset length - fillBuffers initChunkSize chunks + r <- fillBuffers initChunkSize chunks + return r where -- we flick between two states: @@ -300,10 +306,22 @@ Stream.Error code msg -> case code of Stream.BufferError -> fail "BufferError should be impossible!" - Stream.NeedDict -> fail "NeedDict is impossible!" + Stream.NeedDict _ -> fail "NeedDict is impossible!" _ -> fail msg + -- Set the custom dictionary, if we were provided with one + -- and if the format supports it (zlib and raw, not gzip). + 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 a data stream. -- -- It will throw an exception if any error is encountered in the input data. If @@ -331,7 +349,7 @@ -> DecompressParams -> L.ByteString -> DecompressStream -decompressWithErrors format (DecompressParams bits initChunkSize) input = +decompressWithErrors format (DecompressParams bits initChunkSize mdict) input = Stream.run $ do Stream.inflateInit format bits case L.toChunks input of @@ -406,7 +424,11 @@ Stream.Error code msg -> case code of Stream.BufferError -> finish (StreamError TruncatedInput msg') where msg' = "premature end of compressed stream" - Stream.NeedDict -> finish (StreamError DictionaryRequired msg) + Stream.NeedDict adler -> do + err <- setDictionary adler mdict + case err of + Just streamErr -> finish streamErr + Nothing -> drainBuffers inChunks Stream.DataError -> finish (StreamError DataError msg) _ -> fail msg @@ -423,3 +445,17 @@ return (StreamChunk (S.PS outFPtr offset length) end) else do Stream.finalise return end + + setDictionary :: Stream.DictionaryHash -> Maybe S.ByteString + -> Stream (Maybe DecompressStream) + setDictionary _adler Nothing = + return $ Just (StreamError DictionaryRequired "custom dictionary needed") + setDictionary _adler (Just dict) = do + status <- Stream.inflateSetDictionary dict + case status of + Stream.Ok -> return Nothing + Stream.Error Stream.StreamError _ -> + return $ Just (StreamError DictionaryRequired "provided dictionary not valid") + Stream.Error Stream.DataError _ -> + return $ Just (StreamError DictionaryRequired "given dictionary does not match the expected one") + _ -> fail "error when setting inflate dictionary" diff --git a/Codec/Compression/Zlib/Stream.hsc b/Codec/Compression/Zlib/Stream.hsc --- a/Codec/Compression/Zlib/Stream.hsc +++ b/Codec/Compression/Zlib/Stream.hsc @@ -30,6 +30,7 @@ zlibFormat, rawFormat, gzipOrZlibFormat, + formatSupportsDictionary, CompressionLevel(..), defaultCompression, noCompression, @@ -70,6 +71,15 @@ outputBufferSpaceRemaining, outputBufferFull, + -- ** Dictionary + deflateSetDictionary, + inflateSetDictionary, + + -- ** Dictionary hashes + DictionaryHash, + dictionaryHash, + zeroDictionaryHash, + #ifdef DEBUG -- * Debugging consistencyCheck, @@ -95,6 +105,8 @@ #endif import Foreign.C import Data.ByteString.Internal (nullForeignPtr) +import qualified Data.ByteString.Unsafe as B +import Data.ByteString (ByteString) import System.IO.Unsafe (unsafeInterleaveIO) import Control.Monad (liftM) import Control.Exception (assert) @@ -235,7 +247,47 @@ setOutAvail (outAvail + outExtra) return result +deflateSetDictionary :: ByteString -> Stream Status +deflateSetDictionary dict = do + err <- withStreamState $ \zstream -> + B.unsafeUseAsCStringLen dict $ \(ptr, len) -> + c_deflateSetDictionary zstream ptr (fromIntegral len) + toStatus err +inflateSetDictionary :: ByteString -> Stream Status +inflateSetDictionary dict = do + err <- withStreamState $ \zstream -> do + B.unsafeUseAsCStringLen dict $ \(ptr, len) -> + c_inflateSetDictionary zstream ptr (fromIntegral len) + toStatus err + +-- | A hash of a custom compression dictionary. These hashes are used by +-- zlib as dictionary identifiers. +-- (The particular hash function used is Adler32.) +-- +newtype DictionaryHash = DictHash CULong + deriving (Eq, Ord, Read, Show) + +-- | Update a running 'DictionaryHash'. You can generate a 'DictionaryHash' +-- from one or more 'ByteString's by starting from 'zeroDictionaryHash', e.g. +-- +-- > dictionaryHash zeroDictionaryHash :: ByteString -> DictionaryHash +-- +-- or +-- +-- > foldl' dictionaryHash zeroDictionaryHash :: [ByteString] -> DictionaryHash +-- +dictionaryHash :: DictionaryHash -> ByteString -> DictionaryHash +dictionaryHash (DictHash adler) dict = + unsafePerformIO $ + B.unsafeUseAsCStringLen dict $ \(ptr, len) -> + liftM DictHash $ c_adler32 adler ptr (fromIntegral len) + +-- | A zero 'DictionaryHash' to use as the initial value with 'dictionaryHash'. +-- +zeroDictionaryHash :: DictionaryHash +zeroDictionaryHash = DictHash 0 + ---------------------------- -- Stream monad -- @@ -404,7 +456,7 @@ | Error ErrorCode String data ErrorCode = - NeedDict + NeedDict DictionaryHash | FileError | StreamError | DataError @@ -420,7 +472,9 @@ toStatus errno = case errno of (#{const Z_OK}) -> return Ok (#{const Z_STREAM_END}) -> return StreamEnd - (#{const Z_NEED_DICT}) -> err NeedDict "custom dictionary needed" + (#{const Z_NEED_DICT}) -> do + adler <- withStreamPtr (#{peek z_stream, adler}) + err (NeedDict (DictHash adler)) "custom dictionary needed" (#{const Z_BUF_ERROR}) -> err BufferError "buffer error" (#{const Z_ERRNO}) -> err FileError "file error" (#{const Z_STREAM_ERROR}) -> err StreamError "stream error" @@ -498,6 +552,10 @@ gzipOrZlibFormat :: Format gzipOrZlibFormat = GZipOrZlib +formatSupportsDictionary :: Format -> Bool +formatSupportsDictionary Zlib = True +formatSupportsDictionary Raw = True +formatSupportsDictionary _ = False -- | The compression method -- @@ -885,6 +943,18 @@ withCAString #{const_str ZLIB_VERSION} $ \versionStr -> c_deflateInit2_ z a b c d e versionStr (#{const sizeof(z_stream)} :: CInt) +foreign import ccall unsafe "zlib.h deflateSetDictionary" + c_deflateSetDictionary :: StreamState + -> Ptr CChar + -> CUInt + -> IO CInt + +foreign import ccall unsafe "zlib.h inflateSetDictionary" + c_inflateSetDictionary :: StreamState + -> Ptr CChar + -> CUInt + -> IO CInt + foreign import ccall unsafe "zlib.h deflate" c_deflate :: StreamState -> CInt -> IO CInt @@ -893,3 +963,9 @@ foreign import ccall unsafe "zlib.h zlibVersion" c_zlibVersion :: IO CString + +foreign import ccall unsafe "zlib.h adler32" + c_adler32 :: CULong + -> Ptr CChar + -> CUInt + -> IO CULong diff --git a/examples/gzip.hs b/examples/gzip.hs --- a/examples/gzip.hs +++ b/examples/gzip.hs @@ -3,4 +3,6 @@ import qualified Data.ByteString.Lazy as B import qualified Codec.Compression.GZip as GZip -main = B.interact GZip.compress +main = B.interact $ GZip.compressWith GZip.defaultCompressParams { + GZip.compressLevel = GZip.BestCompression + } diff --git a/zlib.cabal b/zlib.cabal --- a/zlib.cabal +++ b/zlib.cabal @@ -1,6 +1,6 @@ name: zlib -version: 0.5.3.3 -copyright: (c) 2006-2008 Duncan Coutts +version: 0.5.4.0 +copyright: (c) 2006-2012 Duncan Coutts license: BSD3 license-file: LICENSE author: Duncan Coutts @@ -17,7 +17,7 @@ tasks and for the few cases where more control is needed it provides access to the full zlib feature set. build-type: Simple -cabal-version: >= 1.6 +cabal-version: >= 1.8 extra-source-files: cbits/crc32.h cbits/inffast.h cbits/inflate.h cbits/trees.h cbits/deflate.h cbits/inffixed.h cbits/inftrees.h cbits/zutil.h @@ -36,7 +36,7 @@ other-modules: Codec.Compression.Zlib.Stream extensions: CPP, ForeignFunctionInterface build-depends: base >= 3 && < 5, - bytestring == 0.9.* + bytestring >= 0.9 && < 0.12 includes: zlib.h ghc-options: -Wall if !os(windows)