module Data.Enumerator (
Stream (..)
, Step (..)
, Iteratee (..)
, Enumerator
, Enumeratee
, returnI
, yield
, continue
, throwError
, catchError
, liftI
, (>>==)
, (==<<)
, ($$)
, (>==>)
, (<==<)
, run
, consume
, isEOF
, liftTrans
, liftFoldL
, liftFoldL'
, liftFoldM
, printChunks
, enumEOF
, enumList
, concatEnums
, checkDone
, Data.Enumerator.map
, Data.Enumerator.sequence
, joinI
, Data.Enumerator.head
, peek
, Data.Enumerator.last
, Data.Enumerator.length
, Data.Enumerator.drop
, Data.Enumerator.dropWhile
, span
, Data.Enumerator.break
) where
import Data.List (genericDrop, genericLength, genericSplitAt)
import Data.Monoid (Monoid, mempty, mappend, mconcat)
import qualified Control.Applicative as A
import Control.Monad (liftM, ap)
import qualified Control.Monad.IO.Class as MIO
import qualified Control.Monad.Trans.Class as MT
import qualified Data.List as DataList
import Control.Monad (foldM)
import qualified Control.Exception as E
import Prelude hiding (span)
import qualified Prelude as Prelude
data Stream a
= Chunks [a]
| EOF
deriving (Show, Eq)
data Step e a m b
= Continue (Stream a -> Iteratee e a m b)
| Yield b (Stream a)
| Error e
newtype Iteratee e a m b = Iteratee
{ runIteratee :: m (Step e a m b)
}
type Enumerator e a m b = Step e a m b -> Iteratee e a m b
type Enumeratee e aOut aIn m b = Step e aIn m b -> Iteratee e aOut m (Step e aIn m b)
instance Monoid (Stream a) where
mempty = Chunks mempty
mappend (Chunks xs) (Chunks ys) = Chunks $ mappend xs ys
mappend _ _ = EOF
instance Functor Stream where
fmap f (Chunks xs) = Chunks $ fmap f xs
fmap _ EOF = EOF
instance Monad Stream where
return = Chunks . return
Chunks xs >>= f = mconcat $ fmap f xs
EOF >>= _ = EOF
instance Monad m => Monad (Iteratee e a m) where
return x = Iteratee . return $ Yield x $ Chunks []
m >>= f = Iteratee $ runIteratee m >>=
\r1 -> case r1 of
Continue k -> return $ Continue ((>>= f) . k)
Error err -> return $ Error err
Yield x (Chunks []) -> runIteratee $ f x
Yield x chunk -> runIteratee (f x) >>=
\r2 -> case r2 of
Continue k -> runIteratee $ k chunk
Error err -> return $ Error err
Yield x' _ -> return $ Yield x' chunk
instance Monad m => Functor (Iteratee e a m) where
fmap = liftM
instance Monad m => A.Applicative (Iteratee e a m) where
pure = return
(<*>) = ap
instance MT.MonadTrans (Iteratee e a) where
lift m = Iteratee $ m >>= runIteratee . return
instance MIO.MonadIO m => MIO.MonadIO (Iteratee e a m) where
liftIO = MT.lift . MIO.liftIO
liftTrans :: (Monad m, MT.MonadTrans t, Monad (t m)) =>
Iteratee e a m b -> Iteratee e a (t m) b
liftTrans iter = Iteratee $ do
step <- MT.lift $ runIteratee iter
return $ case step of
Yield x cs -> Yield x cs
Error err -> Error err
Continue k -> Continue (liftTrans . k)
returnI :: Monad m => Step e a m b -> Iteratee e a m b
returnI = Iteratee . return
yield :: Monad m => b -> Stream a -> Iteratee e a m b
yield x chunk = returnI (Yield x chunk)
continue :: Monad m => (Stream a -> Iteratee e a m b) -> Iteratee e a m b
continue = returnI . Continue
throwError :: Monad m => e -> Iteratee e a m b
throwError = returnI . Error
liftI :: Monad m => (Stream a -> Step e a m b) -> Iteratee e a m b
liftI k = continue $ returnI . k
catchError :: Monad m => Iteratee e a m b -> (e -> Iteratee e a m b) -> Iteratee e a m b
catchError iter h = Iteratee $ do
step <- runIteratee iter
case step of
Error err -> runIteratee (h err)
_ -> return step
infixl 1 >>==
(>>==) :: Monad m =>
Iteratee e a m b ->
(Step e a m b -> Iteratee e a' m b') ->
Iteratee e a' m b'
i >>== f = Iteratee $ runIteratee i >>= runIteratee . f
infixr 1 ==<<
(==<<):: Monad m =>
(Step e a m b -> Iteratee e a' m b') ->
Iteratee e a m b ->
Iteratee e a' m b'
(==<<) = flip (>>==)
infixr 0 $$
($$):: Monad m =>
(Step e a m b -> Iteratee e a' m b') ->
Iteratee e a m b ->
Iteratee e a' m b'
($$) = (==<<)
infixr 1 >==>
(>==>) :: Monad m =>
Enumerator e a m b ->
(Step e a m b -> Iteratee e a' m b') ->
Step e a m b ->
Iteratee e a' m b'
(>==>) e1 e2 s = e1 s >>== e2
infixr 1 <==<
(<==<) :: Monad m =>
(Step e a m b -> Iteratee e a' m b') ->
Enumerator e a m b ->
Step e a m b ->
Iteratee e a' m b'
(<==<) = flip (>==>)
consume :: Monad m => Iteratee e a m [a]
consume = liftI $ step id where
step acc chunk = case chunk of
Chunks [] -> Continue $ returnI . step acc
Chunks xs -> Continue $ returnI . (step $ acc . (xs ++))
EOF -> Yield (acc []) EOF
isEOF :: Monad m => Iteratee e a m Bool
isEOF = liftI $ \c -> case c of
EOF -> Yield True c
_ -> Yield False c
liftFoldL :: Monad m => (b -> a -> b) -> b -> Iteratee e a m b
liftFoldL f = liftI . step where
step acc chunk = case chunk of
Chunks [] -> Continue $ returnI . step acc
Chunks xs -> Continue $ returnI . step (Prelude.foldl f acc xs)
EOF -> Yield acc EOF
liftFoldL' :: Monad m => (b -> a -> b) -> b -> Iteratee e a m b
liftFoldL' f = liftI . step where
fold = DataList.foldl' f
step acc chunk = case chunk of
Chunks [] -> Continue $ returnI . step acc
Chunks xs -> Continue $ returnI . (step $! fold acc xs)
EOF -> Yield acc EOF
liftFoldM :: Monad m => (b -> a -> m b) -> b -> Iteratee e a m b
liftFoldM f = continue . step where
step acc chunk = case chunk of
Chunks [] -> continue $ step acc
Chunks xs -> Iteratee $ liftM (Continue . step) (foldM f acc xs)
EOF -> yield acc EOF
run :: Monad m => Iteratee e a m b -> m (Either e b)
run i = do
mStep <- runIteratee $ enumEOF ==<< i
case mStep of
Error err -> return $ Left err
Yield x _ -> return $ Right x
Continue _ -> error "run: divergent iteratee"
printChunks :: (MIO.MonadIO m, Show a) => Bool -> Iteratee e a m ()
printChunks printEmpty = continue step where
step (Chunks []) | not printEmpty = continue step
step (Chunks xs) = MIO.liftIO (print xs) >> continue step
step EOF = MIO.liftIO (putStrLn "EOF") >> yield () EOF
enumEOF :: Monad m => Enumerator e a m b
enumEOF (Yield x _) = yield x EOF
enumEOF (Error err) = throwError err
enumEOF (Continue k) = k EOF >>== check where
check (Continue _) = error "enumEOF: divergent iteratee"
check s = enumEOF s
enumList :: Monad m => Integer -> [a] -> Enumerator e a m b
enumList n xs (Continue k) | not (null xs) = k chunk >>== loop where
(s1, s2) = genericSplitAt n xs
chunk = Chunks s1
loop = enumList n s2
enumList _ _ step = returnI step
concatEnums :: Monad m => [Enumerator e a m b] -> Enumerator e a m b
concatEnums = foldl (>==>) returnI
joinI :: Monad m => Iteratee e a m (Step e a' m b) -> Iteratee e a m b
joinI outer = outer >>= check where
check (Continue k) = k EOF >>== \s -> case s of
Continue _ -> error "joinI: divergent iteratee"
_ -> check s
check (Yield x _) = return x
check (Error e) = throwError e
checkDone :: Monad m =>
((Stream a -> Iteratee e a m b) -> Iteratee e a' m (Step e a m b)) ->
Enumeratee e a' a m b
checkDone _ (Yield x chunk) = return $ Yield x chunk
checkDone f (Continue k) = f k
checkDone _ (Error err) = throwError err
map :: Monad m => (ao -> ai) -> Enumeratee e ao ai m b
map f = loop where
loop = checkDone $ continue . step
step k EOF = yield (Continue k) EOF
step k (Chunks []) = continue $ step k
step k (Chunks xs) = k (Chunks (Prelude.map f xs)) >>== loop
sequence :: Monad m => Iteratee e ao m ai -> Enumeratee e ao ai m b
sequence i = loop where
loop = checkDone check
check k = isEOF >>= \f -> if f
then yield (Continue k) EOF
else step k
step k = i >>= \v -> k (Chunks [v]) >>== loop
head :: Monad m => Iteratee e a m (Maybe a)
head = liftI step where
step (Chunks []) = Continue $ returnI . step
step (Chunks (x:xs)) = Yield (Just x) (Chunks xs)
step EOF = Yield Nothing EOF
peek :: Monad m => Iteratee e a m (Maybe a)
peek = liftI step where
step (Chunks []) = Continue $ returnI . step
step chunk@(Chunks (x:_)) = Yield (Just x) chunk
step chunk = Yield Nothing chunk
last :: Monad m => Iteratee e a m (Maybe a)
last = liftI $ step Nothing where
step ret (Chunks xs) = let
ret' = case xs of
[] -> ret
_ -> Just $ Prelude.last xs
in Continue $ returnI . step ret'
step ret EOF = Yield ret EOF
length :: Monad m => Iteratee e a m Integer
length = liftI $ step 0 where
step n (Chunks xs) = Continue $ returnI . step (n + genericLength xs)
step n EOF = Yield n EOF
drop :: Monad m => Integer -> Iteratee e a m ()
drop 0 = return ()
drop n = liftI $ step n where
step n' (Chunks xs)
| len xs < n' = Continue $ returnI . step (n' len xs)
| otherwise = Yield () $ Chunks $ genericDrop n' xs
step _ EOF = Yield () EOF
len = genericLength
dropWhile :: Monad m => (a -> Bool) -> Iteratee e a m ()
dropWhile p = liftI step where
step (Chunks xs) = case Prelude.dropWhile p xs of
[] -> Continue $ returnI . step
xs' -> Yield () $ Chunks xs'
step EOF = Yield () EOF
span :: Monad m => (a -> Bool) -> Iteratee e a m [a]
span f = liftI $ step [] where
step acc (Chunks xs) = case Prelude.span f xs of
(_, []) -> Continue $ returnI . step (acc ++ xs)
(head', tail') -> Yield (acc ++ head') (Chunks tail')
step acc EOF = Yield acc EOF
break :: Monad m => (a -> Bool) -> Iteratee e a m [a]
break p = span $ not . p