-- |
-- Module: Data.Enumerator.Internal
-- Copyright: 2010-2011 John Millikin
-- License: MIT
--
-- Maintainer: jmillikin@gmail.com
-- Portability: portable
--
-- Core enumerator types, and some useful primitives.
--
-- Be careful when using the functions defined in this module, as they will
-- allow you to create iteratees which violate the monad laws.
module Data.Enumerator.Internal
	( Stream (..)
	, Iteratee (..)
	, Step (..)
	, Enumerator
	, Enumeratee
	
	-- * Primitives
	, returnI
	, continue
	, yield
	
	-- * Operators
	, (>>==)
	, (==<<)
	, ($$)
	, (>==>)
	, (<==<)
	
	-- * Miscellaneous
	, enumEOF
	, checkContinue0
	, checkContinue1
	, checkDoneEx
	, checkDone
	
	) where

import           Control.Applicative as A
import qualified Control.Exception as Exc
import qualified Control.Monad as CM
import           Control.Monad.IO.Class (MonadIO, liftIO)
import           Control.Monad.Trans.Class (MonadTrans, lift)
import           Data.Function (fix)
import           Data.Monoid (Monoid, mempty, mappend, mconcat)
import           Data.Typeable ( Typeable, typeOf
                               , Typeable1, typeOf1
                               , mkTyConApp, mkTyCon)

-- | 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 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; an iteratee consumes
-- chunks of input from a stream until it either yields a value or
-- encounters an error.
--
-- Compatibility note: @Iteratee@ will become abstract in @enumerator_0.5@. If
-- you depend on internal implementation details, please import
-- @"Data.Enumerator.Internal"@.

-- 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)
	}

instance Monad m => Monad (Iteratee a m) where
	return x = yield x (Chunks [])

	m0 >>= f = ($ m0) $ fix $
		\bind m -> Iteratee $ runIteratee m >>= \r1 ->
			case r1 of
				Continue k -> return (Continue (bind . 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

-- | Enumerators are sources of data, to be consumed by iteratees.
-- Enumerators typically read from an external source (parser, handle,
-- random generator, etc), then feed chunks into an tteratee until:
--
-- * The input source runs out of data.
--
-- * The iteratee yields a result value.
--
-- * The iteratee throws an exception.

-- 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)@.
type Enumerator a m b = Step a m b -> Iteratee a m b

-- | An enumeratee acts as a stream adapter; place one between an enumerator
-- and an iteratee, and it changes the type or contents of the input stream.
--
-- Most users will want to combine enumerators, enumeratees, and iteratees
-- using the stream combinators @joinI@ and @joinE@, or their operator aliases
-- @(=$)@ and @($=)@. These combinators are used to manage how left-over input
-- is passed between elements of the data processing pipeline.
type Enumeratee ao ai m b = Step ai m b -> Iteratee ao m (Step ai m b)

-- | Since: 0.4.8
instance Typeable1 Stream where
	typeOf1 _ = mkTyConApp tyCon [] where
		tyCon = mkTyCon "Data.Enumerator.Stream"

-- | Since: 0.4.6
instance (Typeable a, Typeable1 m) =>
	Typeable1 (Iteratee a m) where
		typeOf1 i = let
			tyCon = mkTyCon "Data.Enumerator.Iteratee"
			(a, m) = peel i
			
			peel :: Iteratee a m b -> (a, m ())
			peel = undefined
			
			in mkTyConApp tyCon [typeOf a, typeOf1 m]

-- | Since: 0.4.8
instance (Typeable a, Typeable1 m) =>
	Typeable1 (Step a m) where
		typeOf1 s = let
			tyCon = mkTyCon "Data.Enumerator.Step"
			(a, m) = peel s
			
			peel :: Step a m b -> (a, m ())
			peel = undefined
			
			in mkTyConApp tyCon [typeOf a, typeOf1 m]

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

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

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

-- | Since: 0.4.5
instance A.Applicative Stream where
	pure = return
	(<*>) = CM.ap

-- | @'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)@
--
-- WARNING: due to the current encoding of iteratees in this library,
-- careless use of the 'yield' primitive may violate the monad laws.
-- To prevent this, always make sure that an iteratee never yields
-- extra data unless it has received at least one input element.
--
-- More strictly, iteratees may not yield data that they did not
-- receive as input. Don't use 'yield' to &#x201c;inject&#x201d; elements
-- into the stream.
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)

infixl 1 >>==
infixr 1 ==<<
infixr 0 $$
infixr 1 >==>
infixr 1 <==<

-- | The most primitive stream operator. @iter >>== enum@ returns a new
-- iteratee which will read from @enum@ before continuing.
(>>==) :: 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)

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

-- | @('$$') = ('==<<')@
--
-- This is somewhat easier to read when constructing an iteratee from many
-- processing stages. You can treat it like @('$')@, and read the data flow
-- from left to right.
--
-- Since: 0.1.1
($$) :: Monad m
     => (Step a m b -> Iteratee a' m b')
     -> Iteratee a m b
     -> Iteratee a' m b'
($$) = (==<<)

-- | @('>==>') enum1 enum2 step = enum1 step '>>==' enum2@
--
-- The moral equivalent of @('CM.>=>')@ for iteratees.
--
-- 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

-- | @('<==<') = 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 (>==>)

-- | Sends 'EOF' to its iteratee. Most clients should use 'run' or 'run_'
-- instead.
enumEOF :: Monad m => Enumerator a m b
enumEOF (Yield x _) = yield x EOF
enumEOF (Error err) = returnI (Error err)
enumEOF (Continue k) = k EOF >>== check where
	check (Continue _) = error "enumEOF: divergent iteratee"
	check s = enumEOF s

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

-- | A common pattern in 'Enumerator' implementations is to check whether
-- the inner 'Iteratee' has finished, and if so, to return its output.
-- 'checkContinue0' passes its parameter a continuation if the 'Iteratee'
-- can still consume input; if not, it returns the iteratee's step.
--
-- The type signature here is a bit crazy, but it's actually very easy to
-- use. Take this code:
--
-- > repeat :: Monad m => a -> Enumerator a m b
-- > repeat x = loop where
-- > 	loop (Continue k) = k (Chunks [x]) >>== loop
-- > 	loop step = returnI step
--
-- And rewrite it without the boilerplate:
--
-- > repeat :: Monad m => a -> Enumerator a m b
-- > repeat x = checkContinue0 $ \loop k -> k (Chunks [x] >>== loop
--
-- Since: 0.4.9
checkContinue0 :: Monad m
               => (Enumerator a m b
                -> (Stream a -> Iteratee a m b)
                -> Iteratee a m b)
               -> Enumerator a m b
checkContinue0 inner = loop where
	loop (Continue k) = inner loop k
	loop step = returnI step


-- | Like 'checkContinue0', but allows each loop step to use a state value:
--
-- > iterate :: Monad m => (a -> a) -> a -> Enumerator a m b
-- > iterate f = checkContinue1 $ \loop a k -> k (Chunks [a]) >>== loop (f a)
--
-- Since: 0.4.9
checkContinue1 :: Monad m
               => ((s1 -> Enumerator a m b)
                -> s1
                -> (Stream a -> Iteratee a m b)
                -> Iteratee a m b)
               -> s1
               -> Enumerator a m b
checkContinue1 inner = loop where
	loop s (Continue k) = inner loop s k
	loop _ step = returnI step