(* Yann Regis-Gianas implementation of co-routines using call/cc in O'Caml. *)

(* We are using Xavier Leroy's experimental Callcc library for O'Caml. *)
open Callcc;;
let unsome = function Some x -> x | None -> assert false

module type S = sig
  (* Coroutines with input 'i and output 'o *)
  type ('i,'o) coroutine
  type ('i,'o) coroutine_body = ('i, 'o) coroutine -> 'i -> 'o
  (* create coroutine from body function *)
  val create : ('i, 'o) coroutine_body -> ('i,'o) coroutine
  (* Asymmetric coroutine operator: *)
  (* resume a suspended coroutine. *)
  val resume: ('i,'o) coroutine -> 'i -> 'o
  (* Complement to resume: suspend the current coroutine, return to caller *)
  val yield : ('i, 'o) coroutine -> 'o -> 'i
end

module Coroutine : S = struct

  (* A co-routine has to remember its own continuation
     when it is suspended as well as the continuation of
     its caller when it is resumed. *)
  type ('i, 'o) coroutine_data = {
    suspension : ('i cont) option ref;
    caller     : ('o cont) option ref;
    body       : 'i -> 'o
  }

  type ('i, 'o) coroutine =
      Coroutine of ('i, 'o) coroutine_data

  (* A co-routine is defined using a function that
     will ultimately produce a value of type ['o]
     from a value of type ['i]. This function can
     interrupt itself using [yield]. *)
  type ('i,'o) coroutine_body =
      ('i, 'o) coroutine -> 'i -> 'o

  let resume (Coroutine co) v =
    match !(co.suspension) with
      (* This is the first this co-routine is runned, so we
     directly use the initial function that was given
     at the creation of the co-routine. *)
      | None ->
    callcc (fun k -> co.caller := Some k; co.body v)

      (* The co-routine has already been suspended. It
     first captures the caller's continuation and
     reinstall its own suspended continuation. *)
      | Some suspension ->
    callcc (fun k -> co.caller := Some k; throw suspension v)

  let yield (Coroutine co) v =
    (* This time, this is the other around: we have to save the
       co-routine continuation and restore the caller's. *)
    callcc (fun k -> co.suspension := Some k; throw (unsome !(co.caller)) v)

  let create body =
    let rec c = Coroutine ({
      suspension = ref None;
      caller     = ref None;
      body       = (fun x -> body c x)
    })
    in
    c

end

open Coroutine;;

(* A first test:
   Two co-routines play ping-pong.
   Pong is the master. *)
let ping : (int, int) coroutine =
  let rec body = fun co x ->
    begin
      Printf.printf "ping %d !\n%!" x;
      body co (yield co (x + 1))
    end
  in
  create body

let max_ping_pong = 11

let pong : (int, int) coroutine =
  let rec body = fun co x ->
    if x < max_ping_pong then begin
      Printf.printf "pong %d !\n%!" x;
      body co (resume ping (x + 1));
    end else x
  in
  create body

let _ = resume pong 0

(* A second test:
   Now ping and pong play pingpong but symmetrically.
   This requires a mutually recursive definition of these co-routines,
   which is realized through a two-steps initialization through a
   reference.
*)

let ping pong : (int, int) coroutine =
  let rec body = fun co x ->
    if x < max_ping_pong then begin
      Printf.printf "ping %d !\n%!" x;
      body co (resume (unsome !pong) (x + 1))
    end else x
  in
  create body

let max_ping_pong = 11

let pong ping : (int, int) coroutine =
  let rec body = fun co x ->
    if x < max_ping_pong then begin
      Printf.printf "pong %d !\n%!" x;
      body co (resume (unsome !ping) (x + 1));
    end else x
  in
  create body

let pingpong =
  let rping = ref None in
  let rpong = ref None in
  let ping  = ping rpong in
  let pong  = pong rping in
  rping := Some ping;
  rpong := Some pong;
  resume ping 0

(* A third test:
   An iteration through an infinite tree.
*)

type 'a tree = Empty | Node of 'a tree * 'a * 'a tree

let fold_tree =
  fun f tree init ->
    let rec body =
      fun tree accu co () ->
    match tree with
      | Empty -> accu
      | Node (l, x, r) ->
        let accu = f accu x in
        yield co accu;
        let accu = body l accu co () in
        body r accu co ()
    in
    create (body tree init)

let rec infinite_tree = Node (infinite_tree, 42, infinite_tree)

let _ =
  let sum accu x = x + accu in
  let co_fold_tree = fold_tree sum infinite_tree 0 in
  for i = 0 to 100 do
    Printf.printf "Sum: %d\n" (resume co_fold_tree ())
  done