enumerator-tutorial-part-3.hs

{-# LANGUAGE ScopedTypeVariables #-}
import Prelude hiding (head, map)
import Data.Enumerator
import Control.Monad.IO.Class
import Control.Monad.Trans.Class

sumIter :: Monad m => Iteratee Int m Int
sumIter = do
    maybeNum <- head
    case maybeNum of
        Nothing -> return 0
        Just i -> do
            rest <- sumIter
            return $ i + rest

lineEnum :: MonadIO m => Enumerator String m b
lineEnum (Continue k) = do
    x <- liftIO getLine
    if x == "q"
        then continue k
        else k (Chunks [x]) >>== lineEnum
lineEnum step = returnI step

sumIterString :: Monad m => Iteratee String m Int
sumIterString = Iteratee $ do
    innerStep <- runIteratee sumIter
    return $ go innerStep
  where
    go :: Monad m => Step Int m Int -> Step String m Int
    go (Yield res _) = Yield res EOF
    go (Error err) = Error err
    go (Continue k) = Continue $ \strings -> Iteratee $ do
        let ints = fmap read strings :: Stream Int
        step <- runIteratee $ k ints
        return $ go step

mapIter :: Monad m => (aOut -> aIn) -> Iteratee aIn m b -> Iteratee aOut m b
mapIter f innerIter = Iteratee $ do
    innerStep <- runIteratee innerIter
    return $ go innerStep
  where
    go (Yield res _) = Yield res EOF
    go (Error err) = Error err
    go (Continue k) = Continue $ \strings -> Iteratee $ do
        let ints = fmap f strings
        step <- runIteratee $ k ints
        return $ go step

unnest :: Monad m => Iteratee String m (Step Int m Int) -> Iteratee String m Int
unnest outer = do -- using the Monad instance of Iteratee
    inner <- outer -- inner :: Step Int m Int
    go inner
  where
    go (Error e) = throwError e
    go (Yield x _) = yield x EOF
    go (Continue k) = k EOF >>== go

skip :: Monad m => Enumeratee a a m b
skip (Continue k) = do
    x <- head
    _ <- head -- the one we're skipping
    case x of
        Nothing -> return $ Continue k
        Just y -> do
            newStep <- lift $ runIteratee $ k $ Chunks [y]
            skip newStep
skip step = return step

main = do
    run_ (testEnum $$ sumIterString) >>= print
    run_ (testEnum $$ mapIter read sumIter) >>= print
    run_ (testEnum $$ unnest $ map read $$ sumIter) >>= print
    run_ (testEnum $$ joinI $ map read $$ sumIter) >>= print
    run_ (testEnum $$ joinI $ skip $$ joinI $ map read $$ sumIter) >>= print
  where
    testEnum = enumList 1 $ words "1 2 3 4 5"