Coroutines

As a final example of programming in CPS, let us replicate the implementation of coroutines via continuations that we illustrated using this piece of code before:

(define (lexy n yield)
  (call/cc (lambda (continue) (yield (cons '() continue))))
  (unless (zero? n)
    (map (lambda (c)
           (let ((result (call/cc (lambda (yield) (lexy (- n 1) yield)))))
             (when (car result)
               (yield (cons (cons c (car result)) (cdr result))))))
         '(#\0 #\1)))
  (cons #f #f))


(let ((result (call/cc (lambda (yield) (lexy 3 yield)))))
  (when (car result)
    (display (list->string (car result)))
    (newline)
    ((cdr result))))

In the Scheme code, lexy signals to the caller that it is done producing values by returning a pair whose first component is #f. Otherwise, the first component of this pair is a list of characters. This works in Scheme because it is dynamically typed. In Haskell, lexy needs to return a value of the same type no matter whether it actually produces a value or it has no more values to produce. The natural way to represent this is that lexy yields Maybe (String, r). This value is Nothing when lexy is done iterating. Otherwise, it is Just a pair composed of the string produced by lexy and a continuation of type r that we can invoke to resume lexy. To allow lexy to yield such a value, lexy needs to be given a continuation which it can call with this value. This makes lexy have the following type:

lexy :: Int -> (forall r. (Maybe (String, r) -> r) -> r)

Here, it isn't as essential to unpack the function of type Maybe a -> r into a pair of functions of type r and a -> r, as we did in our implementation of mapCPS, but there is a certain elegance to having lexy choose the right continuation directly instead of signalling its intention using a value of type Maybe String, and then the continuation implements a case distinction based on this value. So let's replace the function of type Maybe (String, r) -> r with two functions of type r and (String, r) -> r or, unpacking the pair (String, r) with two functions of type r and String -> r -> r.

lexy :: Int -> (forall r. r -> (String -> r -> r) -> r)

The driver code of the Scheme program is the let block in the last five lines of that code. This is essentially a loop. Every time lexy yields a value, the body of this loop should run. At the end of each iteration, we invoke the continuation returned by lexy to resume it. When lexy is done iterating, the main code should exit.

In Haskell, printing strings on screen requires the IO monad. So it is natural to choose r = IO (). We make the main function exit by having lexy invoke the continuation return (). The continuation that lexy should use to yield a value is a function loop that implements the loop body:

main :: IO ()
main = lexy 3 (return ()) loop
  where
    loop str resume = do
        putStrLn str
        resume

Here's the skeleton for implementing lexy:

lexy :: Int -> r -> (String -> r -> r) -> r
lexy n done yield = ...

n is the maximum length of the strings to be produced by lexy. done is the continuation that lexy should invoke to signal that it's done. yield is the continuation it uses to yield a string and a continuation of type r back to its caller.

The first thing that lexy should always do no matter the value of n is to yield the empty string and a continuation that allows lexy to be resumed. We implement this continuation using a function lexy':

lexy :: Int -> r -> (String -> r -> r) -> r
lexy n done yield = yield "" lexy'
  where
    lexy' = ...

If n == 0, then lexy shouldn't do anything after yielding the empty string. Thus, lexy' 0 simply invokes done to signal that it's done:

lexy :: Int -> r -> (String -> r -> r) -> r
lexy n done yield = yield "" lexy'
  where
    lexy' = case n of
        0 -> done
        _ -> ...

For all other values of n, we want to start a loop that iterates over all characters in the alphabet "01". Since that's the outer loop of a pair of nested loops, we call it outerLoop and call it with the alphabet to be iterated over:

lexy :: Int -> r -> (String -> r -> r) -> r
lexy n done yield = yield "" lexy'
  where
    lexy' = case n of
        0 -> done
        _ -> outerLoop "01"

outerLoop implements this iteration using tail-recursion. In general, the argument of outerLoop is the list of characters in the alphabet we haven't iterated over yet. Thus, once that list of characters is empty, the loop is done, that is, lexy is done. We invoke the done continuation to signal this:

lexy :: Int -> r -> (String -> r -> r) -> r
lexy n done yield = yield "" lexy'
  where
    lexy' = case n of
        0 -> done
        _ -> outerLoop "01"

    outerLoop "" = done

If there are still characters to iterate over, we split the list of characters into the first character in this list and the list of the remaining characters:

lexy :: Int -> r -> (String -> r -> r) -> r
lexy n done yield = yield "" lexy'
  where
    lexy' = case n of
        0 -> done
        _ -> outerLoop "01"

    outerLoop ""     = done
    outerLoop (x:xs) = ...

x is the character of interest in the current iteration. What the current iteration should do is to iterate over the strings yielded by lexy (n - 1). For each such string str, we yield x:str and the resume continuation yielded by lexy (n - 1) back to the caller of lexy n. We implement this logic in a function innerLoop that implements the body of this inner loop:

lexy :: Int -> r -> (String -> r -> r) -> r
lexy n done yield = yield "" lexy'
  where
    lexy' = case n of
        0 -> done
        _ -> outerLoop "01"

    outerLoop ""     = done
    outerLoop (x:xs) = lexy (n - 1) ??? innerLoop
      where
        innerLoop str resume = yield (x:str) resume

What should the ??? be? That's the continuation that lexy (n - 1) should invoke to signal that it is done iterating. When that's the case, we need to start another iteration of outerLoop, which we do by invoking outerLoop with the remaining characters in the alphabet, xs:

lexy :: Int -> r -> (String -> r -> r) -> r
lexy n done yield = yield "" lexy'
  where
    lexy' = case n of
        0 -> done
        _ -> outerLoop "01"

    outerLoop ""     = done
    outerLoop (x:xs) = lexy (n - 1) (outerLoop xs) innerLoop
      where
        innerLoop str resume = yield (x:str) resume

And with that, we have a Haskell version of lexy implemented in CPS, just as in Scheme. Let's try it out:

Lexy.hs

{-# LANGUAGE RankNTypes #-}

main :: IO ()
main = lexy 3 (return ()) loop
  where
    loop str resume = do
        putStrLn str
        resume

lexy :: Int -> (forall r. r -> (String -> r -> r) -> r)
lexy n done yield = yield "" lexy'
    where
    lexy' = case n of
        0 -> done
        _ -> outerLoop "01"

    outerLoop ""     = done
    outerLoop (x:xs) = lexy (n - 1) (outerLoop xs) innerLoop
        where
        innerLoop str resume = yield (x:str) resume

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

This was quite a journey. Programming in CPS is not as unnatural as it may seem at first, but we would surely prefer to program using standard functions that return results instead of having to pass a continuation to every function, and then the function calls the continuation for us. In Scheme, we can do this while also using continuations when we need them. In order to the same in Haskell, we need a monad. That's what we discuss next.