module Data.Iteratee.ZLib
(
enumInflate,
enumDeflate,
ZLibParamsException(..),
ZLibException(..),
CompressParams(..),
defaultCompressParams,
DecompressParams(..),
defaultDecompressParams,
Format(..),
CompressionLevel(..),
Method(..),
WindowBits(..),
MemoryLevel(..),
CompressionStrategy(..),
)
where
import Control.Applicative
import Control.Exception
import Control.Monad.Trans
import Data.ByteString as BS
import Data.ByteString.Internal
import Data.Iteratee
import Data.Iteratee.IO
import Data.Typeable
import Foreign
import Foreign.C
data ZLibParamsException
= IncorrectCompressionLevel !Int
| IncorrectWindowBits !Int
| IncorrectMemoryLevel !Int
deriving (Eq,Typeable)
data ZLibException
= NeedDictionary
| BufferError
| StreamError
| DataError
| MemoryError
| VersionError
| Unexpected !CInt
| IncorrectState
deriving (Eq,Typeable)
data ZlibFlush
= SyncFlush
| FullFlush
| Block
deriving (Eq,Typeable)
instance Show ZlibFlush where
show SyncFlush = "zlib: flush requested"
show FullFlush = "zlib: full flush requested"
show Block = "zlib: block flush requested"
instance Exception ZlibFlush
fromFlush :: ZlibFlush -> CInt
fromFlush SyncFlush = 2
fromFlush FullFlush = 3
fromFlush Block = 5
instance Show ZLibParamsException where
show (IncorrectCompressionLevel lvl)
= "zlib: incorrect compression level " ++ show lvl
show (IncorrectWindowBits lvl)
= "zlib: incorrect window bits " ++ show lvl
show (IncorrectMemoryLevel lvl)
= "zlib: incorrect memory level " ++ show lvl
instance Show ZLibException where
show NeedDictionary = "zlib: needs dictionary"
show BufferError = "zlib: no progress is possible (internal error)"
show StreamError = "zlib: stream error"
show DataError = "zlib: data error"
show MemoryError = "zlib: memory error"
show VersionError = "zlib: version error"
show (Unexpected lvl) = "zlib: unknown error " ++ show lvl
show IncorrectState = "zlib: incorrect state"
instance Exception ZLibParamsException
instance Exception ZLibException
newtype ZStream = ZStream (ForeignPtr ZStream)
withZStream :: ZStream -> (Ptr ZStream -> IO a) -> IO a
withZStream (ZStream fptr) = withForeignPtr fptr
mallocZStream :: IO ZStream
mallocZStream = ZStream <$> mallocForeignPtrBytes (56)
data CompressParams = CompressParams {
compressLevel :: !CompressionLevel,
compressMethod :: !Method,
compressWindowBits :: !WindowBits,
compressMemoryLevel :: !MemoryLevel,
compressStrategy :: !CompressionStrategy,
compressBufferSize :: !Int
}
defaultCompressParams
= CompressParams DefaultCompression Deflated DefaultWindowBits
DefaultMemoryLevel DefaultStrategy (8*1024)
data DecompressParams = DecompressParams {
decompressWindowBits :: !WindowBits,
decompressBufferSize :: !Int
}
defaultDecompressParams = DecompressParams DefaultWindowBits (8*1024)
data Format
= GZip
| Zlib
| Raw
| GZipOrZlib
deriving (Eq)
data CompressionLevel
= DefaultCompression
| NoCompression
| BestSpeed
| BestCompression
| CompressionLevel Int
data Method
= Deflated
data WindowBits
= WindowBits Int
| DefaultWindowBits
data MemoryLevel
= DefaultMemoryLevel
| MinMemoryLevel
| MaxMemoryLevel
| MemoryLevel Int
data CompressionStrategy
= DefaultStrategy
| Filtered
| HuffmanOnly
fromMethod :: Method -> CInt
fromMethod Deflated = 8
fromCompressionLevel :: CompressionLevel -> Either ZLibParamsException CInt
fromCompressionLevel DefaultCompression = Right $! 1
fromCompressionLevel NoCompression = Right $! 0
fromCompressionLevel BestSpeed = Right $! 1
fromCompressionLevel BestCompression = Right $! 9
fromCompressionLevel (CompressionLevel n)
| n >= 0 && n <= 9 = Right $! fromIntegral $! n
| otherwise = Left $! IncorrectCompressionLevel n
fromWindowBits :: Format -> WindowBits -> Either ZLibParamsException CInt
fromWindowBits format bits
= formatModifier format <$> checkWindowBits bits
where checkWindowBits DefaultWindowBits = Right $! 15
checkWindowBits (WindowBits n)
| n >= 8 && n <= 15 = Right $! fromIntegral $! n
| otherwise = Left $! IncorrectWindowBits $! n
formatModifier Zlib = id
formatModifier GZip = (+16)
formatModifier GZipOrZlib = (+32)
formatModifier Raw = negate
fromMemoryLevel :: MemoryLevel -> Either ZLibParamsException CInt
fromMemoryLevel DefaultMemoryLevel = Right $! 8
fromMemoryLevel MinMemoryLevel = Right $! 1
fromMemoryLevel MaxMemoryLevel = Right $! 9
fromMemoryLevel (MemoryLevel n)
| n >= 1 && n <= 9 = Right $! fromIntegral n
| otherwise = Left $! IncorrectMemoryLevel $! fromIntegral n
fromCompressionStrategy :: CompressionStrategy -> CInt
fromCompressionStrategy DefaultStrategy = 0
fromCompressionStrategy Filtered = 1
fromCompressionStrategy HuffmanOnly = 2
fromErrno :: CInt -> Either ZLibException Bool
fromErrno (0) = Right $! True
fromErrno (1) = Right $! False
fromErrno (2) = Left $! NeedDictionary
fromErrno (5) = Left $! BufferError
fromErrno (2) = Left $! StreamError
fromErrno (3) = Left $! DataError
fromErrno (4) = Left $! MemoryError
fromErrno (6) = Left $! VersionError
fromErrno n = Left $! Unexpected n
convParam :: Format
-> CompressParams
-> Either ZLibParamsException (CInt, CInt, CInt, CInt, CInt)
convParam f (CompressParams c m w l s _)
= let c' = fromCompressionLevel c
m' = fromMethod m
b' = fromWindowBits f w
l' = fromMemoryLevel l
s' = fromCompressionStrategy s
eit = either Left
r = Right
in eit (\c_ -> eit (\b_ -> eit (\l_ -> r (c_, m', b_, l_, s')) l') b') c'
newtype Initial = Initial ZStream
data EmptyIn = EmptyIn !ZStream !ByteString
data FullOut = FullOut !ZStream !ByteString
data Invalid = Invalid !ZStream !ByteString !ByteString
data Finishing = Finishing !ZStream !ByteString
data Flushing = Flushing !ZStream !ZlibFlush !ByteString
withByteString :: ByteString -> (Ptr Word8 -> Int -> IO a) -> IO a
withByteString (PS ptr off len) f
= withForeignPtr ptr (\ptr' -> f (ptr' `plusPtr` off) len)
mkByteString :: MonadIO m => Int -> m ByteString
mkByteString s = liftIO $ create s (\_ -> return ())
putOutBuffer :: Int -> ZStream -> IO ByteString
putOutBuffer size zstr = do
_out <- mkByteString size
withByteString _out $ \ptr len -> withZStream zstr $ \zptr -> do
(\hsc_ptr -> pokeByteOff hsc_ptr 12) zptr ptr
(\hsc_ptr -> pokeByteOff hsc_ptr 16) zptr len
return _out
putInBuffer :: ZStream -> ByteString -> IO ()
putInBuffer zstr _in
= withByteString _in $ \ptr len -> withZStream zstr $ \zptr -> do
(\hsc_ptr -> pokeByteOff hsc_ptr 0) zptr ptr
(\hsc_ptr -> pokeByteOff hsc_ptr 4) zptr len
pullOutBuffer :: ZStream -> ByteString -> IO ByteString
pullOutBuffer zstr _out = withByteString _out $ \ptr _ -> do
next_out <- withZStream zstr $ \zptr -> (\hsc_ptr -> peekByteOff hsc_ptr 12) zptr
return $! BS.take (next_out `minusPtr` ptr) _out
pullInBuffer :: ZStream -> ByteString -> IO ByteString
pullInBuffer zstr _in = withByteString _in $ \ptr _ -> do
next_in <- withZStream zstr $ \zptr -> (\hsc_ptr -> peekByteOff hsc_ptr 0) zptr
return $! BS.drop (next_in `minusPtr` ptr) _in
insertOut :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> Initial
-> Enumerator ByteString m a
insertOut size run (Initial zstr) iter = return $! do
_out <- liftIO $ putOutBuffer size zstr
joinIM $ fill size run (EmptyIn zstr _out) iter
fill :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> EmptyIn
-> Enumerator ByteString m a
fill size run (EmptyIn zstr _out) iter
= let fill' (Chunk _in)
| not (BS.null _in) = do
liftIO $ putInBuffer zstr _in
joinIM $ doRun size run (Invalid zstr _in _out) iter
| otherwise = fillI
fill' (EOF Nothing)
= joinIM $ finish size run (Finishing zstr BS.empty) iter
fill' (EOF (Just err))
= case fromException err of
Just err' ->
joinIM $ flush size run (Flushing zstr err' _out) iter
Nothing -> throwRecoverableErr err fill'
fillI = liftI fill'
in return $! fillI
swapOut :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> FullOut
-> Enumerator ByteString m a
swapOut size run (FullOut zstr _in) iter = return $! do
_out <- liftIO $ putOutBuffer size zstr
joinIM $ doRun size run (Invalid zstr _in _out) iter
doRun :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> Invalid
-> Enumerator ByteString m a
doRun size run (Invalid zstr _in _out) iter = return $! do
status <- liftIO $ run zstr 0
case fromErrno status of
Left err -> joinIM $ enumErr err iter
Right False -> do
remaining <- liftIO $ pullInBuffer zstr _in
out <- liftIO $ pullOutBuffer zstr _out
iter' <- lift $ enumPure1Chunk out iter
res <- lift $ tryRun iter'
case res of
Left err@(SomeException _) -> throwErr err
Right x -> idone x (Chunk remaining)
Right True -> do
(avail_in, avail_out) <- liftIO $ withZStream zstr $ \zptr -> do
avail_in <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 4) zptr
avail_out <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 16) zptr
return (avail_in, avail_out) :: IO (CInt, CInt)
case avail_out of
0 -> do
out <- liftIO $ pullOutBuffer zstr _out
iter' <- lift $ enumPure1Chunk out iter
joinIM $ case avail_in of
0 -> insertOut size run (Initial zstr) iter'
_ -> swapOut size run (FullOut zstr _in) iter'
_ -> joinIM $ case avail_in of
0 -> fill size run (EmptyIn zstr _out) iter
_ -> enumErr IncorrectState iter
flush :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> Flushing
-> Enumerator ByteString m a
flush size run fin@(Flushing zstr _flush _out) iter = return $! do
status <- liftIO $ run zstr (fromFlush _flush)
case fromErrno status of
Left err -> joinIM $ enumErr err iter
Right False -> do
out <- liftIO $ pullOutBuffer zstr _out
iter' <- lift $ enumPure1Chunk out iter
res <- lift $ tryRun iter'
case res of
Left err@(SomeException _) -> throwErr err
Right x -> idone x (Chunk BS.empty)
Right True -> do
(avail_in, avail_out) <- liftIO $ withZStream zstr $ \zptr -> do
avail_in <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 4) zptr
avail_out <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 16) zptr
return (avail_in, avail_out) :: IO (CInt, CInt)
case avail_out of
0 -> do
out <- liftIO $ pullOutBuffer zstr _out
iter' <- lift $ enumPure1Chunk out iter
out' <- liftIO $ putOutBuffer size zstr
joinIM $ flush size run (Flushing zstr _flush out') iter'
_ -> joinIM $ insertOut size run (Initial zstr) iter
finish :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> Finishing
-> Enumerator ByteString m a
finish size run fin@(Finishing zstr _in) iter = return $! do
_out <- liftIO $ putOutBuffer size zstr
status <- liftIO $ run zstr 4
case fromErrno status of
Left err -> joinIM $ enumErr err iter
Right False -> do
remaining <- liftIO $ pullInBuffer zstr _in
out <- liftIO $ pullOutBuffer zstr _out
iter' <- lift $ enumPure1Chunk out iter
res <- lift $ tryRun iter'
case res of
Left err@(SomeException _) -> throwErr err
Right x -> idone x (Chunk remaining)
Right True -> do
(avail_in, avail_out) <- liftIO $ withZStream zstr $ \zptr -> do
avail_in <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 4) zptr
avail_out <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 16) zptr
return (avail_in, avail_out) :: IO (CInt, CInt)
case avail_out of
0 -> do
out <- liftIO $ withZStream zstr $ \zptr ->
pullOutBuffer zstr _out
iter' <- lift $ enumPure1Chunk out iter
joinIM $ finish size run fin iter'
_ -> throwErr $! SomeException IncorrectState
foreign import ccall unsafe deflateInit2_ :: Ptr ZStream -> CInt -> CInt
-> CInt -> CInt -> CInt
-> CString -> CInt -> IO CInt
foreign import ccall unsafe inflateInit2_ :: Ptr ZStream -> CInt
-> CString -> CInt -> IO CInt
foreign import ccall unsafe inflate :: Ptr ZStream -> CInt -> IO CInt
foreign import ccall unsafe deflate :: Ptr ZStream -> CInt -> IO CInt
foreign import ccall unsafe "&deflateEnd"
deflateEnd :: FunPtr (Ptr ZStream -> IO ())
foreign import ccall unsafe "&inflateEnd"
inflateEnd :: FunPtr (Ptr ZStream -> IO ())
deflateInit2 :: Ptr ZStream -> CInt -> CInt -> CInt -> CInt -> CInt -> IO CInt
deflateInit2 s l m wB mL s'
= withCString "1.2.3.4" $ \v ->
deflateInit2_ s l m wB mL s' v (56)
inflateInit2 :: Ptr ZStream -> CInt -> IO CInt
inflateInit2 s wB
= withCString "1.2.3.4" $ \v ->
inflateInit2_ s wB v (56)
deflate' :: ZStream -> CInt -> IO CInt
deflate' z f = withZStream z $ \p -> deflate p f
inflate' :: ZStream -> CInt -> IO CInt
inflate' z f = withZStream z $ \p -> inflate p f
mkCompress :: Format -> CompressParams
-> IO (Either ZLibParamsException Initial)
mkCompress frm cp
= case convParam frm cp of
Left err -> return $! Left err
Right (c, m, b, l, s) -> do
zstr <- mallocForeignPtrBytes (56)
withForeignPtr zstr $ \zptr -> do
memset (castPtr zptr) 0 (56)
deflateInit2 zptr c m b l s `finally`
addForeignPtrFinalizer deflateEnd zstr
return $! Right $! Initial $ ZStream zstr
mkDecompress :: Format -> DecompressParams
-> IO (Either ZLibParamsException Initial)
mkDecompress frm cp@(DecompressParams wB _)
= case fromWindowBits frm wB of
Left err -> return $! Left err
Right wB' -> do
zstr <- mallocForeignPtrBytes (56)
withForeignPtr zstr $ \zptr -> do
memset (castPtr zptr) 0 (56)
inflateInit2 zptr wB' `finally`
addForeignPtrFinalizer inflateEnd zstr
return $! Right $! Initial $ ZStream zstr
enumDeflate :: MonadIO m
=> Format
-> CompressParams
-> Enumerator ByteString m a
enumDeflate f cp@(CompressParams _ _ _ _ _ size) iter = do
cmp <- liftIO $ mkCompress f cp
case cmp of
Left err -> enumErr err iter
Right init -> insertOut size deflate' init iter
enumInflate :: MonadIO m
=> Format
-> DecompressParams
-> Enumerator ByteString m a
enumInflate f dp@(DecompressParams _ size) iter = do
dcmp <- liftIO $ mkDecompress f dp
case dcmp of
Left err -> enumErr err iter
Right init -> insertOut size deflate' init iter