Implementation of coroutines in Haskell

Posted on January 21, 2015

In my previous blog post, I showed off the usefulness of coroutines with some examples. In this post, I want to go into more detail how coroutines can be implemented.

There are few languages that allow coroutine implementations as a library and not a first-class construct, one of which is Haskell. For this reason I will present code from the Haskell monad-coroutine library.

newtype Coroutine s m r =
    Coroutine { resume :: m (Either (s (Coroutine s m r)) r) }

Ok, this type looks more difficult than it actually is. I explain it from the outside in

The Coroutine type has a Functor, an Applicative and a Monad instance. In my opinion the monad instance is the most interesting part and I will explain it in more detail in this post.

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 fc (Right x) = resume (fc x)
            apply fc (Left s) = return (Left (fmap (>>= fc) s))

The return x function produces a coroutine that terminates immediately with the value x. The t >>= f function runs the coroutine t until it terminates and then calls the continuation f with the termination value.

Suspension functors

The infamous suspension functor is the part that decides what can be done in each step of the coroutine. This can be for example the Yield functor that produces a new value in each step, or the Await functor that consumes a value in each step, or it could be the YieldAwait functor that allows both. The mine coroutine above uses the YieldAwait functor to await a level and yield the updated level. I will give a small overview of common suspension functors and an example how they are used.

data Yield x y = Yield x y

yield :: Monad m => x -> Coroutine (Yield x) m ()
yield x = Coroutine (return (Left (Yield x (return ()))))

data Tree a = Tree a [Tree a]

preOrder :: Monad m => Tree a -> Coroutine (Yield a) m ()
preOrder (Tree a cs) = do
    yield a
    mapM_ preOrder cs

foldr :: (a -> b -> b) -> b -> Coroutine (Yield a) Identity () -> b
foldr f b0 co =
    case runIdentity (resume co) of
        Left (Yield x k) -> foldr f (f x b0) k
        Right ()         -> b0
                
sum :: Tree Int -> Int
sum = foldr (+) 0 . preOrder

The coroutine in this particular example gives the advantage that the caller can decide at which speed to consume the elements of the tree. The caller could decide to consume two elements at a time, print some intermediate result, consume two more and skip the rest of the iteration. An ordinary fold-function doesn’t allow this kind of freedom.

data Await x y = Await (x -> y)

await :: Monad m => Coroutine (Await x) m x
await = Coroutine (return (Left (Await return)))

printer :: (Show a,Monad m) => Coroutine (Await a) m r
printer = do
    x <- await
    lift $ print x
    printer

($$) :: Coroutine (Yield a) m r -> Coroutine (Await a) m r' -> m ()
producer $$ consumer = do
    p <- resume producer
    c <- resume consumer
    case (p,c) of
        (Left (Yield x producer'), Left (Await k)) -> producer' $$ k x
        _                                          -> return ()

printElements tree = preOrder tree $$ printer

The Await functor allows a coroutine to await values and then decide what to do next. The example above defines printer, an awaiting coroutine that prints each element it receives. printElements uses the yielding preOrder coroutine from above to send each yielded element to the printer. If one of the two coroutines of the $$ operator terminates, the whole operation terminates.

Conclusion

Haskell is one of the few languages that allows the implementation of coroutines as a library and it feels like a first-class construct. Coroutines are often used in cases where the language doesn’t support multi-threading like Lua coroutines, Python generators and JavaScript generators.