-----------------------------------------------------------------------------
-- |
-- Module: Data.Enumerator
-- Copyright: 2010 John Millikin
-- License: MIT
--
-- Maintainer: jmillikin@gmail.com
-- Portability: portable
--
-- An implementation of Oleg Kiselyov’s left-fold enumerators
--
-----------------------------------------------------------------------------
module Data.Enumerator (
	  -- * Types
	  Stream (..)
	, Step (..)
	, Iteratee (..)
	, Enumerator
	, Enumeratee
	  -- * Primitives
	  -- ** Combinators
	  -- | These are common patterns which occur whenever iteratees are
	  -- being defined.
	, returnI
	, yield
	, continue
	, throwError
	, catchError
	, liftI
	, (>>==)
	, (==<<)
	, ($$)
	, (>==>)
	, (<==<)
	  -- ** Iteratees
	, run
	, run_
	, consume
	, isEOF
	, liftTrans
	, liftFoldL
	, liftFoldL'
	, liftFoldM
	, printChunks
	  -- ** Enumerators
	, enumEOF
	, enumList
	, concatEnums
	  -- ** Enumeratees
	, checkDone
	, checkDoneEx
	, Data.Enumerator.map
	, Data.Enumerator.concatMap
	, Data.Enumerator.mapM
	, Data.Enumerator.sequence
	, joinI
	  -- * Parser combinators
	  -- | Oleg&#x2019;s original @IterateeM.hs@ includes some basic iteratees
	  -- for parsing, so this section ports them to the new interface. However,
	  -- in practice most parsing will be performed with enumerator-based
	  -- interfaces to existing parser libraries (such as Parsec or Attoparsec).
	, 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
-- | Not to be confused with types from the @Stream@ or
-- @stream-fusion@ packages, a 'Stream' is a sequence of chunks
-- generated by an 'Enumerator'. In contrast to Oleg&#x2019;s implementation,
-- this stream does not support error handling -- errors encountered
-- while generating a stream are reported in the 'Step' type instead.
--
-- @(Chunks [])@ is used to indicate that a stream is still active, but
-- currently has no available data. Iteratees should ignore empty chunks.
data Stream a
	= Chunks [a]
	| EOF
	deriving (Show, Eq)
data Step a m b
	-- | The 'Iteratee' is capable of accepting more input. Note that more input
	-- is not necessarily required; the 'Iteratee' might be able to generate a
	-- value immediately if it receives 'EOF'.
	= Continue (Stream a -> Iteratee a m b)
	
	-- | The 'Iteratee' has received enough input to generate a result.
	-- Included in this value is left-over input, which can be passed to
	-- composed 'Iteratee's.
	| Yield b (Stream a)
	
	-- | The 'Iteratee' encountered an error which prevents it from proceeding
	-- further.
	| Error E.SomeException

-- | The primary data type for this library, which consumes
-- input from a 'Stream' until it either generates a value or encounters
-- an error. Rather than requiring all input at once, an iteratee will
-- return 'Continue' when it is capable of processing more data.
--
-- In general, iteratees begin in the 'Continue' state. As each chunk is
-- passed to the continuation, the iteratee returns the next step:
-- 'Continue' for more data, 'Yield' when it's finished, or 'Error' to
-- abort processing.
newtype Iteratee a m b = Iteratee
	{ runIteratee :: m (Step a m b)
	}
-- | While 'Iteratee's consume data, enumerators generate it. Since
-- @'Iteratee'@ is an alias for @m ('Step' a m b)@, 'Enumerator's can
-- be considered step transformers of type
-- @'Step' a m b -> m ('Step' a m b)@.
--
-- 'Enumerator's typically read from an external source (parser, handle,
-- random generator, etc). They feed chunks into an 'Iteratee' until the
-- source runs out of data (triggering 'EOF') or the iteratee finishes
-- processing ('Yield's a value).
type Enumerator a m b = Step a m b -> Iteratee a m b

-- | In cases where an enumerator acts as both a source and sink, the resulting
-- type is named an 'Enumeratee'. Enumeratees have two input types,
-- &#x201c;outer a&#x201d; (@aOut@) and &#x201c;inner a&#x201d; (@aIn@).
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 []
	{-# INLINE return #-}
	
	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
	{-# INLINE fmap #-}

instance Monad m => A.Applicative (Iteratee a m) where
	pure = return
	{-# INLINE pure #-}
	
	(<*>) = ap
	{-# INLINE (<*>) #-}
instance MT.MonadTrans (Iteratee a) where
	lift m = Iteratee $ m >>= runIteratee . return
	{-# INLINE lift #-}

instance MIO.MonadIO m => MIO.MonadIO (Iteratee a m) where
	liftIO = MT.lift . MIO.liftIO
	{-# INLINE liftIO #-}
-- | Lift an 'Iteratee' onto a monad transformer, re-wrapping the
-- 'Iteratee'&#x2019;s inner monadic values.
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 x = Iteratee (return x)@
returnI :: Monad m => Step a m b -> Iteratee a m b
returnI = Iteratee . return
{-# INLINE returnI #-}

-- | @yield x chunk = returnI (Yield x chunk)@
yield :: Monad m => b -> Stream a -> Iteratee a m b
yield x chunk = returnI (Yield x chunk)
{-# INLINE yield #-}

-- | @continue k = returnI (Continue k)@
continue :: Monad m => (Stream a -> Iteratee a m b) -> Iteratee a m b
continue = returnI . Continue
{-# INLINE continue #-}

-- | @throwError err = returnI (Error err)@
throwError :: (Monad m, E.Exception e) => e -> Iteratee a m b
throwError = returnI . Error . E.toException
{-# INLINE throwError #-}

-- | @liftI f = continue (returnI . f)@
liftI :: Monad m => (Stream a -> Step a m b) -> Iteratee a m b
liftI k = continue $ returnI . k
{-# INLINE liftI #-}
catchError :: Monad m => Iteratee a m b -> (E.SomeException -> Iteratee a m b) -> Iteratee a m b
catchError iter h = iter >>== step where
	step (Yield b as) = yield b as
	step (Error err) = h err
	step (Continue k) = continue (\stream -> k stream >>== step)
infixl 1 >>==

-- | Equivalent to (>>=), but allows 'Iteratee's with different input types
-- to be composed.
(>>==) :: 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
{-# INLINE (>>==) #-}
infixr 1 ==<<

-- | @(==\<\<) = flip (\>\>==)@
(==<<):: Monad m =>
	(Step a m b -> Iteratee a' m b') ->
	Iteratee a m b ->
	Iteratee a' m b'
(==<<) = flip (>>==)
{-# INLINE (==<<) #-}
infixr 0 $$

-- | @($$) = (==\<\<)@
--
-- This might be easier to read when passing a chain of iteratees to an
-- enumerator.
($$):: Monad m =>
	(Step a m b -> Iteratee a' m b') ->
	Iteratee a m b ->
	Iteratee a' m b'
($$) = (==<<)
{-# INLINE ($$) #-}
infixr 1 >==>

-- | @(>==>) e1 e2 s = e1 s >>== e2@
(>==>) :: 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
{-# INLINE (>==>) #-}
infixr 1 <==<

-- | @(\<==\<) = flip (>==>)@
(<==<) :: Monad m =>
	(Step a m b -> Iteratee a' m b') ->
	Enumerator a m b ->
	Step a m b ->
	Iteratee a' m b'
(<==<) = flip (>==>)
{-# INLINE (<==<) #-}
-- | Consume all input until 'EOF', then return consumed input as a list.
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
-- | Return 'True' if the next 'Stream' is 'EOF'.
isEOF :: Monad m => Iteratee a m Bool
isEOF = liftI $ \c -> case c of
	EOF -> Yield True c
	_   -> Yield False c
-- | Lifts a pure left fold into an iteratee.
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
-- | As 'liftFoldL', but strict in its accumulator.
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
-- | Lifts a monadic left fold into an iteratee.
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 an iteratee until it finishes, and return either the final value
-- (if it succeeded) or the error (if it failed).
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"
run_ :: Monad m => Iteratee a m b -> m b
run_ i = run i >>= either E.throw return
-- | Print chunks as they're received from the enumerator, optionally
-- printing empty chunks.
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
-- | The most primitive enumerator; simply sends 'EOF'. The iteratee must
-- either yield a value or throw an error continuing receiving 'EOF' will
-- not terminate with any useful value.
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
-- | Another small, useful enumerator separates an input list into chunks,
-- and sends them to the iteratee. This is useful for testing iteratees in pure
-- code.
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
-- | Compose a list of 'Enumerator's using '(>>==)'
concatEnums :: Monad m => [Enumerator a m b] -> Enumerator a m b
concatEnums = foldl (>==>) returnI
{-# INLINE concatEnums #-}
-- | 'joinI' is used to &#x201C;flatten&#x201D; 'Enumeratee's into an
-- 'Iteratee'.
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
-- | A common pattern in 'Enumeratee' implementations is to check whether
-- the inner 'Iteratee' has finished, and if so, to return its output.
-- 'checkDone' passes its parameter a continuation if the 'Iteratee'
-- can still consume input, or yields otherwise.
checkDoneEx :: Monad m =>
	Stream a' ->
	((Stream a -> Iteratee a m b) -> Iteratee a' m (Step a m b)) ->
	Enumeratee a' a m b
checkDoneEx extra _ (Yield x chunk) = returnI (Yield (Yield x chunk) extra)
checkDoneEx _ f (Continue k) = f k
checkDoneEx _ _ (Error err) = throwError err
{-# INLINE checkDoneEx #-}

-- | @checkDone = checkDoneEx (Chunks [])@
--
-- Use this for enumeratees which do not have an input buffer.
checkDone :: Monad m =>
	((Stream a -> Iteratee a m b) -> Iteratee a' m (Step a m b)) ->
	Enumeratee a' a m b
checkDone = checkDoneEx (Chunks [])
{-# INLINE checkDone #-}
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
concatMap :: Monad m => (ao -> [ai]) -> Enumeratee ao ai m b
concatMap 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.concatMap f xs)) >>== loop
mapM :: Monad m => (ao -> m ai) -> Enumeratee ao ai m b
mapM f = checkDone (continue . step) where
	step k EOF = yield (Continue k) EOF
	step k (Chunks xs) = loop k xs
	
	loop k [] = continue (step k)
	loop k (x:xs) = do
		fx <- MT.lift (f x)
		k (Chunks [fx]) >>== checkDoneEx (Chunks xs) (\k' -> loop k' xs)
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 p = 'span' (not . p)@
break :: Monad m => (a -> Bool) -> Iteratee a m [a]
break p = span $ not . p