Implementing python style generators with 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
shift0 consists of the stack frames between the two (this also means you cannot call
shift0 anywhere outside the scope of
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
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"
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
shift have the same return type). Let’s play!
# !count;; - : int = 1 # k ();; - : t = More <fun> # k ();; - : t = More <fun> # !count;; - : int = 3
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.