coroutine.hs

{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, UndecidableInstances, LambdaCase #-}


import Control.Monad.State
import Control.Monad.Writer
import Data.Foldable

instance Monad m => Functor (CoroutineT m) where
  fmap f m = CoroutineT $ do
    result <- runCoroutineT m
    pure $ case result of
      Left m' -> Left (fmap f m')
      Right x -> Right (f x)

instance Monad m => Applicative (CoroutineT m) where
  pure = CoroutineT . pure . Right
  f <*> x = do
    f' <- f
    x' <- x
    pure $ f' x'

instance Monad m => Monad (CoroutineT m) where
  m >>= f = CoroutineT $ do
    runCoroutineT m >>= \case
      Left m' -> pure . Left $ m' >>= f
      Right x -> runCoroutineT (f x)

instance MonadWriter w m => MonadWriter w (CoroutineT m) where
  tell w = CoroutineT $ do
    tell w
    pure (Right ())
  listen m = CoroutineT $ do
    (next, w) <- listen $ runCoroutineT m
    runCoroutineT . fmap (\x -> (x, w)) . CoroutineT . pure $ next
  pass = undefined

newtype CoroutineT m a =
  CoroutineT { runCoroutineT :: m (Either (CoroutineT m a) a) }

yield :: Monad m => CoroutineT m ()
yield = CoroutineT $ pure (Left (pure ()))

coroutine1 :: CoroutineT (Writer String) ()
coroutine1 = do
  tell "1"
  yield
  yield
  tell "2"

coroutine2 :: CoroutineT (Writer String) ()
coroutine2 = do
  tell "a"
  yield
  tell "b"
  yield
  tell "c"
  yield
  tell "d"
  yield

runCoroutines :: Monad m => CoroutineT m () -> CoroutineT m () -> m ()
runCoroutines = go where
  go m1 m2 = runCoroutineT m1 >>= \case
    Left m1' -> go m2 m1'
    _ -> goSingle m2
  goSingle m = runCoroutineT m >>= \case
    Left next -> goSingle next
    _ -> pure ()