{- |
Module      : Control.Monad.Catana
Copyright   : (c) Dustin DeWeese 2011-2012
License     : BSD3
Maintainer  : dustin.deweese@gmail.com
Stability   : experimental
Portability : portable

[Computation type:] Computations that both consume and produce elements lazily with support for advanced control flow using continuations

[Binding Strategy:] Binding a function to a monadic value produces a continuation which passes the unconsumed input and the combined output function in turn to the next continuation.

[Useful for:] Lazily processing a list with complex control structures.

[Zero and plus:] None.

[Example type:] @'Catana' i o b a@

The Catana monad represents computations that are both catamorphisms and anamorphisms; they both consume and produce values.  In addition, the Catana monad represents the computation in continuation-passing style, and implements callCC.
-}

module Control.Monad.Catana (
    -- * The Catana monad
    Catana(..),
    consume,
    top,
    push,
    produce,
    stop,
    evalCatana,
    evalCatana',
    parallelB,
    parallelE,
    serial
    -- * Example 1: Usage of the Catana monad
    -- $catanaExample1
  ) where

import Control.Applicative
import Control.Monad
import Control.Monad.Cont
import Data.Either
import Data.Maybe

{- $catanaExample1

An example of complex control structure using Catana:

>catanaExample1 =
>  forever $ do
>    z <- callCC $ \exit -> do
>      (x, y) <- (,) <$> consume <*> consume
>      produce (x + y)
>      produce (x * y)
>      if x > 1
>        then exit 0
>        else do
>          produce (x - y)
>          produce (x / y)
>          return 1
>    produce z

Catana monads can be converted into a function over lists using 'evalCatana'.

> evalCatana catanaExample1 [1..4]
> -- result: [3.0,2.0,-1.0,0.5,1.0,7.0,12.0,0.0] 

An example using serial and parallel data flow:

>catanaExample2 = sq `serial` (ev `parallelB` od)
>  where od = forever $ do
>               x <- consume
>               when (odd x) .
>                 produce $ x * 2
>        ev = forever $ do
>               x <- consume
>               when (even x) .
>                 produce $ x + 3
>        sq = forever $ do
>               x <- consume
>               produce x
>               produce $ x ^ 2

> let l = 1 : evalCatana catanaExample2 l
> take 10 l
> -- result: [1,2,4,5,25,7,49,10,100,50]

-}

data Catana i o b a = Catana { runCatana :: (a -> Step i o b) -> Step i o b }

data Step i o a = Yield o (Step i o a)
                | Waiting (i -> Step i o a)
                | Done a

instance Functor (Catana i o b) where
  fmap f (Catana l) = Catana $ l . (. f)

instance Applicative (Catana i o b) where

  Catana fl <*> x = Catana $ \k ->
                      fl $ \f ->
                        runCatana (fmap f x) k
  pure x = Catana $ \k -> k x

instance Monad (Catana i o b) where
  return = pure
  Catana l >>= f = Catana $ \k ->
                     l $ \x ->
                       runCatana (f x) k

instance MonadCont (Catana i o b) where
  callCC f = Catana $ \k -> runCatana (f $ Catana . const . k) k

-- |Consumes an element from the input list, returning it
-- If there is no more input, the chain of continuations ends
-- immediately; no more computations will be processed
consume :: Catana i o a i
consume = Catana Waiting

-- |Returns the next input without consuming it
top :: Catana i o a i
top = Catana $ \k -> Waiting (\i -> feed i (k i))

-- |Feeds an input into the next Waiting step
feed :: i -> Step i o a -> Step i o a
feed i (Yield o s) = Yield o (feed i s)
feed i (Waiting s) = s i
feed _ (Done x) = Done x

-- |Stops computation, ending the continuation chain
stop :: b -> Catana i o b a
stop = Catana . const . Done

-- |Pushes 'x' into the input
push :: i -> Catana i o a ()
push x = Catana $ feed x . ($())

-- |Produces 'x' in the output
produce :: o -> Catana i o a ()
produce x = Catana $ Yield x . ($())

-- |Converts a Catana monad into a function over lists
evalCatana :: Catana i o a a -> [i] -> [o]
evalCatana c = runSteps (runCatana c Done)

