(* CPS and tail-calls *)

(* disclaimer: code has not been thoroughly tested *)

(* Assume you are in a setting where very deep call-stacks lead to run-time
   errors, so you don't want recursion proportional to data-structure depth
   since data structures might be large *)

(* list example: could cause stack overflow for large list *)

let rec length1 xs =
  match xs with
    [] -> 0
  | _::tl -> 1 + length1 tl

(* list example: by using an /accumulator/, the recursive call is a /tail call/
   so functional-language implementations ensure O(1) stack depth 
   -- a common pattern, using a helper function to hide from callers
*)

let length2 xs =
  let rec f acc xs =
    match xs with
      [] -> acc
    | _::tl -> f (acc+1) tl
  in f 0 xs

(* length2 is the best way to do it (fast, correct, straightforward) but 
   it relies on commutativity of addition and "order of list doesn't matter" *)

(* this version just mechanically does CPS instead: *)

let length3 xs =
  let rec f xs k =
    match xs with 
      [] -> k 0
    | _::tl -> f tl (fun i -> k (i+1))
  in f xs (fun x -> x)

(* tree example: if the tree is /balanced/, no problem, but if imbalance is
   possible, this could cause stack overflow *)

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

let rec size1 tree =
  match tree with
    Null -> 0
  | Node(_,l,r) -> 1 + size1 l + size1 r

(* tree example: we can try to use an accumulator, but it won't suffice because
   only one of the two recursive calls can be in tail position *)

let size2 tree =
    let rec f acc tree =
      match tree with
	Null -> acc
      |	Node(_,l,r) -> f (f acc l) r
    in f 0 tree

(* tree example: a tree traversal /requires/ a stack, but we can use an
   explicit heap-allocated stack instead of the meta-language's call-stack *)

let size3 tree =
  let rec f acc stack tree =
    match tree with
      Null -> (match stack with
	         [] -> acc
               | hd::tl -> f acc tl hd)
    | Node(_,l,r) -> f (acc+1) (r::stack) l
  in f 0 [] tree

(* Manual CPS has the same effect using the entirely mechanicaml CPS approach
   rather than the special-case of a linked-list of right-children *)

let size4 tree =
  let rec f tree k = (* f : 'a tree -> (int -> 'b) -> 'b *)
    match tree with
      Null -> k 0
    | Node(_,l,r) -> f l (fun i -> f r (fun j -> k (i+j+1)))
  in f tree (fun x -> x)

(* tree example: while size4 uses only tail calls, it allocates two
closures for each node.  Re-introducing an accumulator cuts this in half --
an accumulator needs no heap allocation *)

let size5 tree =
  let rec f tree acc k =
    match tree with
      Null -> k acc
    | Node(_,l,r) -> f l acc (fun i -> f r (1+acc) k)
  in f tree 0 (fun x -> x)

(* here is tree-product in CPS (not bothering with accumulator for simplicity) *)

let treeprod1 tree =
  let rec f tree k = (* f : int tree -> (int -> 'b) -> 'b *)
    match tree with
      Null -> k 1
    | Node(v,l,r) -> f l (fun i -> f r (fun j -> k (v * i * j)))
  in f tree (fun x -> x)

(* here is a version using a "fast explicit throw" for a 0 *)
let treeprod2 tree =
  let f1 tree k k2 = (* f : int tree -> (int -> 'b) -> (int -> 'b) -> 'b *)
    let rec f2 tree k =
    match tree with
      Null -> k 1
    | Node(v,l,r) -> if v=0
                     then k2 0
                     else f2 l (fun i -> f2 r (fun j -> k (v * i * j)))
    in f2 tree k
  in f1 tree (fun x -> x) (fun x -> x)

(* but actually there is no need for the explicit second continuation
   if we remember we didn't convert the entire program into CPS *)

let treeprod3 tree =
  let rec f tree k = (* f : int tree -> (int -> 'b) -> 'b *)
    match tree with
      Null -> k 1
    | Node(v,l,r) -> if v=0
                     then 0
                     else f l (fun i -> f r (fun j -> k (v * i * j)))
  in f tree (fun x -> x)

(* note this is also what exceptions "do" even if often considered poor-style
   (and it's not all that efficient) *)

exception Zero
let treeprod4 tree =
  let rec f tree = (* int tree -> int *)
    match tree with
      Null -> 1
    | Node(v,l,r) -> 
      if v=0
      then raise Zero
      else v * f l * f r
  in
  try f tree with
    Zero -> 0

let t1 = 
  Node(3,
       Node(4,
            Null,
            Node(5,Null,Null)),
       Node(2,
            Node(5,Null,Null),
            Null))

let t2 = 
    Node(3,
       Node(4,
            Null,
            Node(0,Null,Null)),
       Node(2,
            Node(5,Null,Null),
            Null))


let x1 = [treeprod1 t1; treeprod2 t1; treeprod3 t1; treeprod4 t1]
let x2 = [treeprod1 t2; treeprod2 t2; treeprod3 t2; treeprod4 t2]