{- 
    Copyright 2010 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 #-}

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

import Prelude hiding (foldl, foldr)
import Control.Monad (Monad, liftM)
import Control.Monad.Trans.Class (MonadTrans(..))
import Data.Foldable (Foldable, foldl, foldr)

import Control.Monad.Coroutine
import Control.Cofunctor.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)

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

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

-- | 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 concat return
   where concat s = do chunk <- await
                       concatAwaits (feedAll chunk (suspend s))

-- | 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 lift = 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 $ lift $ Request rest c2)
}

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

-- | Feeds a collection of values to an awaiting coroutine.
feedAll :: (Foldable f, Monad m) => f x -> Coroutine (Await x) m r -> Coroutine (Await x) m r
feedAll chunk c = foldl (flip feed) c chunk