generator.ml

open Printf

module type TREE = sig
  type 'a t
  (** The type of tree. *)

  val leaf : 'a t
  (** A tree with only a leaf. *)

  val node : 'a t -> 'a -> 'a t -> 'a t
  (** [node l x r] constructs a new tree with a new node [x] as the value, with
      [l] and [r] being the left and right sub-trees. *)

  val deep : int -> int t
  (** [deep n] constructs a tree of depth n, in linear time, where every node at
      level [l] has value [l]. *)

  val to_iter : 'a t -> ('a -> unit) -> unit
  (** Iterator function. *)

  val to_gen : 'a t -> (unit -> 'a option)
  (** Generator function. [to_gen t] returns a generator function [g] for the
      tree that traverses the tree in depth-first fashion, returning [Some x]
      for each node when [g] is invoked. [g] returns [None] once the traversal
      is complete. *)

  val to_gen_cps : 'a t -> (unit -> 'a option)
  (** CPS version of the generator function. *)
end

module Tree : TREE = struct

  type 'a t =
  | Leaf
  | Node of 'a t * 'a * 'a t

  let leaf = Leaf

  let node l x r = Node (l,x,r)

  let rec deep = function
    | 0 -> Leaf
    | n -> let t = deep (n-1) in Node (t,n,t)

  let rec iter f = function
    | Leaf -> ()
    | Node (l, x, r) -> iter f l; f x; iter f r

  (* val to_iter : 'a t -> ('a -> unit) -> unit *)
  let to_iter t f = iter f t

  (* val to_gen : 'a t -> (unit -> 'a option) *)
  let to_gen (type a) (t : a t) =
    let module M = struct effect Next : a -> unit end in
    let open M in
    let step = ref (fun () -> assert false) in
    let first_step () =
      try
        iter (fun x -> perform (Next x)) t;
        None
      with effect (Next v) k ->
        step := continue k;
        Some v
    in
      step := first_step;
      fun () -> !step ()

  let to_gen_cps t =
    let next = ref t in
    let cont = ref Leaf in
    let rec iter t k = match t with
      | Leaf -> run k
      | Node (left, x, right) -> iter left (Node (k, x, right))
    and run = function
      | Leaf -> None
      | Node (k, x, right) ->
        next := right;
        cont := k;
        Some x
    in fun () -> iter !next !cont
end

let get_mean_sd l =
  let get_mean l = (List.fold_right (fun a v -> a +. v) l 0.) /.
               (float_of_int @@ List.length l)
  in
  let mean = get_mean l in
  let sd = get_mean @@ List.map (fun v -> abs_float (v -. mean) ** 2.) l in
  (mean, sd)

let benchmark f n =
  let rec run acc = function
  | 0 -> acc
  | n -> let t1 = Sys.time () in
         let () = f () in
         let d = Sys.time () -. t1 in
         run (d::acc) (n-1)
  in
  let r = run [] n in
  get_mean_sd r

(* Main follows *)

let n =
  try
    int_of_string (Sys.argv.(1))
  with
  | _ -> 25

let t = Tree.deep n

let iter_fun () = Tree.to_iter t (fun _ -> ())
let (m,sd) = benchmark iter_fun 5
let () = printf "Iter: mean = %f, sd = %f\n%!" m sd

let rec consume_all f =
  match f () with
  | None -> ()
  | Some _ -> consume_all f

let gen_cps_fun () =
  let f = Tree.to_gen_cps t in
  consume_all f
let (m,sd) = benchmark gen_cps_fun 5
let () = printf "Gen_cps: mean = %f, sd = %f\n%!" m sd

let gen_fun () =
  let f = Tree.to_gen t in
  consume_all f
let (m, sd) = benchmark gen_fun 5
let () = printf "Gen_eff: mean = %f, sd = %f\n%!" m sd