-- |Helper for evalCatana, runs the steps
runSteps :: Step i o a -> [i] -> [o]
runSteps (Yield o s) i = o : runSteps s i
runSteps (Waiting f) (i:is) = runSteps (f i) is
runSteps (Waiting _) [] = []
runSteps (Done x) _ = []

-- |Evaluates a Catana monad over a list returning the result and output
evalCatana' :: Catana i o a a -> [i] -> (Maybe a, [o])
evalCatana' c i = (listToMaybe x, o)
  where (x, o) = partitionEithers $ runSteps' (runCatana c Done) i

-- |Helper for evalCatana', runs the steps
runSteps' :: Step i o a -> [i] -> [Either a o]
runSteps' (Yield o s) i = Right o : runSteps' s i
runSteps' (Waiting f) (i:is) = runSteps' (f i) is
runSteps' (Waiting _) [] = []
runSteps' (Done x) _ = [Left x]

-- |Helper for parallelB, combine steps to consume the same input at the same time, using k as the continuation
parStepsB :: Step i o a -> Step i o b -> ((a,b) -> Step i o c) -> Step i o c

-- Yield when possible
parStepsB (Yield oA sA) (Yield oB sB) k = Yield oA . Yield oB $ parStepsB sA sB k
parStepsB (Yield oA sA) sB k = Yield oA $ parStepsB sA sB k
parStepsB sA (Yield oB sB) k = Yield oB $ parStepsB sA sB k

-- Wait for input
parStepsB (Waiting fA) (Waiting fB) k = Waiting $ \i -> parStepsB (fA i) (fB i) k
parStepsB (Waiting fA) sB k = Waiting $ \i -> parStepsB (fA i) sB k
parStepsB sA (Waiting fB) k = Waiting $ \i -> parStepsB sA (fB i) k

-- Apply continuation to results
parStepsB (Done xA) (Done xB) k = k (xA, xB)

-- |Combine two monads to run in parallel, consuming the same input
parallelB :: Catana i o a a -> Catana i o b b -> Catana i o c (a, b)
parallelB a b = Catana $ parStepsB (runCatana a Done) (runCatana b Done)

-- |Helper for parallelB, combine steps to consume the same input at the same time, using k as the continuation
parStepsE :: Step i o a -> Step i o b -> (Either a b -> Step i o c) -> Step i o c

-- Yield when possible
parStepsE (Yield oA sA) (Yield oB sB) k = Yield oA . Yield oB $ parStepsE sA sB k
parStepsE (Yield oA sA) sB k = Yield oA $ parStepsE sA sB k
parStepsE sA (Yield oB sB) k = Yield oB $ parStepsE sA sB k

-- Apply continuation to result
parStepsE (Done xA) _ k = k (Left xA)
parStepsE _ (Done xB) k = k (Right xB)

-- Wait for input
parStepsE (Waiting fA) (Waiting fB) k = Waiting $ \i -> parStepsE (fA i) (fB i) k

-- |Combine two monads to run in parallel, consuming the same input, stopping when either of them finish.
parallelE :: Catana i o a a -> Catana i o b b -> Catana i o c (Either a b)
parallelE a b = Catana $ parStepsE (runCatana a Done) (runCatana b Done)

serSteps :: Step io o a -> Step i io b -> (Either a b -> Step i o c) -> Step i o c

-- Yield when possible
serSteps (Yield oA sA) sB k = Yield oA $ serSteps sA sB k

-- Apply continuation to results
serSteps (Done xA) _ k = k (Left xA)
serSteps _ (Done xB) k = k (Right xB)

-- Pass output from B to A
serSteps (Waiting fA) (Yield oB sB) k = serSteps (fA oB) sB k

-- Wait for input to B
serSteps sA (Waiting fB) k = Waiting $ \i -> serSteps sA (fB i) k

-- |Combine two monads to run in serial, the first consuming the output of the second
serial :: Catana io o a a -> Catana i io b b -> Catana i o c (Either a b)
serial a b = Catana $ serSteps (runCatana a Done) (runCatana b Done)