Coroutines

Our second example is the implementation of lexy, now using the Cont r monad. Since this example involves performing I/O to print the strings produced by lexy on screen, we can't use the Cont monad. However, there also exists a transformer version of it:

newtype ContT r m a = Cont { runContT :: (a -> m r) -> m r }

Interestingly, for ContT to be a monad, we don't even need m to be a monad, and the implementations of return and (>>=) for ContT r m are identical to those for Cont r.

Here, we choose the monad ContT r IO, as we want to equip the IO monad with continuations.

Lexy.hs

import Control.Monad (forM_, when)
import Control.Monad.Cont (ContT, runContT, callCC)
import Control.Monad.IO.Class (liftIO)

main :: IO ()
main = runContT prog return

prog :: ContT r IO ()
prog = do
    str <- callCC (lexy 3)
    case str of
        Nothing           -> return ()
        Just (xs, resume) -> do
            liftIO $ putStrLn xs
            resume

lexy :: Int
     -> (Maybe (String, ContT r m ()) -> ContT r m ())
     -> ContT r m (Maybe (String, ContT r m ()))
lexy n yield = do
    callCC (\resume -> yield $ Just ("", resume ()))
    when (n > 0) $ do
        forM_ "01" $ \x -> do
            nxt <- callCC (lexy (n - 1))
            case nxt of
                Nothing           -> return ()
                Just (xs, resume) -> yield $ Just (x:xs, resume)
    return Nothing

The main challenge with this implementation is the rather daunting type of lexy. Let's ignore it for now and focus on how the implementation works.

prog requests the first string using callCC (lexy 3). lexy takes two arguments n and yield. n is the length of the strings to be generated. yield is the function lexy can call to yield a string back to its caller. Only lexy doesn't simply yield a string. It yields a value of type Maybe (String, Cont r ()). If lexy yields (actually returns, as we will see) Nothing, then that's the indication that we're done iterating. So prog exits itself in this case. If lexy yields Just a pair (xs, resume), then xs is the string produced by lexy, and resume is an action of type ContT r IO () that can be called to resume lexy. After being resumed, lexy will yield another value, that is, each time we call resume, this will make callCC return once more, with another string to be printed, until we finally obtain Nothing, at which point prog exits.

main simply runs prog, using return (the identity function of any monad) to extract the value returned by prog.

Now the implementation of lexy. We said already that n is the length of the strings to be generated, and yield is the function that lexy can call to return a result to callCC. We also discussed that this result should have type Maybe (String, ContT r m ()). Thus, n has the type Int, and yield has the type Maybe (String, Cont r m ()) -> Cont r m (). The argument type of yield must match the return type of lexy, so the return type of lexy is ContT r m (Maybe (String, Cont T r m ())). And that's where the daunting type signature came from. When using this type of code in production, we would probably introduce type aliases to make these type signatures more readable.

Back to the implementation of lexy. The first thing that lexy does is to call callCC (\resume -> yield $ Just ("", resume ())). Thus, it yields Just the pair ("", resume ()) back to the caller. In other words, it produces the empty string paired with the action that can be called to make this invocation of callCC return. When the caller invokes resume, lexy continues its computation.

If n == 0, then the when block does not run. So lexy returns Nothing and thereby signals to its caller that it is done producing values. That's certainly the correct behaviour for lexy 0.

If n > 0, then lexy iterates over the characters in the alphabet "01". For each such character x, it asks for strings of length n - 1 by calling callCC (lexy (n - 1)). If this produces Nothing, then lexy (n - 1) is done producing values. Thus, the current iteration of the forM_ "loop" terminates as well, and we either start a next iteration or, if this was the last character in "01", the when block terminates. At this point, lexy n returns Nothing and thereby signals that it too is done producing values. If callCC (lexy (n - 1)) produces Just a pair (xs, resume), then lexy n yields Just the pair (x:xs, resume) to its caller. The string is the current character x together with the string produced by lexy (n - 1). When the caller invokes resume, this directly resumes lexy (n - 1), which then yields the next string, which lexy n augments with the current character x, and lexy n once again yields the result back to the caller.

Let's try it all out:

GHCi

>>> :l Lexy.hs
[1 of 1] Compiling Main             ( Lexy.hs, interpreted )
Ok, one module loaded.
>>> main

0
00
000
001
01
010
011
1
10
100
101
11
110
111

Works like a charm.

Of course, thanks to lazy evaluation, none of the above is needed in Haskell. The amount of effort to implement lexy in CPS is quite perverse really. The following implementation is only slightly less efficient, but equally efficient as the Python version of lexy, and is much shorter than even the Python version:

lexy :: n -> [String]
lexy 0 = [""]
lexy n = "" : [x:xs | x <- "01", xs <- lexy (n - 1)]

GHCi

>>> :{
  | lexy 0 = [""]
  | lexy n = "" : [x:xs | x <- "01", xs <- lexy (n - 1)]
  | :}
>>> mapM_ putStrLn (lexy 3)

0
00
000
001
01
010
011
1
10
100
101
11
110
111

The type of this version of lexy suggests that we are producing an entire list of strings, exactly the problem we want to avoid by implementing a coroutine that produces these strings one at a time. However, lazy evaluation ensures that the strings in this list are produced one at a time, just as when using a coroutine. Each string is generated whenever mapM_ putStrLn asks for it to print it. The coroutine that produces these strings on the fly for us is the lazy evaluation engine built into Haskell's runtime system.