Control operators: delimited continuations in OCaml (chapter 8)

open Delimcc
open Printf

(* Section 8.7 *)

(* Control inversion on an iterator *)

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

let rec tree_iter f t =
  match t with
  | Leaf -> ()
  | Node(l, x, r) -> tree_iter f l; f x; tree_iter f r

let my_tree = Node(Node(Leaf, 1, Node(Leaf, 2, Leaf)), 3, Node(Leaf, 4, Leaf))

let rec print_tree (t: int tree) : unit =
  match t with
  | Leaf -> ()
  | Node(l, x, r) -> printf "["; print_tree l; printf "%d" x; print_tree r; printf "]"

type 'a enum = Done | More of 'a * (unit -> 'a enum)

let tree_enum (t: 'a tree) : 'a enum =
  let p = new_prompt() in
  push_prompt p begin fun () ->
    tree_iter
      (fun x -> shift p (fun k -> More(x, k)))
      t;
    Done
  end

let rec iter_on_enum f = function
  | Done -> ()
  | More(x, k) -> f x; iter_on_enum f (k ())

let _ =
  printf "Tree enumeration:";
  iter_on_enum (fun x -> printf " %d" x) (tree_enum my_tree);
  printf "\n"

(* Choice points and backtracking *)

let p = new_prompt()

let choose xs =
  shift p (fun k -> List.iter k xs)

let fail () = shift p (fun k -> ())

let assert_ b = if not b then fail ()

let _ =
  push_prompt p begin fun () ->
    let a = choose [1;2;3;4;5;6;7] in
    let b = choose [1;2;3;4;5;6;7] in
    let c = choose [1;2;3;4;5;6;7] in
    assert_ (c * c = a * a + b * b);
    assert_ (b < a);
    printf "Right triangle found: %d %d %d\n" a b c
  end

(* Generating and counting *)

let p : int prompt = new_prompt()

let bool () = shift p (fun k -> k false + k true)

let int lo hi =
  shift p (fun k ->
    let rec sum i =
      if i <= hi then k i + sum (i + 1) else 0
    in sum lo)

let _ =
  let n =
    push_prompt p begin fun () ->
      let d1 = int 1 6 and d2 = int 1 6 and d3 = int 1 6
      in if d1 + d2 + d3 >= 16 then 1 else 0
    end in
  printf "Number of 3d6 that sum to 16: %d\n" n

let rec bintree h =
  shift p (fun k ->
    if h < 0 then 0 else
    if h = 0 then k Leaf else
    if bool ()
    then k (Node(bintree (h-1), h, bintree (int 0 (h-1))))
    else k (Node(bintree (int 0 (h-2)), h, bintree (h-1))))

let _ =
  let n =
    push_prompt p begin fun () ->
      let t = bintree 3 in
      print_tree t; printf "\n";
      1
    end in
  printf "Number of binary trees of height 3: %d\n" n