{- 
    Copyright 2010-2011 Mario Blazevic

    This file is part of the Streaming Component Combinators (SCC) project.

    The SCC project is free software: you can redistribute it and/or modify it under the terms of the GNU General Public
    License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later
    version.

    SCC is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty
    of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.

    You should have received a copy of the GNU General Public License along with SCC.  If not, see
    <http://www.gnu.org/licenses/>.
-}

-- | This module defines some common suspension functors for use with the "Control.Monad.Coroutine" module.
-- 

{-# LANGUAGE Rank2Types, ExistentialQuantification #-}

module Control.Monad.Coroutine.SuspensionFunctors
   (
    -- * Suspension functors
    Yield(Yield), Await(Await), Request(Request), ParseRequest, EitherFunctor(LeftF, RightF), eitherFunctor,
    yield, await, request, requestParse,
    -- * Utility functions
    concatYields, concatAwaits,
    -- * Resolvers for running pairs of coroutines
    awaitYieldResolver, awaitMaybeYieldResolver, awaitYieldChunkResolver, requestsResolver, 
    tickerYieldResolver, tickerRequestResolver, lazyTickerRequestResolver, 
    parserRequestResolver, lazyParserRequestResolver,
    liftedTickerYieldResolver, liftedTickerRequestResolver, liftedLazyTickerRequestResolver,
    liftedParserRequestResolver, nestedLazyParserRequestResolver,
   )
where

import Prelude hiding (foldl, foldr)
import Control.Monad.Trans.Class (MonadTrans(..))
import Data.Foldable (Foldable, foldl, foldr)
import Data.Monoid (Monoid, mempty)
import Data.Monoid.Null (MonoidNull, mnull)
import Text.ParserCombinators.Incremental (Parser, feed, feedEof, results, (><))

import Control.Monad.Coroutine
import Data.Functor.Contravariant.Ticker (Ticker, splitTicked)

-- | The 'Yield' functor instance is equivalent to (,) but more descriptive.
data Yield x y = Yield x y
instance Functor (Yield x) where
   fmap f (Yield x y) = Yield x (f y)

-- | The 'Await' functor instance is equivalent to (->) but more descriptive.
newtype Await x y = Await (x -> y)
instance Functor (Await x) where
   fmap f (Await g) = Await (f . g)

-- | The 'Request' functor instance combines a 'Yield' of a request with an 'Await' for a response.
data Request request response x = Request request (response -> x)
instance Functor (Request x f) where
   fmap f (Request x g) = Request x (f . g)

data ParseRequest x z = forall a y. MonoidNull y => 
                        ParseRequest ([x] -> [x]) (Parser a [x] y) ((y, Maybe (Parser a [x] y)) -> z)
instance Functor (ParseRequest x) where
   fmap f (ParseRequest b p g) = ParseRequest b p (f . g)

-- | Combines two alternative functors into one, applying one or the other. Used for nested coroutines.
data EitherFunctor l r x = LeftF (l x) | RightF (r x)
instance (Functor l, Functor r) => Functor (EitherFunctor l r) where
   fmap f (LeftF l) = LeftF (fmap f l)
   fmap f (RightF r) = RightF (fmap f r)

-- | Like 'either' for the EitherFunctor data type.
eitherFunctor :: (l x -> y) -> (r x -> y) -> EitherFunctor l r x -> y
eitherFunctor left _ (LeftF f) = left f
eitherFunctor _ right (RightF f) = right f

-- | Suspend the current coroutine yielding a value.
yield :: Monad m => x -> Coroutine (Yield x) m ()
yield x = suspend (Yield x (return ()))

-- | Suspend the current coroutine until a value is provided.
await :: Monad m => Coroutine (Await x) m x
await = suspend (Await return)

-- | Suspend yielding a request and awaiting the response.
request :: Monad m => x -> Coroutine (Request x y) m y
request x = suspend (Request x return)

-- | Suspend yielding a request and awaiting the response.
requestParse :: (Monad m, MonoidNull y) => Parser a [x] y -> Coroutine (ParseRequest x) m (y, Maybe (Parser a [x] y))
requestParse p = suspend (ParseRequest id p return)

-- | Converts a coroutine yielding collections of values into one yielding single values.
concatYields :: (Monad m, Foldable f) => Coroutine (Yield (f x)) m r -> Coroutine (Yield x) m r
concatYields c = Coroutine{resume= resume c >>= foldChunk}
   where foldChunk (Right r) = return (Right r)
         foldChunk (Left (Yield s c')) = foldr f (resume $ concatYields c') s
         f x rest = return (Left $ Yield x (Coroutine rest))

-- | Converts a coroutine awaiting single values into one awaiting collections of values.
concatAwaits :: (Monad m, Foldable f) => Coroutine (Await x) m r -> Coroutine (Await (f x)) m r
concatAwaits c = lift (resume c) >>= either concatenate return
   where concatenate s = do chunk <- await
                            concatAwaits (feedAll chunk (suspend s))
         feedAll :: (Foldable f, Monad m) => f x -> Coroutine (Await x) m r -> Coroutine (Await x) m r
         feedAll chunk c = foldl (flip feedCoroutine) c chunk

-- | A 'SeesawResolver' for running two coroutines in parallel, one of which 'await's values while the other 'yield's
-- them. The yielding coroutine must not terminate before the other one.
awaitYieldResolver :: SeesawResolver (Await x) (Yield x) s1 s2
awaitYieldResolver = SeesawResolver {
   resumeLeft= undefined,
   resumeRight= \(Yield _ c)-> c,
   resumeBoth= \cont (Await f) (Yield x c2)-> cont (f x) c2
}

-- | A 'SeesawResolver' for running two coroutines in parallel, one of which 'await's values while the other 'yield's
-- them. If the yielding coroutine terminates before the awaiting one, the latter will receive 'Nothing'.
awaitMaybeYieldResolver :: SeesawResolver (Await (Maybe x)) (Yield x) s1 s2
awaitMaybeYieldResolver = SeesawResolver {
   resumeLeft= \(Await f)-> f Nothing,
   resumeRight= \(Yield _ c)-> c,
   resumeBoth= \cont (Await f) (Yield x c2)-> cont (f $ Just x) c2
}

-- | A 'SeesawResolver' for running two coroutines in parallel, one of which 'await's non-empty lists of values while
-- the other 'yield's them. If the yielding coroutine dies, the awaiting coroutine receives empty lists.
awaitYieldChunkResolver :: SeesawResolver (Await [x]) (Yield [x]) s1 s2
awaitYieldChunkResolver = SeesawResolver {
   resumeLeft= \(Await f)-> f [],
   resumeRight= \(Yield _ c)-> c,
   resumeBoth= \cont (Await f) (Yield chunk c2)-> cont (f chunk) c2
}

-- | A 'SeesawResolver' for running two 'request'ing coroutines in parallel. One coroutine's request becomes the other's
-- response, and vice versa.
requestsResolver :: SeesawResolver (Request x y) (Request y x) s1 s2
requestsResolver = SeesawResolver {
   resumeLeft= undefined,
   resumeRight= undefined,
   resumeBoth= \cont (Request x c1) (Request y c2)-> cont (c1 y) (c2 x)
}

-- | A 'SeesawResolver' for running two coroutines in parallel. One coroutine produces data in chunks, the other
-- consumes data in chunks. The boundaries of the two kinds of chunks need not be the same, as the consumed chunks
-- are determined by a 'Ticker' provided by the consumer's input request.
tickerYieldResolver :: SeesawResolver (Request (Ticker x) [x]) (Yield [x]) (Request (Ticker x) [x]) (Yield [x])
tickerYieldResolver = liftedTickerYieldResolver id id

-- | A generic version of 'tickerYieldResolver', allowing coroutines with 'Request' and 'Yield' functors embedded in
-- other functors.
liftedTickerYieldResolver :: (Functor s1, Functor s2) =>
                             (forall a. Request (Ticker x) [x] a -> s1 a) -> (forall a. Yield [x] a -> s2 a)
                             -> SeesawResolver (Request (Ticker x) [x]) (Yield [x]) s1 s2
liftedTickerYieldResolver lift1 lift2 = SeesawResolver {
   resumeLeft= \(Request _ c)-> c [],
   resumeRight= \(Yield _ c)-> c,
   resumeBoth= \cont (Request t c1) (Yield xs c2)->
               let (t', chunk, rest) = splitTicked t xs
               in case rest
                  of [] -> cont (suspend $ lift1 $ Request t' c1) c2
                     _ -> cont (c1 chunk) (suspend $ lift2 $ Yield rest c2)
}

-- | Like 'tickerYieldResolver', the only difference being that the producing coroutine sends its chunks using 'request'
-- rather than 'yield'. The feedback received from 'request' is the unconsumed remainder of the chunk, which lets the
-- coroutine know when its sibling terminates.
tickerRequestResolver :: SeesawResolver (Request (Ticker x) [x]) (Request [x] [x])
                                        (Request (Ticker x) [x]) (Request [x] [x])
tickerRequestResolver = liftedTickerRequestResolver id id

-- | A generic version of 'tickerRequestResolver', allowing coroutines with 'Request' functors embedded in other
-- functors.
liftedTickerRequestResolver :: (Functor s1, Functor s2) =>
                               (forall a. Request (Ticker x) [x] a -> s1 a) -> (forall a. Request [x] [x] a -> s2 a)
                               -> SeesawResolver (Request (Ticker x) [x]) (Request [x] [x]) s1 s2
liftedTickerRequestResolver lift1 lift2 = SeesawResolver {
   resumeLeft= \(Request _ c)-> c [],
   resumeRight= \(Request chunk c)-> c chunk,
   resumeBoth= \cont (Request t c1) (Request xs c2)->
               let (t', chunk, rest) = splitTicked t xs
               in case rest
                  of [] -> cont (suspend $ lift1 $ Request t' c1) (c2 [])
                     _ -> cont (c1 chunk) (suspend $ lift2 $ Request rest c2)
}

-- | Like 'tickerRequestResolver', except the consuming coroutine requests receive both the selected prefix of the input
-- chunk and a peek at either the next unconsumed input item, if any, or the final 'Ticker' value. Chunks sent by the
-- producing coroutine never get combined for the consuming coroutine. This allows better synchronization between the
-- two coroutines. It also leaks the information about the produced chunk boundaries into the consuming coroutine, so
-- this resolver should be used with caution.
lazyTickerRequestResolver :: SeesawResolver (Request (Ticker x) ([x], Either x (Ticker x))) (Request [x] [x])
                                            (Request (Ticker x) ([x], Either x (Ticker x))) (Request [x] [x])
lazyTickerRequestResolver = liftedLazyTickerRequestResolver id

-- | A generic version of 'lazyTickerRequestResolver', allowing coroutines with 'Request' functors embedded in other
-- functors.
liftedLazyTickerRequestResolver :: 
   (Functor s1, Functor s2) => 
   (forall a. Request [x] [x] a -> s2 a)
   -> SeesawResolver (Request (Ticker x) ([x], Either x (Ticker x))) (Request [x] [x]) s1 s2
liftedLazyTickerRequestResolver lifter = SeesawResolver {
   resumeLeft= \(Request t c)-> c ([], Right t),
   resumeRight= \(Request chunk c)-> c chunk,
   resumeBoth= \cont (Request t c1) (Request xs c2)->
               let (t', chunk, rest) = splitTicked t xs
               in case rest
                  of [] -> cont (c1 (chunk, Right t')) (c2 [])
                     next:_ -> cont (c1 (chunk, Left next)) (suspend $ lifter $ Request rest c2)
}

-- | Like 'parserYieldResolver', the only difference being that the producing coroutine sends its chunks using 'request'
-- rather than 'yield'. The feedback received from 'request' is the unconsumed remainder of the chunk, which lets the
-- coroutine know when its sibling terminates.
parserRequestResolver :: Monoid y => SeesawResolver (Request (Parser a [x] y) y) (Request [x] [x])
                                                    (Request (Parser a [x] y) y) (Request [x] [x])
parserRequestResolver = liftedParserRequestResolver id id

-- | A generic version of 'parserRequestResolver', allowing coroutines with 'Request' functors embedded in other
-- functors.
liftedParserRequestResolver :: (Functor s1, Functor s2, Monoid y) =>
                               (forall b. Request (Parser a [x] y) y b -> s1 b) -> (forall b. Request [x] [x] b -> s2 b)
                               -> SeesawResolver (Request (Parser a [x] y) y) (Request [x] [x]) s1 s2
liftedParserRequestResolver lift1 lift2 = SeesawResolver {
   resumeLeft= \(Request _ c)-> c mempty,
   resumeRight= \(Request chunk c)-> c chunk,
   resumeBoth= \cont (Request p c1) (Request xs c2)->
               case results (feed xs p)
               of ([], Just (r, p')) -> cont (suspend $ lift1 $ Request (return r >< p') c1) (c2 [])
                  ([(r, [])], Nothing) -> cont (c1 r) (c2 [])
                  ([(r, rest)], Nothing) -> cont (c1 r) (suspend $ lift2 $ Request rest c2)
                  _ -> error "Multiple results!"
}

-- | Like 'parserRequestResolver', except the consuming coroutine requests receive both the selected prefix of the input
-- chunk and a peek at either the next unconsumed input item, if any, or the final 'Parser' value. Chunks sent by the
-- producing coroutine never get combined for the consuming coroutine. This allows better synchronization between the
-- two coroutines. It also leaks the information about the produced chunk boundaries into the consuming coroutine, so
-- this resolver should be used with caution.
lazyParserRequestResolver :: SeesawResolver (ParseRequest x) (Request [x] [x]) (ParseRequest x) (Request [x] [x])
lazyParserRequestResolver = SeesawResolver {
   resumeLeft= \(ParseRequest b p c)-> mapFirstSuspension (prependToParseRequest b) $
                                       case results (feedEof p)
                                       of ([], Nothing) -> c (mempty, Nothing)
                                          ([(r, _)], Nothing) -> c (r, Nothing)
                                          _ -> error "Multiple results!",
   resumeRight= \(Request chunk c)-> c chunk,
   resumeBoth= \cont (ParseRequest b p c1) (Request xs c2)->
               case results (if null xs then feedEof p else feed xs p)
                    of ([], Nothing) -> 
                          cont (c1 (mempty, Nothing)) (if null xs then c2 $ b [] else suspend $ Request (b xs) c2)
                       ([], Just (r, p')) -> 
                          cont (if mnull r then suspend $ ParseRequest (b . (xs ++)) p' c1 else c1 (r, Just p')) (c2 [])
                       ([(r, [])], Nothing) -> cont (c1 (r, Nothing)) (c2 [])
                       ([(r, rest)], Nothing) -> cont (c1 (r, Nothing)) (suspend $ Request rest c2)
                       (_, Nothing) -> error "Multiple results!"
}

-- | A generic version of 'lazyParserRequestResolver', allowing coroutines with 'Request' functors embedded in other
-- functors.
nestedLazyParserRequestResolver ::
   (Functor s1, Functor s2) => 
   SeesawResolver (ParseRequest x) (Request [x] [x])
                  (EitherFunctor s1 (ParseRequest x)) (EitherFunctor s2 (Request [x] [x]))
nestedLazyParserRequestResolver = SeesawResolver {
   resumeLeft= \(ParseRequest b p c)-> case results (feedEof p)
                                       of ([], Nothing) -> mapFirstSuspension (retry b) $ c (mempty, Nothing)
                                          ([(r, t)], Nothing) -> mapFirstSuspension (retry (t ++)) $ c (r, Nothing)
                                          (_, Nothing) -> error "Multiple results!",
   resumeRight= \(Request chunk c)-> c chunk,
   resumeBoth= \cont (ParseRequest b p c1) (Request xs c2)->
               case results (if null xs then feedEof p else feed xs p)
               of ([], Nothing) ->
                     cont (c1 (mempty, Nothing)) (if null xs then c2 (b []) else suspend $ RightF $ Request (b xs) c2)
                  ([], Just (r, p')) ->
                     cont
                        (if mnull r then suspend $ RightF $ ParseRequest (b . (xs ++)) p' c1 else c1 (r, Just p'))
                        (c2 [])
                  ([(r, rest)], Nothing) ->
                     cont (c1 (r, Nothing)) (if null rest then c2 [] else suspend $ RightF $ Request rest c2)
                  _ -> error "Multiple results!"
}
   where retry prefix = eitherFunctor LeftF (RightF . feedList prefix)
         feedList b (ParseRequest b' p c) = ParseRequest (b' . b) (feed (b []) p) c

-- | Feeds a single value to an awaiting coroutine.
feedCoroutine :: Monad m => x -> Coroutine (Await x) m r -> Coroutine (Await x) m r
feedCoroutine x c = bounce (\(Await f)-> f x) c

prependToParseRequest b (ParseRequest b' p' c') = ParseRequest (b . b') p' c'