-----------------------------------------------------------------------------
-- |
-- Module: Data.Enumerator
-- Copyright: 2010 John Millikin
-- License: MIT
--
-- Maintainer: jmillikin@gmail.com
-- Portability: portable
--
-- Core enumerator types, and some useful primitives.
--
-- This module is intended to be imported qualified:
--
-- @
-- import qualified Data.Enumerator as E
-- @
--
-----------------------------------------------------------------------------

module Data.Enumerator (

	-- * Core
	-- ** Types
	  Stream (..)
	, Iteratee (..)
	, Step (..)
	, Enumerator
	, Enumeratee

	, returnI
	, yield
	, continue

	-- ** Operators
	, (>>==)
	, (==<<)
	, ($$)
	, (>==>)
	, (<==<)

	-- * Primitives

	-- ** Error handling
	, throwError
	, catchError

	-- ** Iteratees
	, Data.Enumerator.foldl
	, Data.Enumerator.foldl'
	, Data.Enumerator.foldM

	-- ** Enumerators
	, Data.Enumerator.iterate
	, iterateM
	, Data.Enumerator.repeat
	, repeatM
	, Data.Enumerator.replicate
	, replicateM
	, generateM

	-- ** Enumeratees
	, Data.Enumerator.map
	, Data.Enumerator.concatMap
	, Data.Enumerator.filter
	, Data.Enumerator.mapM
	, concatMapM
	, Data.Enumerator.filterM

	-- ** Debugging
	, printChunks

	-- * Misc. utilities

	, concatEnums

	, joinI
	, joinE
	, Data.Enumerator.sequence
	, enumList

	, enumEOF
	, run
	, run_

	, checkDone
	, checkDoneEx

	, isEOF

	-- * Compatibility

	-- ** Obsolete functions
	, liftTrans
	, liftI
	, peek
	, Data.Enumerator.last
	, Data.Enumerator.length

	-- ** Deprecated aliases
	, Data.Enumerator.head
	, Data.Enumerator.drop
	, Data.Enumerator.dropWhile
	, Data.Enumerator.span
	, Data.Enumerator.break
	, Data.Enumerator.consume

	, liftFoldL
	, liftFoldL'
	, liftFoldM

	) where

import qualified Prelude as Prelude
import Prelude hiding (

	concatMap,

	)

import Data.Monoid (Monoid, mempty, mappend, mconcat)

import qualified Control.Exception as Exc

import Control.Monad.Trans.Class (MonadTrans, lift)
import Control.Monad.IO.Class (MonadIO, liftIO)

import qualified Control.Applicative as A
import qualified Control.Monad as CM

import Data.List (foldl')

import Data.List (genericSplitAt)

import Data.List (genericLength)

import {-# SOURCE #-} qualified Data.Enumerator.List as EL


-- | A 'Stream' is a sequence of chunks generated by an 'Enumerator'.
--
-- @('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)

instance Monad Stream where
	return = Chunks . return
	Chunks xs >>= f = mconcat (fmap f xs)
	EOF >>= _ = EOF

instance Functor Stream where
	fmap f (Chunks xs) = Chunks (fmap f xs)
	fmap _ EOF = EOF

instance A.Applicative Stream where
	pure = return
	(<*>) = CM.ap

instance Monoid (Stream a) where
	mempty = Chunks mempty
	mappend (Chunks xs) (Chunks ys) = Chunks (xs ++ ys)
	mappend _ _ = EOF

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' cannot receive any more input, and has generated 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 Exc.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)
	}


-- | @returnI step = 'Iteratee' (return step)@

returnI :: Monad m => Step a m b -> Iteratee a m b
returnI step = Iteratee (return step)


-- | @yield x extra = 'returnI' ('Yield' x extra)@

yield :: Monad m => b -> Stream a -> Iteratee a m b
yield x extra = returnI (Yield x extra)


-- | @continue k = 'returnI' ('Continue' k)@

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


-- | 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 ao ai m b = Step ai m b
          -> Iteratee ao m (Step ai m b)

infixl 1 >>==


-- | Equivalent to '(>>=)' for @m ('Step' a m b)@; 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)

infixr 1 ==<<


-- | @(==\<\<) = flip (\>\>==)@

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

infixr 0 $$


-- | @($$) = (==\<\<)@
--
-- This might be easier to read when passing a chain of iteratees to an
-- enumerator.
--
-- Since: 0.1.1

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

infixr 1 >==>


-- | @(>==>) e1 e2 s = e1 s >>== e2@
--
-- Since: 0.1.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 <==<


-- | @(\<==\<) = flip (>==>)@
--
-- Since: 0.1.1

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

instance Monad m => Monad (Iteratee a m) where
	return x = 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 extra -> runIteratee (f x) >>=
				\r2 -> case r2 of
					Continue k -> runIteratee (k extra)
					Error err -> return (Error err)
					Yield x' _ -> return (Yield x' extra)

