(*		Generators in OCaml                       *)
(* Examples: in-order and post-order traversals of a tree
   and some of the examples from Icon's Overview.
*)

(* At the top-level, first do:
#load "generator.cmo";;
*)

open Generator;;


(* A few preliminaries: define the tree and build a sample one *)

type label = int
type tree = Leaf | Node of label * tree * tree;;

let make_full_tree depth =
  let rec loop label = function
    | 0 -> Leaf
    | n -> Node (label, loop (2*label) (pred n), loop (2*label+1) (pred n))
  in loop 1 depth
;;
(* val make_full_tree : int -> tree = <fun> *)

let tree1 = make_full_tree 3;;

(*
val tree1 : tree =
  Node (1, Node (2, Node (4, Leaf, Leaf), Node (5, Leaf, Leaf)),
           Node (3, Node (6, Leaf, Leaf), Node (7, Leaf, Leaf)))
*)

(* Declare the `exception' to carry the yielded value *)

let pint = new_prompt ();;


(* The traversal code itself *)
let rec in_order2 = function
  | Leaf -> ()
  | Node (label, left, right) ->
      in_order2 left;
      Printf.printf "traversing: %d " label;
      yield pint label;
      in_order2 right
;;
(* val in_order2 : tree -> unit = <fun> *)

(* The first test

 Do the complete traversal, printing the labels as they are yielded.
 The trace shows that we indeed yield as we visit.
*)

for_loop pint (fun () -> in_order2 tree1)
  (fun x -> Printf.printf "Generated: %d\n" x)
;;

(*
traversing: 4 Generated: 4
traversing: 2 Generated: 2
traversing: 5 Generated: 5
traversing: 1 Generated: 1
traversing: 6 Generated: 6
traversing: 3 Generated: 3
traversing: 7 Generated: 7
- : unit = ()
*)

(*
  Stopping the generator earlier: request only two generated values.
  The trace shows that we stop the traversal after consuming
  the needed two values.
  We indeed traverse on-demand.
*)

(* For convenience, we define a new loop `operator' *)
let for_loop_n p n gen f = stream_iter f (stream_take n (msplit p gen));;

for_loop_n pint 2 (fun () -> in_order2 tree1)
  (fun x -> Printf.printf "Generated: %d\n" x)
;;

(*
traversing: 4 Generated: 4
traversing: 2 Generated: 2
- : unit = ()
*)


(* More challenging and interesting is the post_order example,
  an accumulating traversal yielding intermediate results.

  We traverse a tree post-order, yielding the running sum of
  the current label and the labels in the left and the right
  branches (that is, the sum of the labels for the subtree rooted 
  at the current node).

 Now the generator has to return a useful value.
*)

let rec post_order = function
  | Leaf -> 0
  | Node (label, left, right) ->
      let sum_left  = post_order left in
      let sum_right = post_order right in
      let sum = sum_left + sum_right + label in
      yield pint sum; sum
;;
(* val post_order : tree -> int = <fun> *)

for_loop pint (fun () -> Printf.printf "Final: %d\n" (post_order tree1))
   (fun x -> Printf.printf "Got %d\n" x)
;;

(*
Got 4
Got 5
Got 11
Got 6
Got 7
Got 16
Got 28
Final: 28
- : unit = ()
*)


(* Re-implementing Icon's example:
     procedure findodd(s1, s2)
       every i := find(s1, s2) do
          if i % 2 = 1 then suspend i
     end
*)


let sentence = "Store it in the neighboring harbor";;

(* Find all occurrences and accumulate them a list *)
let [2;22;32] = stream_mapl (fun x -> x) (msplit pfind (fun () ->
  find "or" sentence));;

let findodd s1 s2 =
  for_loop pfind (fun () -> find s1 s2)
  (fun i -> if i mod 2 == 1 then yield pfind i)
;;

(* The test. It demonstrates nested loop *)
for_loop pfind (fun () -> findodd "a" "abracadabra")
 (fun i -> Printf.printf "found index %d\n" i)
;;


(*
found index 3
found index 5
found index 7
- : unit = ()
*)


(* Running generators side-by-side

  We solve the famous same-fringe problem, checking to see if two binary 
  trees have the same fringe (that is, the same sequence of labels when 
  traversed in a particular order).
  We must stop the traversal as soon as the mismatch is found,
  returning it. To demonstrate this property, we print labels
  as they are reached by the traversal. 
*)


(* For generality, we define a zipWith-like function for streams.
  To handle the case when two streams have different lengths,
  the zipping function receives 'a option values (with None
 signifying the end of the stream).
*)

let singleton v : 'a stream = Cons (v,fun () -> Nil);;

let rec stream_zip_with f s1 s2 =
  match (s1,s2) with
  | (Nil, Nil)        -> singleton (f None None)
  | (Cons (v1,_),Nil) -> singleton (f (Some v1) None)
  | (Nil,Cons (v2,_)) -> singleton (f None (Some v2))
  | (Cons (v1,t1), Cons (v2,t2)) -> 
      Cons (f (Some v1) (Some v2), 
	    fun () -> stream_zip_with f (t1 ()) (t2 ()))
;;

(*
val stream_zip_with :
  ('a option -> 'b option -> 'c) ->
  'a Generator.stream -> 'b Generator.stream -> 'c Generator.stream = <fun>
*)

(* We may also define the parallel loop directly, whose
   body does not return any value. It may yield.
*)

let for_loop_2 p gen1 gen2 body : unit =
  let rec loop = function
    | (Nil,Nil)         -> body None None
    | (Cons (v1,_),Nil) -> body (Some v1) None
    | (Nil,Cons (v2,_)) -> body None (Some v2)
    | (Cons (v1,t1), Cons (v2,t2)) -> 
	body (Some v1) (Some v2); loop (t1(),t2())
  in loop (msplit p gen1, msplit p gen2)
;;

(*
val for_loop_2 :
  'a Generator.stream Delimcc.prompt ->
  (unit -> unit) -> (unit -> unit) -> ('a option -> 'a option -> unit) -> unit
*)


let pmismatch : ('a option * 'a option) stream Delimcc.prompt
    = new_prompt ();;

let same_fringe p gen1 gen2 =
  for_loop_2 p gen1 gen2 (fun x y ->
    match (x,y) with
    | (None,_) | (_,None) as r              -> yield pmismatch r
    | (Some x,Some y) as r when not (x = y) -> yield pmismatch r
    | _  -> ())
;;

(* Another loop combinator: returning at most n results of the generator *)
let bagof_n n p gen = 
  stream_mapl (fun x -> x) (stream_take n (msplit p gen));;

(*
val bagof_n :
  int -> 'a Generator.stream Delimcc.prompt -> (unit -> unit) -> 'a list =
*)

(*
 First test: comparing the identical trees
 We have to traverse them completely. The trace shows the traversal
 done side-by-side
*)

let [(None, None)] = bagof_n 1 pmismatch 
 (fun () -> same_fringe pint 
     (fun () -> in_order2 tree1)
     (fun () -> in_order2 tree1))
;;

(* Comparing two different complete trees in-order.
   The trees differ already in the first label.
   The trace demonstrates the traversal is finished immediately.
*)

let [(Some 4, Some 8)] = bagof_n 1 pmismatch 
 (fun () -> same_fringe pint 
     (fun () -> in_order2 tree1)
     (fun () -> in_order2 (make_full_tree 4)))
;;

(*
        traversing: 8 traversing: 4
*)


let () = Printf.printf "\nAll done\n";;