Implementing python style generators with delimited continuations

Jake’s article on implementing LWT-compatible fibers with Oleg’s delimited continuation library got me thinking about one of my favorite Python features: generators.

Delimited continuations

Delimited continuations have a very simple API (note: I’m going to consider only the high level API provided by delimcc here consistent with the abstract view of delimited continuations. Jake describes the low-level API in his article). The core API is:

val new_prompt   : unit -> 'a prompt
val push_prompt  : 'a prompt -> (unit -> 'a) -> 'a
val shift0       : 'a prompt -> (('b -> 'a) -> 'a) -> 'b

new_prompt simply creates the “handle”, push_prompt (traditionally called reset) sets the limit of the continuation (how far up the stack we will go), shift0 captures the continuation and gives it to the passed function. Crucially, the continuation captured by shift0 is delimited by the outer push_prompt. The “region” of the continuation delimited by push_prompt and shift0 consists of the stack frames between the two (this also means you cannot call shift0 anywhere outside the scope of push_prompt).

An example should clarify. Let’s first do something illegal.

# open Delimcc;;
# let p = new_prompt ();;
value p : Delimcc.prompt '_a = <abstr>
# push_prompt p identity;;
- : unit = ()
# shift0 p (fun k -> k ());;
Exception: Failure "No prompt was set".

Here, we tried to shift0 outside of the push_prompt. Nope, can’t do that.

# push_prompt p (fun () -> shift0 p (fun k -> k 123));;
- : int = 123

Here, the inner function (passed to shift0) was given the captured continuation k (as delimited by push_prompt), calling it with the argument 123 which was subsequently returned.

The neat thing about delimited continuations is that they capture the state entire delimited stack. That is, calling k ARG from within shift0 is equivalent to replacing the shift0 statement with ARG, returning the result of push_prompt. Furthermore, it’s a first class continuation, so we may use it multiple times. To wit:

# push_prompt p (fun () -> 10 * shift0 p (fun k -> k 1 + k 3));;
- : int = 40

Here, k is equivalent to the function ( * ) 10. The argument & return types needn’t be uniform:

# push_prompt p (fun () -> string_of_int (shift0 p (fun k ->  "hey " ^ k 123)));;
- : string = "hey 123"

Leaking continuations

What makes delimited continuations yet more interesting is that the continuations themselves can escape the scope of shift0, and be restarted from outside.

Because of the type signatures of shift0, we need to unify the types, so let’s first define an ADT to do so:

# type t = Done | More of (unit -> t);;

Note that we can’t just use option here, it seems, because we need a recursive type.

# let count = ref 0;;
# let More k = push_prompt p (fun () -> 
  while true do 
    count := !count + 1; 
    shift0 p (fun k -> More k)
  done; 
  Done);;

Here we are using our unified type trick to return the continuation itself (note how push_prompt and shift have the same return type). Let’s play!

# !count;;
- : int = 1
# k ();;
- : t = More <fun>
# k ();;
- : t = More <fun>
# !count;;
- : int = 3

Generators

We now have a pretty good idea of what’s needed to implement generators with an API similar to:

let producer () =
  let state = ... in
  gen begin fun yield ->
    yield state; ...; yield state; ...; yield state
  end

So we need a function gen that takes a “yielder” (in Python terminology) that needs to capture the continuation, communicate the value and resume execution when the next value is asked for by the consumer:

let consumer () = 
  let g = producer () in
  let v0 = g () in
  ...
  let vN = g ()

It turns out that this is a very natural fit, and we can implement full duplex generators in just a few lines, using the ideas developed above. Here is the full source code:

type ('a, 'b) t = Done | More of 'a * ('b -> ('a, 'b) t)

let gen f =
  (*
   * Note: the first value to yield gets thrown away as the generator
   * has not yet started.
   *)
  let start _ =
    let p = Delimcc.new_prompt () in
    Delimcc.push_prompt p begin fun () ->
      f (fun x -> Delimcc.shift0 p (fun k -> More (x, k))); Done
    end in
  let next = ref start in

  fun rv ->
    match !next rv with
      | More (x, k) -> next := k; Some x
      | Done        -> None

And sample use:

let rec take_all t count =
  match t count with
    | Some i -> printf "took: %d\n" i; take_all t (count + 1)
    | None   -> ()

let () =
  let take = gen begin fun yield ->
    for i = 0 to 10 do
      let rv = yield i in
      printf "got: %d\n" rv
    done
  end in
  take_all take 100

The only additional work we have to do is to type it for bi-directional communication, and to store the next continuation in a ref cell for the next invocation. Given that the generators developed here are full duplex, we can implement co-routines as outlined in PEP 342.

You can also get it as a gist here. I also highly recommend Jake’s article which explores delimcc from a lower level of abstraction.