module Data.Enumerator (
Stream (..)
, Step (..)
, Iteratee (..)
, Enumerator
, Enumeratee
, returnI
, yield
, continue
, throwError
, liftI
, (>>==)
, (==<<)
, consume
, isEOF
, 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
, run
, printChunks
) 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 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
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
(>>==) :: 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
(==<<):: Monad m =>
(Step e a m b -> Iteratee e a' m b') ->
Iteratee e a m b ->
Iteratee e a' m b'
(==<<) = flip (>>==)
consume :: Monad m => Iteratee e a m [a]
consume = liftI $ step [] 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
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 [] s = returnI s
concatEnums (e:es) s = e s >>== concatEnums es
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
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