tube_coro.hs

tube_coro.hs

{-# LANGUAGE Rank2Types #-}

import Prelude hiding (break, take)
import Tubes
import Data.Functor.Identity
import Control.Monad (forever)

{- |
This just holds a value and hands it over when the tube asks, or replaces it
with whatever a tube yields.

When paired with a tube it effectively models the state monad, with 'yield'
corresponding to @put@ and 'await' corresponding to @get@.
-}
box :: a -> Pump a a Identity a
box x = mkPump (Identity x)
    (\(Identity x) -> (x, Identity x))
    (\(Identity _) x' -> Identity x')

while :: (Monad m) => a -> (a -> Bool) -> (a -> m (Maybe a)) -> m a
while thing pred fn = pump const (box thing) loop where
    loop = do
        val <- await
        if pred val
            then do
                result <- lift $ fn val
                case result of
                    Just val' -> (yield val') >> loop
                    Nothing   -> return ()
            else return ()

-- | This makes while loops more readable
next :: Monad m => a -> m (Maybe a)
next = return . Just

-- | This also makes while loops more readable
break :: Monad m => m (Maybe a)
break = return Nothing

{- |
And here we loop over some state, modifying it and continuing until we break
or the predicate no longer holds true.
-}
while_test :: IO ()
while_test = do
    let x = 1
    while x (<= 10) $ \v -> do
        putStrLn $ "v is " ++ (show v)
        if v == 5
            then break
            else next $ v + 1
    return ()

-- * Coroutines!

type Coro m a = a -> m (Maybe a)

{- |
But isn't a while loop just a special case of coroutines where you only have 1?

If we have a list of coroutines as well as an algorithm for scheduling them,
we can execute numerous while loops concurrently, all sharing the same state.
-}

roundRobin :: Monad m => a -> [Coro m a] -> m a
roundRobin thing cs = pump const (box thing) $ loop 0 where
    loop n = do
        v <- await
        r <- lift $ (cs !! n) v
        case r of
            Just v' -> do
                yield v'
                loop $ (n + 1) `mod` len
            Nothing -> return ()

    len = length cs

-- | A coroutine which logs and monitors the current state.
coro1 :: Coro IO Int
coro1 n = do
    putStrLn $ "[1]: " ++ (show n)
    if n == 1
        then break
        else next n

-- | A coroutine which performs a Collatz sequence (for demo purposes)
coro2 :: Coro IO Int
coro2 n = do
    if (n `mod` 2 == 0)
        then next $ n `div` 2
        else next $ (3 * n) + 1

-- | The two coroutines, completely decoupled from one another, are composed.
coro_test :: IO ()
coro_test = do
    roundRobin 50 [coro1, coro2]
    return ()