IterateeIO.hs

module IterateeIO where

import IterateeM 
import System.FilePath
import Data.Maybe
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Lazy.Char8 as L8


readLine :: Monad m => Iteratee e L.ByteString m (Maybe L.ByteString)
readLine = continue (go L.empty)
    where
      go :: Monad m => L.ByteString -> Stream L.ByteString -> Iteratee e L.ByteString m (Maybe L.ByteString)
      go acc EOF = yield Nothing (Chunks [acc])
      go acc (Chunks bs) | not (L.null b) = yield (Just acca) (Chunks [btail])
                         | otherwise      = yield Nothing (Chunks [acca])
                         where
                           (a, b) = L8.break (== '\n') $ L.concat bs
                           acca   = L.append acc a
                           btail  = L.tail b


enumerateFile :: FilePath -> Enumerator e L.ByteString IO L.ByteString
enumerateFile path = 
    bracket (openFile path ReadMode) hClose $ \h -> 
        let go (

IterateeM.hs

{-# LANGUAGE UnicodeSyntax #-}
module IterateeM where

    import Data.Monoid
    import Control.Monad
    import Control.Monad.Trans

    -- | Stream of data either in chunks or EOF if at the end of input  
    data Stream a = Chunks [a] | EOF
                    deriving Show

    instance Monoid (Stream a) where

                 mempty = Chunks []

                 (Chunks xs) `mappend` (Chunks ys) = Chunks (xs ++ ys)
                 (mappend)  _ _ = EOF


    instance Functor Stream where
                 fmap _ EOF = EOF
                 fmap f (Chunks [a]) = Chunks $ f `fmap` [a]

                 
    -- | Step through the computation by consuming chunk of input at a
    -- time.   
    data Step e a m b = Continue (Stream a -> Iteratee e a m b)
                      | Yield b (Stream a)
                      | Error e

    instance (Show e, Show b, Show a) ⇒ Show (Step e a m b)
        where

          showsPrec d step = showParen (d > 10) $
                             case step of
                               (Continue _ ) -> s "Continue"
                               (Yield b _)   -> s "Yield"
                               (Error e)     -> s "Error"

              where
                s = showString

    -- | Iteratee represents a computation that produces a step
    newtype Iteratee e a m b = Iteratee
        {
          runIteratee :: m ( Step e a m b )
        }

    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 b chunks = returnI (Yield b chunks)

    continue :: Monad m => (Stream a -> Iteratee e a m b) -> Iteratee e a m b
    continue = returnI . Continue


    instance Monad m ⇒ Monad (Iteratee e a m) where

                 return = flip yield (Chunks [])

                 i >>= f = Iteratee $ do
                             s <- runIteratee i
                             case s of
                               Continue k -> return $ Continue ((>>= f) . k)
                               Error e    -> return $ Error e
                               Yield b (Chunks []) -> runIteratee $ f b
                               Yield b chunks      -> do
                                                s' <- runIteratee $ f b
                                                case s' of
                                                  Continue k' -> runIteratee $ k' chunks
                                                  Error e'    -> return $ Error e'
                                                  Yield b' _  -> return $ Yield b' chunks
        
    instance MonadTrans (Iteratee e a) where

                 lift m = Iteratee $ m >>= runIteratee . return 

    instance MonadIO m => MonadIO (Iteratee e a m) where

                 liftIO = lift . liftIO


    instance Monad m => Functor (Iteratee e a m) where

                 fmap f i = i >>= return . f 



    type Enumerator e a m b = Step e a m b -> Iteratee e a m b

    isum :: (Num a, Monad m) => a -> Stream a -> Iteratee e a m a
    isum csum EOF         = yield csum EOF
    isum csum (Chunks xs) = returnI $ Continue $ isum (csum + sum xs)

    esum :: (Num a, Monad m) => Iteratee e a m a
    esum = returnI $ Continue (isum 0)


    run :: (Show e, Monad m) => Iteratee e a m b -> m (Either String b)
    run i = do
      s <- runIteratee i
      case s of
        Continue k -> do
                   let Iteratee i' = k EOF
                   s' <- i'
                   case s' of
                     Continue _ -> return $ Left "misbehaving iteratee"
                     Yield b _  -> return $ Right b
                     Error e    -> return $ Left (show e)
        Yield b _ -> return $ Right b
        Error e   -> return $ Left (show e)
                     
    (>>==) :: 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 $ do
      s <- runIteratee i
      runIteratee $ f s

    (==<<) :: Monad m => (Step e a m b -> Iteratee e a' m b') -> Iteratee e a m b -> Iteratee e a' m b'
    (==<<) = flip (>>==)

    enumGen :: Monad m => [a] -> Enumerator e a m a
    enumGen xs (Continue k) = loop xs 
        where loop (x:xs)   = (k (Chunks [x])) >>== enumGen xs 
              loop _        = k (Chunks [])
    enumGen _  (Error e)    = returnI $ Error e
    enumGen _  (Yield b _)  = yield b (Chunks [])

    
    type Enumeratee e aOut aIn m b = Step e aIn m b -> Iteratee e aOut m (Step e aIn m b)

    drive :: Monad m => Iteratee e aIn m b -> Enumeratee e aOut aIn m b -> Iteratee e aOut m (Step e aIn m b)
    iter `drive` enum = iter >>== enum

    senumeratee :: Monad m => Enumeratee e aOut aIn m b
    senumeratee (Error e)     = returnI $ Error e
    senumeratee s@(Yield b _) = yield s (Chunks [])
    senumeratee s@(Continue k)  = undefined