{- 
    Copyright 2009-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 the 'Coroutine' monad transformer.
-- 
-- A 'Coroutine' monadic computation can 'suspend' its execution at any time, returning control to its invoker. The
-- returned coroutine suspension is a 'Functor' containing the resumption of the coroutine. Here is an example of a
-- coroutine in the 'IO' monad that suspends computation using the functor 'Yield' from the
-- "Control.Monad.Coroutine.SuspensionFunctors" module:
-- 
-- @
-- producer :: Coroutine (Yield Int) IO String
-- producer = do yield 1
--               lift (putStrLn \"Produced one, next is four.\")
--               yield 4
--               return \"Finished\"
-- @
-- 
-- To continue the execution of a suspended 'Coroutine', apply its 'resume' method. The easiest way to run a coroutine
-- to completion is by using the 'pogoStick' function, which keeps resuming the coroutine in trampolined style until it
-- completes. Here is an example of 'pogoStick' applied to the /producer/ example above:
-- 
-- @
-- printProduce :: Show x => Coroutine (Yield x) IO r -> IO r
-- printProduce producer = pogoStick (\\(Yield x cont) -> lift (print x) >> cont) producer
-- @
-- 
-- Multiple concurrent coroutines can be run as well, and this module provides two different ways. Functions 'seesaw'
-- and 'seesawSteps' can be used to run two interleaved computations. Another possible way is to use the functions
-- 'couple' or 'merge' to weave together steps of different coroutines into a single coroutine, which can then be
-- executed by 'pogoStick'.
-- 
-- For other uses of trampoline-style coroutines, see
-- 
-- > Trampolined Style - Ganz, S. E. Friedman, D. P. Wand, M, ACM SIGPLAN NOTICES, 1999, VOL 34; NUMBER 9, pages 18-27
-- 
-- and
-- 
-- > The Essence of Multitasking - William L. Harrison, Proceedings of the 11th International Conference on Algebraic
-- > Methodology and Software Technology, volume 4019 of Lecture Notes in Computer Science, 2006

{-# LANGUAGE ScopedTypeVariables, Rank2Types, EmptyDataDecls #-}

module Control.Monad.Coroutine
   (
    -- * Coroutine definition
    Coroutine(Coroutine, resume), CoroutineStepResult, suspend,
    -- * Coroutine operations
    mapMonad, mapSuspension, 
    -- * Running Coroutine computations
    Naught, runCoroutine, bounce, pogoStick, foldRun, seesaw, SeesawResolver(..), seesawSteps,
    -- * Coupled Coroutine computations
    PairBinder, sequentialBinder, parallelBinder, liftBinder, SomeFunctor(..), composePair,
    couple, merge
   )
where

import Control.Monad (liftM, when)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Data.Either (either, partitionEithers)
import Data.Functor.Compose (Compose(..))

import Control.Monad.Parallel

-- | Suspending, resumable monadic computations.
newtype Coroutine s m r = Coroutine {
   -- | Run the next step of a `Coroutine` computation. The result of the step execution will be either a suspension or
   -- the final coroutine result.
   resume :: m (Either (s (Coroutine s m r)) r)
   }

type CoroutineStepResult s m r = Either (s (Coroutine s m r)) r

instance (Functor s, Monad m) => Monad (Coroutine s m) where
   return x = Coroutine (return (Right x))
   t >>= f = Coroutine (resume t >>= apply f)
      where apply f (Right x) = resume (f x)
            apply f (Left s) = return (Left (fmap (>>= f) s))
--   t >>= f = Coroutine (resume t >>= either (return . Left . fmap (>>= f)) (resume . f))

instance (Functor s, MonadParallel m) => MonadParallel (Coroutine s m) where
   bindM2 = liftBinder bindM2

instance Functor s => MonadTrans (Coroutine s) where
   lift = Coroutine . liftM Right

instance (Functor s, MonadIO m) => MonadIO (Coroutine s m) where
   liftIO = lift . liftIO

-- | The 'Naught' functor instance doesn't contain anything and cannot be constructed. Used for building non-suspendable
-- coroutines.
data Naught x
instance Functor Naught where
   fmap f _ = undefined

-- | Combines two functors into one, applying either or both of them. Used for coupled coroutines.
data SomeFunctor l r x = LeftSome (l x) | RightSome (r x) | Both (Compose l r x)
instance (Functor l, Functor r) => Functor (SomeFunctor l r) where
   fmap f (LeftSome l) = LeftSome (fmap f l)
   fmap f (RightSome r) = RightSome (fmap f r)
   fmap f (Both lr) = Both (fmap f lr)

-- | Combines two values under two functors into a pair of values under a single 'Compose'.
composePair :: (Functor a, Functor b) => a x -> b y -> Compose a b (x, y)
composePair a b = Compose $ fmap (\x-> fmap ((,) x) b) a

-- | Suspend the current 'Coroutine'.
suspend :: (Monad m, Functor s) => s (Coroutine s m x) -> Coroutine s m x
suspend s = Coroutine (return (Left s))

-- | Change the base monad of a 'Coroutine'.
mapMonad :: forall s m m' x. (Functor s, Monad m, Monad m') =>
            (forall x. m x -> m' x) -> Coroutine s m x -> Coroutine s m' x
mapMonad f cort = Coroutine {resume= liftM map' (f $ resume cort)}
   where map' (Right r) = Right r
         map' (Left s) = Left (fmap (mapMonad f) s)

-- | Change the suspension functor of a 'Coroutine'.
mapSuspension :: forall s s' m x. (Functor s, Monad m) => (forall x. s x -> s' x) -> Coroutine s m x -> Coroutine s' m x
mapSuspension f cort = Coroutine {resume= liftM map' (resume cort)}
   where map' (Right r) = Right r
         map' (Left s) = Left (f $ fmap (mapSuspension f) s)

-- | Convert a non-suspending 'Coroutine' to the base monad.
runCoroutine :: Monad m => Coroutine Naught m x -> m x
runCoroutine = pogoStick (error "runCoroutine can run only a non-suspending coroutine!")

-- | Runs a single step of a suspendable 'Coroutine', using a function that extracts the coroutine resumption from its
-- suspension functor.
bounce :: (Monad m, Functor s) => (s (Coroutine s m x) -> Coroutine s m x) -> Coroutine s m x -> Coroutine s m x
bounce spring c = lift (resume c) >>= either spring return

-- | Runs a suspendable 'Coroutine' to its completion.
pogoStick :: Monad m => (s (Coroutine s m x) -> Coroutine s m x) -> Coroutine s m x -> m x
pogoStick spring c = resume c >>= either (pogoStick spring . spring) return

-- | Runs a suspendable coroutine much like 'pogoStick', but allows the resumption function to thread an arbitrary
-- state as well.
foldRun :: Monad m => (a -> s (Coroutine s m x) -> (a, Coroutine s m x)) -> a -> Coroutine s m x -> m (a, x)
foldRun f a c = resume c
                >>= \s-> case s 
                         of Right result -> return (a, result)
                            Left c -> uncurry (foldRun f) (f a c)

-- | Type of functions that can bind two monadic values together; used to combine two coroutines' step results.
type PairBinder m = forall x y r. (x -> y -> m r) -> m x -> m y -> m r

-- | A 'PairBinder' that runs the two steps sequentially before combining their results.
sequentialBinder :: Monad m => PairBinder m
sequentialBinder f mx my = do {x <- mx; y <- my; f x y}

-- | A 'PairBinder' that runs the two steps in parallel.
parallelBinder :: MonadParallel m => PairBinder m
parallelBinder = bindM2

-- | Lifting a 'PairBinder' onto a 'Coroutine' monad transformer.
liftBinder :: forall s m. (Functor s, Monad m) => PairBinder m -> PairBinder (Coroutine s m)
liftBinder binder f t1 t2 = Coroutine (binder combine (resume t1) (resume t2)) where
   combine (Right x) (Right y) = resume (f x y)
   combine (Left s) (Right y) = return $ Left (fmap (flip f y =<<) s)
   combine (Right x) (Left s) = return $ Left (fmap (f x =<<) s)
   combine (Left s1) (Left s2) = return $ Left (fmap (liftBinder binder f $ suspend s1) s2)

-- | Weaves two coroutines into one. The two coroutines suspend and resume in lockstep. The combined coroutine suspends
-- as long as either argument coroutine suspends, and it completes execution when both arguments do.
couple :: forall s1 s2 m x y r. (Monad m, Functor s1, Functor s2) => 
          PairBinder m -> Coroutine s1 m x -> Coroutine s2 m y -> Coroutine (SomeFunctor s1 s2) m (x, y)
couple runPair t1 t2 = Coroutine{resume= runPair proceed (resume t1) (resume t2)} where
   proceed :: CoroutineStepResult s1 m x -> CoroutineStepResult s2 m y
           -> m (CoroutineStepResult (SomeFunctor s1 s2) m (x, y))
   proceed (Right x) (Right y) = return $ Right (x, y)
   proceed (Left s1) (Left s2) = return $ Left
                                 $ fmap (uncurry (couple runPair)) (Both $ composePair s1 s2)
   proceed (Right x) (Left s2) = return $ Left $ fmap (couple runPair (return x)) (RightSome s2)
   proceed (Left s1) (Right y) = return $ Left
                                 $ fmap (flip (couple runPair) (return y)) (LeftSome s1)

-- | Weaves a list of coroutines with the same suspension functor type into a single coroutine. The coroutines suspend
-- and resume in lockstep.
merge :: forall s m x. (Monad m, Functor s) =>
         (forall x. [m x] -> m [x]) -> (forall x. [s x] -> s [x])
      -> [Coroutine s m x] -> Coroutine s m [x]
merge sequence1 sequence2 corts = Coroutine{resume= liftM step $ sequence1 (map resume corts)} where
   step :: [CoroutineStepResult s m x] -> CoroutineStepResult s m [x]
   step list = case partitionEithers list
               of ([], ends) -> Right ends
                  (suspensions, ends) -> Left $ fmap (merge sequence1 sequence2 . (map return ends ++)) $
                                         sequence2 suspensions

-- | A simple record containing the resolver functions for all possible coroutine pair suspensions.
data SeesawResolver s1 s2 s1' s2' = SeesawResolver {
   resumeLeft  :: forall m t. (Monad m) => s1 (Coroutine s1' m t) -> Coroutine s1' m t,
   -- ^ resolves the left suspension functor into the resumption it contains
   resumeRight :: forall m t. (Monad m) => s2 (Coroutine s2' m t) -> Coroutine s2' m t,
   -- ^ resolves the right suspension into its resumption
   resumeBoth  :: forall m t1 t2 r. (Monad m) =>
                  (Coroutine s1' m t1 -> Coroutine s2' m t2 -> r) --  ^ continuation to resume both coroutines
               -> s1 (Coroutine s1' m t1)                         --  ^ left suspension
               -> s2 (Coroutine s2' m t2)                         --  ^ right suspension
               -> r
   -- ^ invoked when both coroutines are suspended, resolves both suspensions or either one
}

-- | Runs two coroutines concurrently. The first argument is used to run the next step of each coroutine, the next to
-- convert the left, right, or both suspensions into the corresponding resumptions.
seesaw :: (Monad m, Functor s1, Functor s2) => 
          PairBinder m -> SeesawResolver s1 s2 s1 s2 -> Coroutine s1 m x -> Coroutine s2 m y -> m (x, y)
seesaw runPair resolver t1 t2 = seesawSteps runPair proceed t1 t2 where
   proceed cont (Left s1) (Left s2) = resumeBoth resolver cont s1 s2
   proceed _ (Right x) (Left s2) = liftM ((,) x) $ pogoStick (resumeRight resolver) (resumeRight resolver s2)
   proceed _ (Left s1) (Right y) = liftM (flip (,) y) $ pogoStick (resumeLeft resolver) (resumeLeft resolver s1)
   proceed _ (Right x) (Right y) = return (x, y)

-- | Runs two coroutines concurrently. The first argument is used to run the next step of each coroutine, the next to
-- convert their step results into the corresponding resumptions.
seesawSteps :: (Monad m, Functor s1, Functor s2) => 
               PairBinder m
            -> ((Coroutine s1 m x -> Coroutine s2 m y -> m (x, y)) 
                -> CoroutineStepResult s1 m x -> CoroutineStepResult s2 m y -> m (x, y))
            -> Coroutine s1 m x -> Coroutine s2 m y -> m (x, y)
seesawSteps runPair proceed t1 t2 = seesaw' t1 t2 where
   seesaw' t1 t2 = runPair (proceed seesaw') (resume t1) (resume t2)