instance MonadTrans (Iteratee a) where
	lift m = Iteratee (m >>= runIteratee . return)

instance MonadIO m => MonadIO (Iteratee a m) where
	liftIO = lift . liftIO

instance Monad m => Functor (Iteratee a m) where
	fmap = CM.liftM

instance Monad m => A.Applicative (Iteratee a m) where
	pure = return
	(<*>) = CM.ap


-- | @throwError exc = 'returnI' ('Error' ('Exc.toException' exc))@

throwError :: (Monad m, Exc.Exception e) => e
           -> Iteratee a m b
throwError exc = returnI (Error (Exc.toException exc))


-- | Runs the iteratee, and calls an exception handler if an 'Error' is
-- returned. By handling errors within the enumerator library, and requiring
-- all errors to be represented by 'Exc.SomeException', libraries with
-- varying error types can be easily composed.
--
-- Since: 0.1.1

catchError :: Monad m => Iteratee a m b
           -> (Exc.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 (\s -> k s >>== step)


-- | Run the entire input stream through a pure left fold, yielding when
-- there is no more input.
--
-- Since: 0.4.5

foldl :: Monad m => (b -> a -> b) -> b
      -> Iteratee a m b
foldl step = continue . loop where
	fold = Prelude.foldl step
	loop acc stream = case stream of
		Chunks [] -> continue (loop acc)
		Chunks xs -> continue (loop (fold acc xs))
		EOF -> yield acc EOF


-- | Run the entire input stream through a pure strict left fold, yielding
-- when there is no more input.
--
-- Since: 0.4.5

foldl' :: Monad m => (b -> a -> b) -> b
       -> Iteratee a m b
foldl' step = continue . loop where
	fold = Data.List.foldl' step
	loop acc stream = case stream of
		Chunks [] -> continue (loop acc)
		Chunks xs -> continue (loop (fold acc xs))
		EOF -> yield acc EOF


-- | Run the entire input stream through a monadic left fold, yielding
-- when there is no more input.
--
-- Since: 0.4.5

foldM :: Monad m => (b -> a -> m b) -> b
      -> Iteratee a m b
foldM step = continue . loop where
	fold acc = lift . CM.foldM step acc
	
	loop acc stream = case stream of
		Chunks [] -> continue (loop acc)
		Chunks xs -> fold acc xs >>= continue . loop
		EOF -> yield acc EOF


-- | @iterate f x@ enumerates an infinite stream of repeated applications
-- of /f/ to /x/.
--
-- Analogous to 'Prelude.iterate'.
--
-- Since: 0.4.5

iterate :: Monad m => (a -> a) -> a -> Enumerator a m b
iterate f = loop where
	loop a (Continue k) = k (Chunks [a]) >>== loop (f a)
	loop _ step = returnI step


-- | Similar to 'iterate', except the iteration function is monadic.
--
-- Since: 0.4.5

iterateM :: Monad m => (a -> m a) -> a
         -> Enumerator a m b
iterateM f base = loop (return base) where
	loop m_a (Continue k) = do
		a <- lift m_a
		k (Chunks [a]) >>== loop (f a)
	loop _ step = returnI step


-- | Enumerates an infinite stream of the provided value.
--
-- Analogous to 'Prelude.repeat'.
--
-- Since: 0.4.5

repeat :: Monad m => a -> Enumerator a m b
repeat a = Data.Enumerator.iterate (const a) a


-- | Enumerates an infinite stream by running the provided computation and
-- passing each result to the iteratee.
--
-- Since: 0.4.5

repeatM :: Monad m => m a -> Enumerator a m b
repeatM m_a step = do
	a <- lift m_a
	iterateM (const m_a) a step


-- | @replicateM n m_x@ enumerates a stream of /n/ input elements; each
-- element is generated by running the input computation /m_x/ once.
--
-- Since: 0.4.5

replicateM :: Monad m => Integer -> m a
           -> Enumerator a m b
replicateM maxCount getNext = loop maxCount where
	loop 0 step = returnI step
	loop n (Continue k) = do
		next <- lift getNext
		k (Chunks [next]) >>== loop (n - 1)
	loop _ step = returnI step


-- | @replicate n x = 'replicateM' n (return x)@
--
-- Analogous to 'Prelude.replicate'.
--
-- Since: 0.4.5

replicate :: Monad m => Integer -> a
          -> Enumerator a m b
replicate maxCount a = replicateM maxCount (return a)


-- | Like 'repeatM', except the computation may terminate the stream by
-- returning 'Nothing'.
--
-- Since: 0.4.5

generateM :: Monad m => m (Maybe a)
          -> Enumerator a m b
generateM getNext = loop where
	loop (Continue k) = do
		next <- lift getNext
		case next of
			Nothing -> continue k
			Just x -> k (Chunks [x]) >>== loop
	loop step = returnI step


-- | @concatMapM f@ applies /f/ to each input element and feeds the
-- resulting outputs to the inner iteratee.
--
-- Since: 0.4.5

concatMapM :: Monad m => (ao -> m [ai])
           -> Enumeratee ao ai m b
concatMapM 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 <- lift (f x)
		k (Chunks fx) >>==
			checkDoneEx (Chunks xs) (\k' -> loop k' xs)


-- | @concatMap f = 'concatMapM' (return . f)@
--
-- Since: 0.4.3

concatMap :: Monad m => (ao -> [ai])
          -> Enumeratee ao ai m b
concatMap f = concatMapM (return . f)


-- | @map f = 'concatMap' (\x -> 'Prelude.map' f [x])@

map :: Monad m => (ao -> ai)
    -> Enumeratee ao ai m b
map f = concatMap (\x -> Prelude.map f [x])


-- | @filter p = 'concatMap' (\x -> 'Prelude.filter' p [x])@
--
-- Since: 0.4.5

filter :: Monad m => (a -> Bool)
       -> Enumeratee a a m b
filter p = concatMap (\x -> Prelude.filter p [x])


-- | @mapM f = 'concatMapM' (\x -> 'Prelude.mapM' f [x])@
--
-- Since: 0.4.3

mapM :: Monad m => (ao -> m ai)
     -> Enumeratee ao ai m b
mapM f = concatMapM (\x -> Prelude.mapM f [x])


-- | @filterM p = 'concatMapM' (\x -> 'CM.filterM' p [x])@
--
-- Since: 0.4.5

filterM :: Monad m => (a -> m Bool)
        -> Enumeratee a a m b
filterM p = concatMapM (\x -> CM.filterM p [x])


-- | Print chunks as they're received from the enumerator, optionally
-- printing empty chunks.

printChunks :: (MonadIO m, Show a)
            => Bool -- ^ Print empty chunks
            -> Iteratee a m ()
printChunks printEmpty = continue loop where
	loop (Chunks xs) = do
		let hide = null xs && not printEmpty
		CM.unless hide (liftIO (print xs))
		continue loop
	
	loop EOF = do
		liftIO (putStrLn "EOF")
		yield () EOF


-- | Compose a list of 'Enumerator's using @'(>>==)'@

concatEnums :: Monad m => [Enumerator a m b]
            -> Enumerator a m b
concatEnums = Prelude.foldl (>==>) returnI


-- | '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


-- | Flatten an enumerator/enumeratee pair into a single enumerator.

joinE :: Monad m
      => Enumerator ao m (Step ai m b)
      -> Enumeratee ao ai m b
      -> Enumerator ai m b
joinE enum enee s = Iteratee $ do
	step <- runIteratee (enumEOF $$ enum $$ enee s)
	case step of
		Error err -> return (Error err)
		Yield x _ -> return x
		Continue _ -> error "joinE: divergent iteratee"


-- | Feeds outer input elements into the provided iteratee until it yields
-- an inner input, passes that to the inner iteratee, and then loops.

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


-- | @enumList n xs@ enumerates /xs/ as a stream, passing /n/ inputs per
-- chunk.
--
-- Primarily useful for testing and debugging.

enumList :: Monad m => Integer -> [a] -> Enumerator a m b
enumList n = loop where
	loop xs (Continue k) | not (null xs) = let
		(s1, s2) = genericSplitAt n xs
		in k (Chunks s1) >>== loop s2
	loop _ step = returnI step


-- | 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 Exc.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"


-- | docs TODO

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


-- | Like 'run', except errors are converted to exceptions and thrown.
-- Primarily useful for small scripts or other simple cases.
--
-- Since: 0.4.1

run_ :: Monad m => Iteratee a m b -> m b
run_ i = run i >>= either Exc.throw return


-- | 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.
--
-- Since: 0.4.3

checkDoneEx :: Monad m =>
	Stream a' ->
	((Stream a -> Iteratee a m b) -> Iteratee a' m (Step a m b)) ->
	Enumeratee a' a m b
checkDoneEx _     f (Continue k) = f k
checkDoneEx extra _ step         = yield step extra


-- | @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 [])


-- | docs TODO

isEOF :: Monad m => Iteratee a m Bool
isEOF = continue $ \s -> case s of
	EOF -> yield True s
	_ -> yield False s


-- | Lift an 'Iteratee' onto a monad transformer, re-wrapping the
-- 'Iteratee'&#x2019;s inner monadic values.
--
-- Since: 0.1.1

liftTrans :: (Monad m, MonadTrans t, Monad (t m)) =>
             Iteratee a m b -> Iteratee a (t m) b
liftTrans iter = Iteratee $ do
	step <- lift (runIteratee iter)
	return $ case step of
		Yield x cs -> Yield x cs
		Error err -> Error err
		Continue k -> Continue (liftTrans . k)

{-# DEPRECATED liftI
     "Use 'Data.Enumerator.continue' instead" #-}

-- | Deprecated in 0.4.5: use 'Data.Enumerator.continue' instead

liftI :: Monad m => (Stream a -> Step a m b)
      -> Iteratee a m b
liftI k = continue (returnI . k)


-- | Peek at the next element in the stream, or 'Nothing' if the stream
-- has ended.

peek :: Monad m => Iteratee a m (Maybe a)
peek = continue loop where
	loop (Chunks []) = continue loop
	loop chunk@(Chunks (x:_)) = yield (Just x) chunk
	loop EOF = yield Nothing EOF


-- | Get the last element in the stream, or 'Nothing' if the stream
-- has ended.
--
-- Consumes the entire stream.

last :: Monad m => Iteratee a m (Maybe a)
last = continue (loop Nothing) where
	loop ret (Chunks xs) = continue . loop $ case xs of
		[] -> ret
		_ -> Just (Prelude.last xs)
	loop ret EOF = yield ret EOF


-- | Get how many elements remained in the stream.
--
-- Consumes the entire stream.

length :: Monad m => Iteratee a m Integer
length = continue (loop 0) where
	len = genericLength
	loop n (Chunks xs) = continue (loop (n + len xs))
	loop n EOF = yield n EOF

{-# DEPRECATED head
     "Use 'Data.Enumerator.List.head' instead" #-}

-- | Deprecated in 0.4.5: use 'Data.Enumerator.List.head' instead

head :: Monad m => Iteratee a m (Maybe a)
head = EL.head

{-# DEPRECATED drop
     "Use 'Data.Enumerator.List.drop' instead" #-}

-- | Deprecated in 0.4.5: use 'Data.Enumerator.List.drop' instead

drop :: Monad m => Integer -> Iteratee a m ()
drop = EL.drop

{-# DEPRECATED dropWhile
     "Use 'Data.Enumerator.List.dropWhile' instead" #-}

-- | Deprecated in 0.4.5: use 'Data.Enumerator.List.dropWhile' instead

dropWhile :: Monad m => (a -> Bool) -> Iteratee a m ()
dropWhile = EL.dropWhile

{-# DEPRECATED span
     "Use 'Data.Enumerator.List.takeWhile' instead" #-}

-- | Deprecated in 0.4.5: use 'Data.Enumerator.List.takeWhile' instead

span :: Monad m => (a -> Bool) -> Iteratee a m [a]
span = EL.takeWhile

{-# DEPRECATED break
     "Use 'Data.Enumerator.List.takeWhile' instead" #-}

-- | Deprecated in 0.4.5: use 'Data.Enumerator.List.takeWhile' instead

break :: Monad m => (a -> Bool) -> Iteratee a m [a]
break p = EL.takeWhile (not . p)

{-# DEPRECATED consume
     "Use 'Data.Enumerator.List.consume' instead" #-}

-- | Deprecated in 0.4.5: use 'Data.Enumerator.List.consume' instead

consume :: Monad m => Iteratee a m [a]
consume = EL.consume

{-# DEPRECATED liftFoldL
     "Use 'Data.Enumerator.foldl' instead" #-}

-- | Deprecated in 0.4.5: use 'Data.Enumerator.foldl' instead
--
-- Since: 0.1.1

liftFoldL :: Monad m => (b -> a -> b) -> b
          -> Iteratee a m b
liftFoldL = Data.Enumerator.foldl

{-# DEPRECATED liftFoldL'
     "Use 'Data.Enumerator.foldl' ' instead" #-}

-- | Deprecated in 0.4.5: use 'Data.Enumerator.foldl'' instead
--
-- Since: 0.1.1

liftFoldL' :: Monad m => (b -> a -> b) -> b
           -> Iteratee a m b
liftFoldL' = Data.Enumerator.foldl'

{-# DEPRECATED liftFoldM
     "Use 'Data.Enumerator.foldM' instead" #-}

-- | Deprecated in 0.4.5: use 'Data.Enumerator.foldM' instead
--
-- Since: 0.1.1

liftFoldM :: Monad m => (b -> a -> m b) -> b
          -> Iteratee a m b
liftFoldM = Data.Enumerator.foldM