python generators in ocaml using delimcc

gen.ml

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

(* ==EXAMPLE USE== *)
open Printf
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