Yield, Continue

Roshan P. James and Amr Sabry show in "Yield: Mainstream Delimited Continuations" the interdefinability of yield-style generators and delimited continuations. Their encoding is at the same time simple and general, and even if the examples given in the paper are in Haskell, their translation into OCaml is straightforward. So much so that the result is essentially equivalent to ASAI Kenichi's OchaCaml (Edit: this claim of mine is certainly unsubstantiated and quite possibly wrong. See Zerny's comment).

James and Sabry generalize the mechanics of yield to a three-ported construct represented by the type (ι, ο, ρ) Yield:

This type encapsulates the communication between an iterator and its calling context, where the iterator yields values of type ο, receives inputs of type ι and terminates (or returns) with a final result of type ρ. This communication is mediated by a delimited context that can be activated with run which … marks the boundary of an iterator and delimits the action of yield. This communication is effected by a reified continuation given by a concrete data type with which the calling context can interact:

type ('i, 'o, 'r) iterator =
| Result of 'r
| Susp of 'o * ('i -> ('i, 'o, 'r) iterator)

In effect, run converts a monadic producer that uses yield into a CPS-transformed consumer that invokes the continuation given by an iterator. These patterns of interaction can be abstracted somewhat. The most general consumer is given by foreach:

let rec foreach (m : ('i, 'o, 'r) iterator) (f : 'o -> 'i) : 'r = match m with
| Susp (v, k) -> foreach (k (f v)) f
| Result r    -> r

It applies f to each value yielded by the iterator, feeding the result back to it. If the consumer is interested just in the yielded values and not in the result of the iteration, it can fold over them:

let rec fold (m : ('i, 'i, 'r) iterator) (f : 'i -> 'j -> 'j) (e : 'j) : 'j = match m with
| Susp (v, k) -> f v (fold (k v) f e)
| Result _    -> e

The essence of the iterator is given by an abstract signature:

module type YIELD = sig
  type ('i, 'o, 'r) yield
  val return : 'r -> ('i, 'o, 'r) yield
  val (>>=) : ('i, 'o, 'r) yield -> ('r -> ('i, 'o, 's) yield) -> ('i, 'o, 's) yield
  val yield : 'o -> ('i, 'o, 'i) yield
  val run : ('i, 'o, 'r) yield -> ('i, 'o, 'r) iterator
end

which gives a multi-parameter monad together with a pair of operations: yield, that returns a computation returning the yielded value (note the difference with return); and run, that captures the computation's context and reifies it into an iterator. The paper gives two possible implementations. The first "grabs" each invocation frame turning it directly into an iterator:

module FrameGrabbingYield : YIELD = struct
  type ('i, 'o, 'r) yield = ('i, 'o, 'r) iterator

  let return e = Result e

  let rec (>>=) m f = match m with
  | Result v    -> f v
  | Susp (v, k) -> Susp (v, fun x -> k x >>= f)

  let yield v = Susp (v, return)

  let run e = e
end

The seconds uses the CPS-encoded delimited continuation monad directly:

module CPSYield : YIELD = struct
  type ('i, 'o, 'r) yield =
    { cont : 'b . ('r -> ('i, 'o, 'b) iterator) -> ('i, 'o, 'b) iterator }

  let return x = { cont = fun k -> k x }

  let (>>=) m f = { cont = fun k -> m.cont (fun x -> (f x).cont k) }

  let yield v = { cont = fun k -> Susp (v, k) }

  let run e = e.cont (fun r -> Result r)
end

This is the standard CPS monad with answer type polymorphism, as given by Kiselyov. Now yield e is shift $ return . Susp e, and run e is equivalent to reset $ e >>= return . Result). This is sufficient but a bit bare-bones. Let's build from here:

module YieldT (Y : YIELD) = struct
  include Y

In the simplest case, generators simply yield successive values. The result of the computation is the value itself, that can be updated for the next cycle:

let rec repeat x = yield x >>= repeat
  let rec from i = yield i >>= fun j -> from (succ j)

Transformers are a bit more involved in that they must consume the iterator and yield new values, in effect delimiting the control of the iterator they consume. The simplest transformer is obviously map:

let rec map f y =
    let rec go = function
    | Result r    -> return r
    | Susp (v, k) -> yield (f v) >>= fun _ -> go (k v)
    in go (run y)

(Note that the monadic fmap would only act on the result and not on the generated values.) In this case, the result of the computation is the mapped value of the original iterator, that must be continued with the original value. Truncating an iterator is straightforward:

let rec take n y =
    let rec go n = function
    | Result r -> return (Some r)
    | Susp (_, _) when n = 0 -> return None
    | Susp (v, k) -> yield v >>= fun x -> go (n - 1) (k x)
    in go n (run y)

Combining two generators is also straightforward:

let zip y1 y2 =
    let rec go = function
    | Result r1, Result r2 -> return (r1, r2)
    | Susp (v1, k1), Susp (v2, k2) ->
      yield (v1, v2) >>= fun (x1, x2) -> go (k1 x1, k2 x2)
    | _ -> failwith "zip"
    in go (run y1, run y2)
end

(Iterators that return early must be dealt with in a somewhat arbitrary way.) With this it is relatively straightforward to use iterators:

let ex1 y =
  let module Y = YieldT( (val y : YIELD) ) in
  foreach Y.(run (take 10 (map succ (from 0)))) (Printf.printf "%d ")

And both implementations give equivalent results:

# let _ = ex1 (module FrameGrabbingYield : YIELD) ;;
1 2 3 4 5 6 7 8 9 10 - : 'a option = None
# let _ = ex1 (module CPSYield : YIELD) ;;
1 2 3 4 5 6 7 8 9 10 - : 'a option = None

Furthermore, Asai's examples (as given in this Reddit thread) can be easily duplicated as well:

module Tree (Y : YIELD) = struct
  type 'a tree = E | N of 'a tree * 'a * 'a tree

  open Y

  let rec depth_walk : 'a tree -> ('b, 'a, 'b tree) yield = function
  | N (l, n, r) ->
    depth_walk l >>= fun l' ->
    yield n      >>= fun n' ->
    depth_walk r >>= fun r' ->
    return (N (l', n', r'))
  | E -> return E

  let to_list t = fold (run (depth_walk t)) (fun x xs -> x :: xs) []

  let map f t = foreach (run (depth_walk t)) f

  let samefringe l r =
    let rec visit l r = match l, r with
    | Result _, Result _ -> true
    | Susp (a, ka), Susp (b, kb)
      when a = b ->
      visit (ka a) (kb b)
    | _ -> false
    in visit (run (depth_walk l)) (run (depth_walk r))

  let swap l r =
    let rec visit l r = match l, r with
    | Susp (a, ka), Susp (b, kb) -> visit (ka b) (kb a)
    | Result t1, Result t2 -> (t1, t2)
    | _ -> failwith "Unequal number of leaves"
    in visit (run (depth_walk l)) (run (depth_walk r))
end

Note that, except for the return type polymorphism, these versions are exactly the same. To prove that all works properly, here are a number of tests:

# module T = Tree(CPSYield) ;;
# open T ;;

# let t1 = N (N (E, 10, E), 20, N (E, 30, N (E, 40, E)))
and t2 = N (N (E, 10, N (E, 20, E)), 30, N (E, 40, E))
and t3 = N (N (E, 'a', N (E, 'b', E)), 'c', N (E, 'd', E))
;;

(I omit the output of these for clarity.)

# let _ = map succ t1 ;;
- : int T.tree = N (N (E, 11, E), 21, N (E, 31, N (E, 41, E)))
# let _ = to_list t1 ;;
- : int list = [10; 20; 30; 40]
# let _ = samefringe t1 t2 ;;
- : bool = true
# let _ = swap t1 t3 ;;
- : char T.tree * int T.tree =
(N (N (E, 'a', E), 'b', N (E, 'c', N (E, 'd', E))),
 N (N (E, 10, N (E, 20, E)), 30, N (E, 40, E)))

Note that in the last example the trees retain their respective shapes but interchange the values of their leaves.