forked from tibbe/hyena
-
Notifications
You must be signed in to change notification settings - Fork 1
/
Enumerator.hs
146 lines (128 loc) · 4.41 KB
/
Enumerator.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
{-# LANGUAGE Rank2Types #-}
module Data.Enumerator
( EnumeratorM,
IterateeM,
-- Enumerators
bytesEnum,
lazyBytesEnum,
chunkEnum,
partialSocketEnum,
socketEnum,
-- Combining enumerators
compose
) where
import Control.Monad (liftM)
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Char8 as C (unpack)
import Data.Word (Word8)
import Network.Socket (Socket)
import Network.Socket.ByteString (recv)
import Numeric (readHex)
type IterateeM a m = a -> S.ByteString -> m (Either a a)
type EnumeratorM m = forall a. IterateeM a m -> a -> m a
-- -----------------------------------------------------------
-- Enumerators
-- | Enumerates a 'ByteString'.
bytesEnum :: Monad m => S.ByteString -> EnumeratorM m
bytesEnum bs f seed =
let block = S.take blockSize bs
in if S.null block
then return seed
else do
seed' <- f seed block
case seed' of
Left seed'' -> return seed''
Right seed'' -> bytesEnum (S.drop blockSize bs) f seed''
-- | Enumerates a lazy 'ByteString'.
lazyBytesEnum :: Monad m => L.ByteString -> EnumeratorM m
lazyBytesEnum bs f seed = go (L.toChunks bs) f seed
where
go [] _ sd = return sd
go (x:xs) g sd = do
sd' <- bytesEnum x g sd
go xs g sd'
nl :: Word8
nl = 10
-- | Enumerates chunks of data encoded using HTTP chunked encoding.
chunkEnum :: Monad m => EnumeratorM m -> EnumeratorM m
chunkEnum enum f initSeed = fst `liftM` enum go (initSeed, Left S.empty)
where
go (seed, Left acc) bs =
case S.elemIndex nl bs of
Just ix -> let (line, rest) = S.splitAt (ix + 1) bs
hdr = S.append acc line
chunkLen = pHeader hdr
in case chunkLen of
Just n -> go (seed, Right n) rest
Nothing -> error $ "malformed header" ++ show hdr
Nothing -> return $ Right (seed, Left (S.append acc bs))
go (seed, Right n) bs =
let len = S.length bs
in if len < n
then do
seed' <- f seed bs
case seed' of
Right seed'' -> return $ Right (seed'', Right $! n - len)
Left seed'' -> return $ Left (seed'', Right $! n - len)
else let (bs', rest) = S.splitAt n bs
in do
seed' <- f seed bs'
case seed' of
Right seed'' -> go (seed'', Left S.empty) rest
Left seed'' -> return $ Left (seed'', Left rest)
-- TODO: Ignore header.
pHeader :: S.ByteString -> Maybe Int
pHeader bs =
case readHex $ C.unpack hdr of
[(n, "")] -> Just n
_ -> Nothing
where
hdr = S.take (S.length bs - 2) bs
-- | Maximum number of bytes sent or received in every socket
-- operation.
blockSize :: Int
blockSize = 4 * 1024
-- | @partialSocketEnum sock numBytes@ enumerates @numBytes@ bytes
-- received through the given socket. Does not close the socket.
partialSocketEnum :: Socket -> Int -> EnumeratorM IO
partialSocketEnum sock numBytes f initSeed = go initSeed numBytes
where
go seed 0 = return seed
go seed n = do
bs <- recv sock blockSize
if S.null bs
then return seed
else do
seed' <- f seed bs
case seed' of
Right seed'' -> go seed'' $! n - S.length bs
Left seed'' -> return seed''
-- | Enumerates data received through the given socket. Does not
-- close the socket.
socketEnum :: Socket -> EnumeratorM IO
socketEnum sock f initSeed = go initSeed
where
go seed = do
bs <- recv sock blockSize
if S.null bs
then return seed
else do
seed' <- f seed bs
case seed' of
Right seed'' -> go seed''
Left seed'' -> return seed''
-- -----------------------------------------------------------
-- Combining enumerators
-- Make two enumerators behave like one.
compose :: Monad m => EnumeratorM m -> EnumeratorM m -> EnumeratorM m
compose enum1 enum2 f initSeed = enum1 f1 (Right initSeed) >>= k
where
f1 (Right seed) bs = do
r <- f seed bs
case r of
x@(Right _) -> return $ Right x
x -> return $ Left x
f1 x _ = return $ Left x -- Cannot happen.
k (Left seed) = return seed
k (Right seed) = enum2 f seed