module Bio.Iteratee.ZLib
(
enumInflate,
enumInflateAny,
enumDeflate,
ZLibParamsException(..),
ZLibException(..),
CompressParams(..),
defaultCompressParams,
DecompressParams(..),
defaultDecompressParams,
Format(..),
CompressionLevel(..),
Method(..),
WindowBits(..),
MemoryLevel(..),
CompressionStrategy(..),
enumSyncFlush,
enumFullFlush,
enumBlockFlush,
)
where
import Bio.Iteratee
import Control.Applicative
import Control.Exception
import Control.Monad ( liftM, liftM2 )
import Data.ByteString as BS
import Data.ByteString.Internal
import Data.Foldable
import Data.Typeable
import Foreign
import Foreign.C
import Prelude
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
data CompressParams = CompressParams {
compressLevel :: !CompressionLevel,
compressMethod :: !Method,
compressWindowBits :: !WindowBits,
compressMemoryLevel :: !MemoryLevel,
compressStrategy :: !CompressionStrategy,
compressBufferSize :: !Int,
compressDictionary :: !(Maybe ByteString)
}
defaultCompressParams :: CompressParams
defaultCompressParams
= CompressParams DefaultCompression Deflated DefaultWindowBits
DefaultMemoryLevel DefaultStrategy (8*1024) Nothing
data DecompressParams = DecompressParams {
decompressWindowBits :: !WindowBits,
decompressBufferSize :: !Int,
decompressDictionary :: !(Maybe ByteString)
}
defaultDecompressParams :: DecompressParams
defaultDecompressParams = DecompressParams DefaultWindowBits (8*1024) Nothing
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 24) zptr ptr
(\hsc_ptr -> pokeByteOff hsc_ptr 32) 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 8) zptr len
pullOutBuffer :: ZStream -> ByteString -> IO ByteString
pullOutBuffer zstr _out = withByteString _out $ \ptr _ -> do
next_out <- withZStream zstr $ \zptr -> (\hsc_ptr -> peekByteOff hsc_ptr 24) 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
type EnumerateeS eli elo m a = (Stream eli -> Iteratee eli m a) -> Iteratee elo m (Iteratee eli m a)
eneeErr :: (Monad m, Exception err, Nullable elo)
=> (Stream eli -> Iteratee eli m a) -> err -> Iteratee elo m ()
eneeErr iter = liftM (const ()) . lift . run . iter . EOF . Just . toException
insertOut :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> Initial
-> Enumeratee ByteString ByteString m a
insertOut size runf (Initial zstr) iter = do
_out <- liftIO $ putOutBuffer size zstr
eneeCheckIfDone (fill size runf (EmptyIn zstr _out)) iter
fill :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> EmptyIn
-> EnumerateeS ByteString ByteString m a
fill size run' (EmptyIn zstr _out) iter
= let fill' (Chunk _in)
| not (BS.null _in) = do
liftIO $ putInBuffer zstr _in
doRun size run' (Invalid zstr _in _out) iter
| otherwise = fillI
fill' (EOF Nothing) = do
out <- liftIO $ pullOutBuffer zstr _out
eneeCheckIfDone (finish size run' (Finishing zstr BS.empty)) $ iter (Chunk out)
fill' (EOF (Just err))
= case fromException err of
Just err' -> flush size run' (Flushing zstr err' _out) iter
Nothing -> throwRecoverableErr err fill'
fillI = liftI fill'
in fillI
swapOut :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> FullOut
-> Enumeratee ByteString ByteString m a
swapOut size run' (FullOut zstr _in) iter = do
_out <- liftIO $ putOutBuffer size zstr
eneeCheckIfDone (doRun size run' (Invalid zstr _in _out)) iter
doRun :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> Invalid
-> EnumerateeS ByteString ByteString m a
doRun size run' (Invalid zstr _in _out) iter = do
status <- liftIO $ run' zstr 0
case fromErrno status of
Left err -> do
eneeErr iter err
throwErr (toException err)
Right False -> do
remaining <- liftIO $ pullInBuffer zstr _in
out <- liftIO $ pullOutBuffer zstr _out
idone (iter (Chunk out)) (Chunk remaining)
Right True -> do
(avail_in, avail_out) <- liftIO $ withZStream zstr $ \zptr -> do
avail_in <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 8) zptr
avail_out <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 32) zptr
return (avail_in, avail_out) :: IO (CInt, CInt)
case avail_out of
0 -> do
out <- liftIO $ pullOutBuffer zstr _out
case avail_in of
0 -> insertOut size run' (Initial zstr) $ iter (Chunk out)
_ -> swapOut size run' (FullOut zstr _in) $ iter (Chunk out)
_ -> case avail_in of
0 -> fill size run' (EmptyIn zstr _out) iter
_ -> do
eneeErr iter IncorrectState
throwErr (toException IncorrectState)
flush :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> Flushing
-> EnumerateeS ByteString ByteString m a
flush size run' (Flushing zstr _flush _out) iter = do
status <- liftIO $ run' zstr (fromFlush _flush)
case fromErrno status of
Left err -> do
eneeErr iter err
throwErr (toException err)
Right False -> do
out <- liftIO $ pullOutBuffer zstr _out
idone (iter (Chunk out)) (Chunk BS.empty)
Right True -> do
avail_out <- liftIO $ withZStream zstr (\hsc_ptr -> peekByteOff hsc_ptr 32)
case avail_out :: CInt of
0 -> do
out <- liftIO $ pullOutBuffer zstr _out
out' <- liftIO $ putOutBuffer size zstr
eneeCheckIfDone (flush size run' (Flushing zstr _flush out')) $ iter (Chunk out)
_ -> insertOut size run' (Initial zstr) (liftI iter)
finish :: MonadIO m
=> Int
-> (ZStream -> CInt -> IO CInt)
-> Finishing
-> EnumerateeS ByteString ByteString m a
finish size run' fin@(Finishing zstr _in) iter = do
_out <- liftIO $ putOutBuffer size zstr
status <- liftIO $ run' zstr 4
case fromErrno status of
Left err -> do
eneeErr iter err
throwErr (toException err)
Right False -> do
remaining <- liftIO $ pullInBuffer zstr _in
out <- liftIO $ pullOutBuffer zstr _out
idone (iter (Chunk out)) (Chunk remaining)
Right True -> do
avail_out <- liftIO $ withZStream zstr (\hsc_ptr -> peekByteOff hsc_ptr 32)
case avail_out :: CInt of
0 -> do
out <- liftIO $ pullOutBuffer zstr _out
eneeCheckIfDone (finish size run' fin) $ iter (Chunk out)
_ -> do
eneeErr iter IncorrectState
throwErr $! toException 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 ())
foreign import ccall unsafe deflateSetDictionary :: Ptr ZStream -> Ptr Word8
-> CUInt -> IO CInt
foreign import ccall unsafe inflateSetDictionary :: Ptr ZStream -> Ptr Word8
-> CUInt -> IO CInt
deflateInit2 :: Ptr ZStream -> CInt -> CInt -> CInt -> CInt -> CInt -> IO CInt
deflateInit2 s l m wB mL s'
= withCString "1.2.8" $ \v ->
deflateInit2_ s l m wB mL s' v (112)
inflateInit2 :: Ptr ZStream -> CInt -> IO CInt
inflateInit2 s wB
= withCString "1.2.8" $ \v ->
inflateInit2_ s wB v (112)
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 (112)
withForeignPtr zstr $ \zptr -> do
_ <- memset (castPtr zptr) 0 (112)
_ <- deflateInit2 zptr c m b l s `finally`
addForeignPtrFinalizer deflateEnd zstr
for_ (compressDictionary cp) $ \(PS fp off len) ->
withForeignPtr fp $ \ptr ->
deflateSetDictionary zptr (ptr `plusPtr` off)
(fromIntegral len)
return $! Right $! Initial $ ZStream zstr
mkDecompress :: Format -> DecompressParams
-> IO (Either ZLibParamsException (Initial, Maybe ByteString))
mkDecompress frm (DecompressParams w _ md)
= case fromWindowBits frm w of
Left err -> return $! Left err
Right wB' -> do
zstr <- mallocForeignPtrBytes (112)
v <- withForeignPtr zstr $ \zptr -> do
_ <- memset (castPtr zptr) 0 (112)
_ <- inflateInit2 zptr wB' `finally`
addForeignPtrFinalizer inflateEnd zstr
case (md, frm) of
(Just (PS fp off len), Raw) -> do
_ <- withForeignPtr fp $ \ptr ->
inflateSetDictionary zptr (ptr `plusPtr` off)
(fromIntegral len)
return Nothing
(Nothing, _) -> return $! Nothing
(Just bs, _) -> return $! (Just bs)
return $! Right $! (Initial $ ZStream zstr, v)
enumDeflate :: MonadIO m
=> Format
-> CompressParams
-> Enumeratee ByteString ByteString m a
enumDeflate f cp@(CompressParams _ _ _ _ _ size _) iter = do
cmp <- liftIO $ mkCompress f cp
case cmp of
Left err -> do
_ <- lift $ enumErr err iter
throwErr (toException err)
Right init' -> insertOut size deflate' init' iter
enumInflate :: MonadIO m
=> Format
-> DecompressParams
-> Enumeratee ByteString ByteString m a
enumInflate f dp@(DecompressParams _ size _md) iter = do
dcmp <- liftIO $ mkDecompress f dp
case dcmp of
Left err -> do
_ <- lift $ enumErr err iter
throwErr (toException err)
Right (init', Nothing) -> insertOut size inflate' init' iter
Right (init', (Just (PS fp off len))) ->
let inflate'' zstr param = do
ret <- inflate' zstr param
case fromErrno ret of
Left NeedDictionary -> do
_ <- withForeignPtr fp $ \ptr ->
withZStream zstr $ \zptr ->
inflateSetDictionary zptr (ptr `plusPtr` off)
(fromIntegral len)
inflate' zstr param
_ -> return ret
in insertOut size inflate'' init' iter
enumInflateAny :: MonadIO m => Enumeratee ByteString ByteString m a
enumInflateAny it = do magic <- iLookAhead $ liftM2 (,) tryHeadBS tryHeadBS
case magic of
(Just 0x1f, Just 0x8b) ->
enumInflate GZip defaultDecompressParams it
>>= enumInflateAny
_ -> mapChunks id it
enumSyncFlush :: Monad m => Enumerator ByteString m a
enumSyncFlush = enumErr SyncFlush
enumFullFlush :: Monad m => Enumerator ByteString m a
enumFullFlush = enumErr FullFlush
enumBlockFlush :: Monad m => Enumerator ByteString m a
enumBlockFlush = enumErr Block