module Data.IterIO.Zlib (
ZState, deflateInit2, inflateInit2
, inumZState, inumZlib, inumGzip, inumGunzip
, max_wbits, max_mem_level, def_mem_level, zlib_version
, z_DEFAULT_COMPRESSION
, ZStrategy, z_FILTERED, z_HUFFMAN_ONLY, z_RLE
, z_FIXED, z_DEFAULT_STRATEGY
, ZMethod, z_DEFLATED
) where
import Prelude hiding (null)
import Control.Exception (throwIO, ErrorCall(..))
import Control.Monad.State.Strict
import qualified Data.ByteString.Internal as S
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Lazy.Internal as L
import Foreign
import Foreign.C
import Data.IterIO.Iter
import Data.IterIO.Inum
import Data.IterIO.ZlibInt
data ZState = ZState { zStream :: (ForeignPtr ZStream)
, zOp :: (ZFlush -> IO CInt)
, zFinish :: !ZFlush
, zInChunk :: !(ForeignPtr Word8)
, zOutChunk :: !(ForeignPtr Word8)
, zOut :: L.ByteString -> L.ByteString
}
defaultZState :: ZState
defaultZState = ZState { zStream = error "must allocate zStream"
, zOp = error "must define zOp"
, zFinish = z_FINISH
, zInChunk = S.nullForeignPtr
, zOutChunk = S.nullForeignPtr
, zOut = id
}
newZStream :: (Ptr ZStream -> IO CInt) -> IO (ForeignPtr ZStream)
newZStream initfn = do
zs <- mallocForeignPtrBytes z_stream_size
withForeignPtr zs $ \ptr ->
do _ <- S.memset (castPtr ptr) 0 z_stream_size
err <- initfn ptr
when (err /= z_OK) $ throwIO $ ErrorCall "newZStream: init failed"
return zs
deflateInit2 :: CInt
-> ZMethod
-> CInt
-> CInt
-> ZStrategy
-> IO ZState
deflateInit2 level method windowBits memLevel strategy = do
z <- newZStream $ \ptr -> (c_deflateInit2 ptr level method windowBits
memLevel strategy zlib_version z_stream_size)
addForeignPtrFinalizer c_deflateEnd z
return defaultZState { zStream = z
, zOp = \flush -> withForeignPtr z $ \zp ->
c_deflate zp flush
}
inflateInit2 :: CInt
-> IO ZState
inflateInit2 windowBits = do
z <- newZStream $ \ptr -> (c_inflateInit2 ptr windowBits
zlib_version z_stream_size)
addForeignPtrFinalizer c_inflateEnd z
return defaultZState { zStream = z
, zOp = \flush -> withForeignPtr z $ \zp ->
c_inflate zp flush
, zFinish = z_NO_FLUSH
}
type ZM = StateT ZState IO
withZFP :: (ZState -> ForeignPtr a) -> (Ptr a -> ZM b) -> ZM b
withZFP field k = StateT $ \zs ->
withForeignPtr (field zs) $ \v -> (runStateT $ k v) zs
zPeek :: (Storable a) => (Ptr ZStream -> Ptr a) -> ZM a
zPeek f = withZFP zStream $ liftIO . peek . f
zPoke :: (Storable a) => (Ptr ZStream -> Ptr a) -> a -> ZM ()
zPoke f a = withZFP zStream $ liftIO . flip poke a . f
zPokeFP :: (Ptr ZStream -> Ptr (Ptr Word8)) -> ForeignPtr Word8 -> Int -> ZM ()
zPokeFP f fp offset = withZFP zStream $ \z ->
liftIO $ withForeignPtr fp $ \p ->
poke (f z) $ p `plusPtr` offset
zMinusPtr :: (Ptr ZStream -> Ptr (Ptr Word8))
-> (ZState -> ForeignPtr Word8)
-> ZM Int
zMinusPtr curf basef = withZFP basef $ \base ->
if base == nullPtr
then return 0
else do
cur <- zPeek curf
return $ cur `minusPtr` base
zPushIn :: L.ByteString -> ZM L.ByteString
zPushIn s = do
avail <- zPeek avail_in
if avail > 0 then return s else pushit s
where
pushit (L.Chunk h t) = do
let (fp, offset, len) = S.toForeignPtr h
modify $ \zs -> zs { zInChunk = fp }
zPokeFP next_in fp offset
zPoke avail_in $ fromIntegral len
return t
pushit L.Empty = return L.Empty
zPopIn :: L.ByteString -> ZM L.ByteString
zPopIn s = do
len <- zPeek avail_in
if len <= 0
then return s
else do
fptr <- gets zInChunk
offset <- zMinusPtr next_in zInChunk
zPoke avail_in 0
return $ L.chunk (S.fromForeignPtr fptr offset $ fromIntegral len) s
zOutLen :: ZM Int
zOutLen = zMinusPtr next_out zOutChunk
zPopOut :: ZM ()
zPopOut = do
len <- zOutLen
when (len > 0) $ do
ochunk <- liftM (\c -> S.fromForeignPtr c 0 len) $ gets zOutChunk
out <- liftM (. L.chunk ochunk) $ gets zOut
modify $ \zs -> zs { zOutChunk = S.nullForeignPtr
, zOut = out }
zPoke avail_out 0
zMkSpace :: ZM ()
zMkSpace = do
avail <- zPeek avail_out
when (avail <= 0) $ do
zPopOut
nchunk <- liftIO $ S.mallocByteString L.defaultChunkSize
zPokeFP next_out nchunk 0
zPoke avail_out $ fromIntegral L.defaultChunkSize
modify $ \zs -> zs { zOutChunk = nchunk }
zExec :: ZFlush -> ZM CInt
zExec flush = do
zMkSpace
op <- gets zOp
r <- withZFP zInChunk $ \_ -> liftIO $ op flush
avail <- zPeek avail_out
case () of
_ | r == z_OK && avail == 0 -> zExec flush
_ | r == z_NEED_DICT -> liftIO $ throwIO $ ErrorCall "zlib NEED_DICT"
_ | r == z_STREAM_END -> do zPopOut
return r
_ | r < 0 -> do cm <- zPeek msg
m <- if cm == nullPtr
then return $ "zlib failed ("
++ show r ++ ")"
else liftIO $ peekCString cm
liftIO $ throwIO $ ErrorCall m
_ | otherwise -> return r
inumZState :: (MonadIO m) =>
ZState
-> Inum L.ByteString L.ByteString m a
inumZState = mkInumM . loop
where
loop zs0 = do
(Chunk dat eof) <- chunkI
((r, rest), zs) <- liftIO (runStateT (runz eof dat) zs0)
ungetI rest
done <- ifeed $ zOut zs L.empty
unless (done || eof || r == z_STREAM_END) $ loop zs { zOut = id }
runz False L.Empty = return (z_OK, L.Empty)
runz eof s0 = do
s <- zPushIn s0
flush <- if eof && L.null s then gets zFinish else return z_NO_FLUSH
r <- zExec flush
if r == z_STREAM_END || L.null s
then do s' <- zPopIn s; return (r, s')
else runz eof s
inumZlib :: (MonadIO m) => Inum L.ByteString L.ByteString m a
inumZlib iter = do
zs <- liftIO (deflateInit2 z_DEFAULT_COMPRESSION z_DEFLATED max_wbits
def_mem_level z_DEFAULT_STRATEGY)
inumZState zs iter
inumGzip :: (MonadIO m) => Inum L.ByteString L.ByteString m a
inumGzip iter = do
zs <- liftIO (deflateInit2 z_DEFAULT_COMPRESSION z_DEFLATED (16 + max_wbits)
def_mem_level z_DEFAULT_STRATEGY)
inumZState zs iter
inumGunzip :: (MonadIO m) => Inum L.ByteString L.ByteString m a
inumGunzip iter = do
zs <- liftIO $ inflateInit2 (32 + max_wbits)
inumZState zs iter