Coroutine.hs

{- 
    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. The function 'seesaw'
-- 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, suspend,
    -- * Coroutine operations
    mapMonad, mapSuspension, 
    -- * Running Coroutine computations
    Naught, runCoroutine, pogoStick, foldRun, seesaw, SeesawResolver(..),
    -- * Nested and coupled Coroutine computations
    NestedFunctor (NestedFunctor), SomeFunctor(..), nest,
    couple, merge
   )
where

import Control.Monad (liftM, when)
import Control.Monad.Trans (MonadTrans(..), MonadIO(..))
import Data.Either (either, partitionEithers)
import Data.Traversable (Traversable, sequence)

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 f t1 t2 = Coroutine (bindM2 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 (bindM2 f $ suspend s1) s2)

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 both.
newtype NestedFunctor l r x = NestedFunctor (l (r x))
instance (Functor l, Functor r) => Functor (NestedFunctor l r) where
   fmap f (NestedFunctor lr) = NestedFunctor ((fmap . fmap) f lr)

-- | 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 (NestedFunctor 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 'NestedFunctor'.
nest :: (Functor a, Functor b) => a x -> b y -> NestedFunctor a b (x, y)
nest a b = NestedFunctor $ 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!")

-- | Run a suspendable 'Coroutine', using a function that extracts the coroutine resumption from each suspension.
pogoStick :: Monad m => (s (Coroutine s m x) -> Coroutine s m x) -> Coroutine s m x -> m x
pogoStick reveal t = resume t
                     >>= \s-> case s 
                              of Right result -> return result
                                 Left c -> pogoStick reveal (reveal c)

-- | 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 t = resume t
                >>= \s-> case s 
                         of Right result -> return (a, result)
                            Left c -> uncurry (foldRun f) (f a c)

-- | Weaves two coroutines into one. The two coroutines suspend and resume in lockstep.
couple :: forall s1 s2 m x y r. (Monad m, Functor s1, Functor s2) => 
          (forall x y r. (x -> y -> m r) -> m x -> m y -> m r)
       -> 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 $ nest 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 = SeesawResolver {
   resumeLeft  :: forall t. s1 t -> t,    -- ^ resolves the left suspension functor into the resumption it contains
   resumeRight :: forall t. s2 t -> t,    -- ^ resolves the right suspension into its resumption
   resumeAny   :: forall t1 t2 r.
                  (t1 -> r)       --  ^ continuation to resume only the left suspended coroutine
               -> (t2 -> r)       --  ^ continuation to resume the right coroutine only
               -> (t1 -> t2 -> r) --  ^ continuation to resume both coroutines
               -> s1 t1           --  ^ left suspension
               -> s2 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) => 
          (forall x y r. (x -> y -> m r) -> m x -> m y -> m r)
       -> SeesawResolver s1 s2
       -> Coroutine s1 m x -> Coroutine s2 m y -> m (x, y)
seesaw runPair resolver t1 t2 = seesaw' t1 t2 where
   seesaw' t1 t2 = runPair proceed (resume t1) (resume t2)
   proceed (Right x) (Right y) = return (x, y)
   proceed (Right x) (Left s2) = seesaw' (return x) (resumeRight resolver s2)
   proceed (Left s1) (Right y) = seesaw' (resumeLeft resolver s1) (return y)
   proceed (Left s1) (Left s2) =
      resumeAny resolver (flip seesaw' (suspend s2)) (seesaw' (suspend s1)) seesaw' s1 s2