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 qualified Control.Exception as E
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 Prelude hiding (span)
import qualified Prelude as Prelude
data Stream a
= Chunks [a]
| EOF
deriving (Show, Eq)
data Step a m b
= Continue (Stream a -> Iteratee a m b)
| Yield b (Stream a)
| Error E.SomeException
newtype Iteratee a m b = Iteratee
{ runIteratee :: m (Step a m b)
}
type Enumerator a m b = Step a m b -> Iteratee a m b
type Enumeratee aOut aIn m b = Step aIn m b -> Iteratee aOut m (Step 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 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 a m) where
fmap = liftM
instance Monad m => A.Applicative (Iteratee a m) where
pure = return
(<*>) = ap
instance MT.MonadTrans (Iteratee a) where
lift m = Iteratee $ m >>= runIteratee . return
instance MIO.MonadIO m => MIO.MonadIO (Iteratee a m) where
liftIO = MT.lift . MIO.liftIO
liftTrans :: (Monad m, MT.MonadTrans t, Monad (t m)) =>
Iteratee a m b -> Iteratee 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 a m b -> Iteratee a m b
returnI = Iteratee . return
yield :: Monad m => b -> Stream a -> Iteratee a m b
yield x chunk = returnI (Yield x chunk)
continue :: Monad m => (Stream a -> Iteratee a m b) -> Iteratee a m b
continue = returnI . Continue
throwError :: (Monad m, E.Exception e) => e -> Iteratee a m b
throwError = returnI . Error . E.SomeException
liftI :: Monad m => (Stream a -> Step a m b) -> Iteratee a m b
liftI k = continue $ returnI . k
catchError :: Monad m => Iteratee a m b -> (E.SomeException -> Iteratee a m b) -> Iteratee 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 a m b ->
(Step a m b -> Iteratee a' m b') ->
Iteratee a' m b'
i >>== f = Iteratee $ runIteratee i >>= runIteratee . f
infixr 1 ==<<
(==<<):: Monad m =>
(Step a m b -> Iteratee a' m b') ->
Iteratee a m b ->
Iteratee a' m b'
(==<<) = flip (>>==)
infixr 0 $$
($$):: Monad m =>
(Step a m b -> Iteratee a' m b') ->
Iteratee a m b ->
Iteratee a' m b'
($$) = (==<<)
infixr 1 >==>
(>==>) :: Monad m =>
Enumerator a m b ->
(Step a m b -> Iteratee a' m b') ->
Step a m b ->
Iteratee a' m b'
(>==>) e1 e2 s = e1 s >>== e2
infixr 1 <==<
(<==<) :: Monad m =>
(Step a m b -> Iteratee a' m b') ->
Enumerator a m b ->
Step a m b ->
Iteratee a' m b'
(<==<) = flip (>==>)
consume :: Monad m => Iteratee 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 a m Bool
isEOF = liftI $ \c -> case c of
EOF -> Yield True c
_ -> Yield False c
liftFoldL :: Monad m => (b -> a -> b) -> b -> Iteratee 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 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 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 a m b -> m (Either E.SomeException 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 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 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 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 a m b] -> Enumerator a m b
concatEnums = foldl (>==>) returnI
joinI :: Monad m => Iteratee a m (Step a' m b) -> Iteratee 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 a m b) -> Iteratee a' m (Step a m b)) ->
Enumeratee 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 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 ao m ai -> Enumeratee 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 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 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 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 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 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 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 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 a m [a]
break p = span $ not . p