@@ -33,28 +33,42 @@ compressInner :: (acc -> ByteString -> IO (Either acc acc))
3333 -> ZStream
3434 -> IO (Either acc acc)
3535compressInner iter acc enum stream = allocaBytes chunkSize $ \outbuff -> do
36− when (stream == nullPtr) $ fail "Error from zlib library"
36+ when (stream == nullPtr) $ error "Error from zlib library"
37+ c_set_avail_out stream outbuff $ fromIntegral chunkSize
3738 eacc <- runEnumerator enum (compressIter iter stream outbuff) acc
3839 case eacc of
3940 Left acc' -> return $ Left acc'
4041 Right acc' -> finishStream iter stream outbuff acc'
4142
43+drainOutput :: Bool -- ^ are we finishing?
44+ -> ZStream
45+ -> Ptr CChar -- ^ output buffer
46+ -> (acc -> ByteString -> IO (Either acc acc))
47+ -> acc
48+ -> IO (Either acc acc)
49+drainOutput isFinish stream outbuff iter acc = do
50+ if isFinish
51+ then c_call_deflate_finish stream >> return ()
52+ else c_call_deflate_noflush stream
53+ newAvail <- fromIntegral `fmap` c_get_avail_out stream
54+ if newAvail == 0 || (isFinish && newAvail /= chunkSize)
55+ then do
56+ bs <- unsafePackCStringLen (outbuff, chunkSize - newAvail)
57+ eacc <- iter acc bs
58+ case eacc of
59+ Left acc' -> return $ Left acc'
60+ Right acc' -> do
61+ c_set_avail_out stream outbuff $ fromIntegral chunkSize
62+ drainOutput isFinish stream outbuff iter acc'
63+ else return $ Right acc
64+
4265finishStream :: (acc -> ByteString -> IO (Either acc acc))
4366 -> ZStream
4467 -> Ptr CChar
4568 -> acc
4669 -> IO (Either acc acc)
47−finishStream iter stream outbuff acc = do
48− c_set_avail_out stream outbuff $ fromIntegral chunkSize
49− status <- c_call_deflate_finish stream
50− newAvail <- fromIntegral `fmap` c_get_avail_out stream
51− bs <- unsafePackCStringLen (outbuff, chunkSize - newAvail)
52− eacc <- iter acc bs
53− if status == 0 || newAvail == 0
54− then case eacc of
55− Left acc' -> return $ Left acc'
56− Right acc' -> finishStream iter stream outbuff acc'
57− else return eacc
70+finishStream iter stream outbuff acc =
71+ drainOutput True stream outbuff iter acc
5872
5973compressIter :: (acc -> ByteString -> IO (Either acc acc))
6074 -> ZStream
@@ -62,22 +76,10 @@ compressIter :: (acc -> ByteString -> IO (Either acc acc))
6276 -> acc
6377 -> ByteString
6478 -> IO (Either acc acc)
65−compressIter iter stream outbuff accInit bsInput = do
79+compressIter iter stream outbuff acc bsInput = do
6680 unsafeUseAsCStringLen bsInput $ \(cstr, len) ->
6781 c_set_avail_in stream cstr $ fromIntegral len
68− go accInit
69− where
70− go acc = do
71− c_set_avail_out stream outbuff $ fromIntegral chunkSize
72− c_call_deflate_noflush stream
73− newAvail <- fromIntegral `fmap` c_get_avail_out stream
74− bs <- unsafePackCStringLen (outbuff, chunkSize - newAvail)
75− eacc <- iter acc bs
76− if newAvail == 0
77− then case eacc of
78− Left acc' -> return $ Left acc'
79− Right acc' -> go acc'
80− else return eacc
82+ drainOutput False stream outbuff iter acc
8183
8284data ZStreamStruct
8385type ZStream = Ptr ZStreamStruct