@@ -0,0 +1,104 @@
1+{-# LANGUAGE CPP, ForeignFunctionInterface #-}
2+{-# LANGUAGE EmptyDataDecls #-}
3+module Network.Wai.Zlib (compress) where
4+
5+import Prelude hiding (length)
6+import Network.Wai hiding (status)
7+
8+import Data.ByteString.Unsafe
9+import Data.ByteString.Lazy.Internal (defaultChunkSize)
10+
11+import Foreign.C
12+import Foreign.Ptr
13+import Foreign.Marshal.Alloc
14+import Control.Exception
15+import Control.Monad
16+
17+import Data.ByteString.Char8 hiding (putStrLn)
18+
19+chunkSize :: Int
20+chunkSize = defaultChunkSize
21+
22+compress :: Enumerator -> Enumerator
23+compress enum =
24+ Enumerator $ \iter acc ->
25+ bracket
26+ c_create_z_stream
27+ c_free_z_stream
28+ (compressInner iter acc enum)
29+
30+compressInner :: (acc -> ByteString -> IO (Either acc acc))
31+ -> acc
32+ -> Enumerator
33+ -> ZStream
34+ -> IO (Either acc acc)
35+compressInner iter acc enum stream = allocaBytes chunkSize $ \outbuff -> do
36+ when (stream == nullPtr) $ fail "Error from zlib library"
37+ eacc <- runEnumerator enum (compressIter iter stream outbuff) acc
38+ case eacc of
39+ Left acc' -> return $ Left acc'
40+ Right acc' -> finishStream iter stream outbuff acc'
41+
42+finishStream :: (acc -> ByteString -> IO (Either acc acc))
43+ -> ZStream
44+ -> Ptr CChar
45+ -> acc
46+ -> 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 <- packCStringLen (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
58+
59+compressIter :: (acc -> ByteString -> IO (Either acc acc))
60+ -> ZStream
61+ -> Ptr CChar
62+ -> acc
63+ -> ByteString
64+ -> IO (Either acc acc)
65+compressIter iter stream outbuff accInit bsInput = do
66+ unsafeUseAsCStringLen bsInput $ \(cstr, len) ->
67+ 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 <- packCStringLen (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
81+
82+data ZStreamStruct
83+type ZStream = Ptr ZStreamStruct
84+
85+foreign import ccall unsafe "create_z_stream"
86+ c_create_z_stream :: IO ZStream
87+
88+foreign import ccall unsafe "free_z_stream"
89+ c_free_z_stream :: ZStream -> IO ()
90+
91+foreign import ccall unsafe "set_avail_in"
92+ c_set_avail_in :: ZStream -> Ptr CChar -> CUInt -> IO ()
93+
94+foreign import ccall unsafe "set_avail_out"
95+ c_set_avail_out :: ZStream -> Ptr CChar -> CUInt -> IO ()
96+
97+foreign import ccall unsafe "call_deflate_noflush"
98+ c_call_deflate_noflush :: ZStream -> IO ()
99+
100+foreign import ccall unsafe "call_deflate_finish"
101+ c_call_deflate_finish :: ZStream -> IO CInt
102+
103+foreign import ccall unsafe "get_avail_out"
104+ c_get_avail_out :: ZStream -> IO CUInt