module System.IterateeM where
import System.Posix
import Foreign.C
import Foreign.Ptr
import Foreign.Marshal.Alloc
import Data.List (splitAt)
import Data.Char (isHexDigit, digitToInt, isSpace)
import Control.Monad.Trans
import Control.Monad.Identity
import System.LowLevelIO
data StreamG a = EOF | Err String | Chunk [a] deriving Show
type Stream = StreamG Char
data IterateeG el m a = IE_done a (StreamG el)
| IE_cont (StreamG el -> IterateeGM el m a)
newtype IterateeGM el m a = IM{unIM:: m (IterateeG el m a)}
type Iteratee m a = IterateeG Char m a
type IterateeM m a = IterateeGM Char m a
liftI :: Monad m => IterateeG el m a -> IterateeGM el m a
liftI = IM . return
infixl 1 >>==
(>>==):: Monad m =>
IterateeGM el m a ->
(IterateeG el m a -> IterateeGM el' m b) ->
IterateeGM el' m b
m >>== f = IM (unIM m >>= unIM . f)
infixr 1 ==<<
(==<<) :: Monad m =>
(IterateeG el m a -> IterateeGM el' m b) ->
IterateeGM el m a ->
IterateeGM el' m b
f ==<< m = m >>== f
joinI :: Monad m => IterateeGM el m (IterateeG el' m a) -> IterateeGM el m a
joinI m = m >>= (\iter -> enum_eof iter >>== check)
where
check (IE_done x (Err str)) = liftI $ (IE_done x (Err str))
check (IE_done x _) = liftI $ (IE_done x EOF)
check (IE_cont _) = error "joinI: can't happen: EOF didn't terminate"
instance Monad m => Monad (IterateeGM el m) where
return x = liftI $ IE_done x (Chunk [])
m >>= f = m >>== docase
where
docase (IE_done a (Chunk [])) = f a
docase (IE_done a stream) = f a >>== (\r -> case r of
IE_done x _ -> liftI $ IE_done x stream
IE_cont k -> k stream)
docase (IE_cont k) = liftI $ IE_cont ((>>= f) . k)
instance MonadTrans (IterateeGM el) where
lift m = IM (m >>= unIM . return)
stream2list :: Monad m => IterateeGM el m [el]
stream2list = liftI $ IE_cont (step [])
where
step acc (Chunk []) = liftI $ IE_cont (step acc)
step acc (Chunk ls) = liftI $ IE_cont (step $ acc ++ ls)
step acc stream = liftI $ IE_done acc stream
iter_report_err :: Monad m => IterateeGM el m (Maybe String)
iter_report_err = liftI $ IE_cont step
where step s@(Err str) = liftI $ IE_done (Just str) s
step s = liftI $ IE_done Nothing s
sbreak :: Monad m => (el -> Bool) -> IterateeGM el m ([el],Maybe el)
sbreak cpred = liftI $ IE_cont (liftI . step [])
where
step before (Chunk []) = IE_cont (liftI . step before)
step before (Chunk str) =
case break cpred str of
(_,[]) -> IE_cont (liftI . step (before ++ str))
(str,c:tail) -> done (before ++ str) (Just c) (Chunk tail)
step before stream = done before Nothing stream
done line char stream = IE_done (line,char) stream
sdropWhile :: Monad m => (el -> Bool) -> IterateeGM el m ()
sdropWhile cpred = liftI $ IE_cont step
where
step (Chunk []) = sdropWhile cpred
step (Chunk str) =
case dropWhile cpred str of
[] -> sdropWhile cpred
str -> liftI $ IE_done () (Chunk str)
step stream = liftI $ IE_done () stream
snext :: Monad m => IterateeGM el m (Maybe el)
snext = liftI $ IE_cont step
where
step (Chunk []) = snext
step (Chunk (c:t)) = liftI $ IE_done (Just c) (Chunk t)
step stream = liftI $ IE_done Nothing stream
speek :: Monad m => IterateeGM el m (Maybe el)
speek = liftI $ IE_cont step
where
step (Chunk []) = speek
step s@(Chunk (c:_)) = liftI $ IE_done (Just c) s
step stream = liftI $ IE_done Nothing stream
skip_till_eof :: Monad m => IterateeGM el m ()
skip_till_eof = liftI $ IE_cont step
where
step (Chunk _) = skip_till_eof
step _ = return ()
sdrop :: Monad m => Int -> IterateeGM el m ()
sdrop 0 = return ()
sdrop n = liftI $ IE_cont step
where
step (Chunk str) | length str <= n = sdrop (n length str)
step (Chunk str) = liftI $ IE_done () (Chunk s2)
where (s1,s2) = splitAt n str
step stream = liftI $ IE_done () stream
type EnumeratorN el_outer el_inner m a =
IterateeG el_inner m a -> IterateeGM el_outer m (IterateeG el_inner m a)
stake :: Monad m => Int -> EnumeratorN el el m a
stake 0 iter = return iter
stake n iter@IE_done{} = sdrop n >> return iter
stake n (IE_cont k) = liftI $ IE_cont step
where
step (Chunk []) = liftI $ IE_cont step
step chunk@(Chunk str) | length str <= n =
stake (n length str) ==<< k chunk
step (Chunk str) = done (Chunk s1) (Chunk s2)
where (s1,s2) = splitAt n str
step stream = done stream stream
done s1 s2 = k s1 >>== \r -> liftI $ IE_done r s2
map_stream :: Monad m => (el -> el') -> EnumeratorN el el' m a
map_stream f iter@IE_done{} = return iter
map_stream f (IE_cont k) = liftI $ IE_cont step
where
step (Chunk []) = liftI $ IE_cont step
step (Chunk str) = k (Chunk (map f str)) >>== map_stream f
step EOF = k EOF >>== \r -> liftI $ IE_done r EOF
step (Err err) = k (Err err) >>== \r -> liftI $ IE_done r (Err err)
conv_stream :: Monad m =>
IterateeGM el m (Maybe [el']) -> EnumeratorN el el' m a
conv_stream fi iter@IE_done{} = return iter
conv_stream fi (IE_cont k) =
fi >>= (conv_stream fi ==<<) . k . maybe (Err "conv: stream error") Chunk
type Line = String
line :: Monad m => IterateeM m (Either Line Line)
line = sbreak (\c -> c == '\r' || c == '\n') >>= check_next
where
check_next (line,Just '\r') = speek >>= \c ->
case c of
Just '\n' -> snext >> return (Right line)
Just _ -> return (Right line)
Nothing -> return (Left line)
check_next (line,Just _) = return (Right line)
check_next (line,Nothing) = return (Left line)
print_lines :: IterateeGM Line IO ()
print_lines = liftI $ IE_cont step
where
step (Chunk []) = print_lines
step (Chunk ls) = lift (mapM_ pr_line ls) >> print_lines
step EOF = lift (putStrLn ">> natural end") >> liftI (IE_done () EOF)
step stream = lift (putStrLn ">> unnatural end") >>
liftI (IE_done () stream)
pr_line line = putStrLn $ ">> read line: " ++ line
enum_lines :: Monad m => EnumeratorN Char Line m a
enum_lines iter@IE_done{} = return iter
enum_lines (IE_cont k) = line >>= check_line k
where
check_line k (Right "") = enum_lines ==<< k EOF
check_line k (Right l) = enum_lines ==<< k (Chunk [l])
check_line k _ = enum_lines ==<< k (Err "EOF")
enum_words :: Monad m => EnumeratorN Char String m a
enum_words iter@IE_done{} = return iter
enum_words (IE_cont k) = sdropWhile isSpace >> sbreak isSpace >>= check_word k
where
check_word k ("",_) = enum_words ==<< k EOF
check_word k (str,_) = enum_words ==<< k (Chunk [str])
type EnumeratorGM el m a = IterateeG el m a -> IterateeGM el m a
type EnumeratorM m a = EnumeratorGM Char m a
enum_eof :: Monad m => EnumeratorGM el m a
enum_eof iter@(IE_done _ (Err _)) = liftI iter
enum_eof (IE_done x _) = liftI $ IE_done x EOF
enum_eof (IE_cont k) = k EOF
enum_err :: Monad m => String -> EnumeratorGM el m a
enum_err str iter@(IE_done _ (Err _)) = liftI iter
enum_err str (IE_done x _) = liftI $ IE_done x (Err str)
enum_err str (IE_cont k) = k (Err str)
(>.):: Monad m =>
EnumeratorGM el m a -> EnumeratorGM el m a -> EnumeratorGM el m a
e1 >. e2 = (e2 ==<<) . e1
enum_pure_1chunk :: Monad m => [el] -> EnumeratorGM el m a
enum_pure_1chunk str iter@IE_done{} = liftI $ iter
enum_pure_1chunk str (IE_cont k) = k (Chunk str)
enum_pure_nchunk :: Monad m => [el] -> Int -> EnumeratorGM el m a
enum_pure_nchunk str n iter@IE_done{} = liftI $ iter
enum_pure_nchunk [] n iter = liftI $ iter
enum_pure_nchunk str n (IE_cont k) = enum_pure_nchunk s2 n ==<< k (Chunk s1)
where (s1,s2) = splitAt n str
enum_fd :: Fd -> EnumeratorM IO a
enum_fd fd iter = IM $ allocaBytes (fromIntegral buffer_size) (loop iter)
where
buffer_size = 5
loop iter@IE_done{} p = return iter
loop iter@(IE_cont step) p = do
n <- myfdRead fd p buffer_size
putStrLn $ "Read buffer, size " ++ either (const "IO err") show n
case n of
Left errno -> unIM $ step (Err "IO error")
Right 0 -> return iter
Right n -> do
str <- peekCAStringLen (p,fromIntegral n)
im <- unIM $ step (Chunk str)
loop im p
enum_file :: FilePath -> EnumeratorM IO a
enum_file filepath iter = IM $ do
putStrLn $ "opened file " ++ filepath
fd <- openFd filepath ReadOnly Nothing defaultFileFlags
r <- unIM $ enum_fd fd iter
closeFd fd
putStrLn $ "closed file " ++ filepath
return r
enum_chunk_decoded :: Monad m => Iteratee m a -> IterateeM m a
enum_chunk_decoded = docase
where
docase iter@IE_done{} =
liftI iter >>= (\r -> (enum_chunk_decoded ==<< skip_till_eof) >> return r)
docase iter@(IE_cont k) = line >>= check_size
where
check_size (Right "0") = line >> k EOF
check_size (Right str) =
maybe (k . Err $ "Bad chunk size: " ++ str) (read_chunk iter)
$ read_hex 0 str
check_size _ = k (Err "Error reading chunk size")
read_chunk iter size =
do
r <- stake size iter
c1 <- snext
c2 <- snext
case (c1,c2) of
(Just '\r',Just '\n') -> docase r
_ -> (enum_chunk_decoded ==<< skip_till_eof) >>
enum_err "Bad chunk trailer" r
read_hex acc "" = Just acc
read_hex acc (d:rest) | isHexDigit d = read_hex (16*acc + digitToInt d) rest
read_hex acc _ = Nothing
test_str1 =
"header1: v1\rheader2: v2\r\nheader3: v3\nheader4: v4\n" ++
"header5: v5\r\nheader6: v6\r\nheader7: v7\r\n\nrest\n"
testp1 =
let IE_done (IE_done lines EOF) (Chunk rest)
= runIdentity . unIM $ enum_pure_1chunk test_str1 ==<<
(enum_lines ==<< stream2list)
in
lines == ["header1: v1","header2: v2","header3: v3","header4: v4",
"header5: v5","header6: v6","header7: v7"]
&& rest == "rest\n"
testp2 =
let IE_done (IE_done lines EOF) (Chunk rest)
= runIdentity . unIM $ enum_pure_nchunk test_str1 5 ==<<
(enum_lines ==<< stream2list)
in
lines == ["header1: v1","header2: v2","header3: v3","header4: v4",
"header5: v5","header6: v6","header7: v7"]
&& rest == "r"
testw1 =
let test_str = "header1: v1\rheader2: v2\r\nheader3:\t v3"
expected = ["header1:","v1","header2:","v2","header3:","v3"] in
let run_test test_str =
let IE_done (IE_done words EOF) EOF
= runIdentity . unIM $ (enum_pure_nchunk test_str 5 >. enum_eof)
==<< (enum_words ==<< stream2list)
in words
in
and [run_test test_str == expected,
run_test (test_str ++ " ") == expected]
test_driver line_collector filepath = do
fd <- openFd filepath ReadOnly Nothing defaultFileFlags
putStrLn "About to read headers"
result <- unIM $ (enum_fd fd >. enum_eof) ==<< read_lines_and_one_more_line
closeFd fd
putStrLn "Finished reading headers"
case result of
IE_done (IE_done headers EOF,after) _ ->
do
putStrLn $ "The line after headers is: " ++ show after
putStrLn "Complete headers"
print headers
IE_done (IE_done headers err,_) stream ->
do
putStrLn $ "Problem " ++ show stream
putStrLn "Incomplete headers"
print headers
where
read_lines_and_one_more_line = do
lines <- enum_lines ==<< line_collector
after <- line
return (lines,after)
test11 = test_driver stream2list "test1.txt"
test12 = test_driver stream2list "test2.txt"
test13 = test_driver stream2list "test3.txt"
test14 = test_driver stream2list "/dev/null"
test21 = test_driver print_lines "test1.txt"
test22 = test_driver print_lines "test2.txt"
test23 = test_driver print_lines "test3.txt"
test24 = test_driver print_lines "/dev/null"
line_printer = enum_lines ==<< print_lines
read_headers_print_body = do
headers <- enum_lines ==<< stream2list
case headers of
IE_done headers EOF -> lift $ do
putStrLn "Complete headers"
print headers
IE_done headers (Err err) -> lift $ do
putStrLn $ "Incomplete headers due to " ++ err
print headers
lift $ putStrLn "\nLines of the body follow"
enum_chunk_decoded ==<< line_printer
print_headers_print_body = do
lift $ putStrLn "\nLines of the headers follow"
line_printer
lift $ putStrLn "\nLines of the body follow"
enum_chunk_decoded ==<< line_printer
test_driver_full iter filepath = do
fd <- openFd filepath ReadOnly Nothing defaultFileFlags
putStrLn "About to read headers"
unIM $ (enum_fd fd >. enum_eof) ==<< iter
closeFd fd
putStrLn "Finished reading"
test31 = test_driver_full read_headers_print_body "test_full1.txt"
test32 = test_driver_full read_headers_print_body "test_full2.txt"
test33 = test_driver_full read_headers_print_body "test_full3.txt"
test34 = test_driver_full print_headers_print_body "test_full3.txt"
test_driver_mux iter fpath1 fpath2 = do
fd1 <- openFd fpath1 ReadOnly Nothing defaultFileFlags
fd2 <- openFd fpath2 ReadOnly Nothing defaultFileFlags
let fds = [fd1,fd2]
putStrLn $ "Opened file descriptors: " ++ show fds
mapM (\(fd,reader) -> unIM reader >>= return . ((,) fd))
(zip fds (repeat iter)) >>=
allocaBytes (fromIntegral buffer_size) . loop
mapM_ closeFd fds
putStrLn $ "Closed file descriptors. All done"
where
buffer_size = 5
loop fjque buf = do
let fds = get_fds fjque
if null fds then return ()
else do
selected <- select'read'pending fds
case selected of
Left errno -> putStrLn "IO Err" >>
tell_iteratee_err "IO Err" fjque >>
return ()
Right [] -> loop fjque buf
Right sel -> process buf sel fjque
get_fds = foldr (\ (fd,iter) acc ->
case iter of {IE_cont _ -> fd:acc; _ -> acc}) []
get_ready selected jq = (e, before ++ after)
where (before,e:after) = break (\(fd,_) -> fd `elem` selected) jq
process buf selected fjque = do
let ((fd,IE_cont step),fjrest) = get_ready selected fjque
n <- myfdRead fd buf buffer_size
putStrLn $ unwords ["Read buffer, size", either (const "IO err") show n,
"from fd", show fd]
case n of
Left errno -> unIM (step (Err "IO error")) >>
loop fjrest buf
Right 0 -> unIM (step EOF) >>
loop fjrest buf
Right n -> do
str <- peekCAStringLen (buf,fromIntegral n)
im <- unIM $ step (Chunk str)
loop (fjrest ++ [(fd,im)]) buf
tell_iteratee_err err = mapM_ (\ (_,iter) -> unIM (enum_err err iter))
testm1 = test_driver_mux line_printer "test1.txt" "test3.txt"
testm2 = test_driver_mux print_headers_print_body
"test_full2.txt" "test_full3.txt"