continuations.ml

(* two ways of implementing factorial: one with unbounded control context and one with
   iterative control context *)

let rec fact_unbounded n =
  if n = 0 then 1 else n * (fact_unbounded (n - 1))

(* ocaml will issue a warning if you use the [@tailcall] decorator on a
   non-tail-call form *)

let rec fact_iter_acc (n:int) (acc:int) : int =
  if n = 0 then acc else ((fact_iter_acc [@tailcall]) (n - 1) (n * acc))

let fact_iter n = fact_iter_acc n 1

let rec fact_cont_h (n:int) (cont:int -> int) =
  if n = 0 then cont 1 else fact_cont_h (n-1) (fun result -> cont (n * result))

let fact_cont n = fact_cont_h n (fun result -> result)

(* Fibonacci sequence:
   a0 = 0
   a1 = 1
   an = a_{n-1} + a_{n - 2}*)

let rec fib a =
  if a = 0 then 0
  else if a = 1 then 1
  else fib (a-1) + fib (a-2)

(* continuation-passing-style fibonacci *)
let rec fib_iter_h a (cont : int -> int) =
  if a = 0 then cont 0
  else if a = 1 then cont 1
  else (fib_iter_h [@tailcall]) (a-1)
      (fun res -> (fib_iter_h [@tailcall]) (a-2)
          (fun x -> x + (cont res)))

let fib_iter a = fib_iter_h a (fun x -> x)

let () =
  assert (fib 5 = fib_iter 5);
  assert (fib 6 = fib_iter 6);

(* a continuation-passing interpreter for a tiny calculator language *)
type calc =
  Num of int
  | Add of calc * calc

(* naive implementation: unbounded control context *)

let rec interp_recursive c =
  match c with
  | Num(n) -> n
  | Add(c1, c2) -> (interp_recursive c1) + (interp_recursive c2)

(* continuation-passing form *)
let rec interp_iter_h (c:calc) (cont: int -> int) : int =
  match c with
  | Num(n) -> cont n
  | Add(c1, c2) ->
    interp_iter_h c1
      (fun n1 -> interp_iter_h c2 (fun n2 -> cont (n1 + n2)))

let interp_iter c =
  interp_iter_h c (fun n -> n)


(**********************************************************************************)
(* adding `return` *)

(* a continuation-passing interpreter for a tiny calculator language *)
type return_calc =
    Num of int
  | Add of return_calc * return_calc
  | Return of return_calc

(* continuation-passing form *)
let rec interp_return_h (c:return_calc) (cont: int -> int) : int =
  match c with
  | Num(n) -> cont n
  | Add(c1, c2) ->
    interp_return_h c1
      (fun n1 -> interp_return_h c2 (fun n2 -> cont (n1 + n2)))
  | Return (calc) ->
    (* discard the current continuation and simply run `calc` *)
    interp_return_h calc (fun x -> x)

let interp_return c =
  interp_return_h c (fun n -> n)

let () =
  let prog : return_calc = (Add(Num 10, Return(Num 40))) in
  assert (interp_return (Add(Num 10, Num 10)) = 20);
  assert (interp_return prog = 40);


(**********************************************************************************)
(* adding other types *)
type value =
  VBool of bool
  | VNum of int

type bcalc =
  Num of int
  | Bool of bool
  | And of bcalc * bcalc
  | Add of bcalc * bcalc

exception Runtime

let rec interp_bcalc_iter_h c (cont : value -> value) : value =
  match c with
  | Num(n) -> cont (VNum n)
  | Bool(b) -> cont (VBool b)
  | Add(e1, e2) ->
    interp_bcalc_iter_h e1
      (fun v1 -> interp_bcalc_iter_h e2 (fun v2 ->
         let sum = match (v1, v2) with
           | (VNum(n1), VNum(n2)) -> VNum(n1 + n2)
           | _ -> raise Runtime in
         cont sum))
  | And(e1, e2) ->
    interp_bcalc_iter_h e1
      (fun v1 -> interp_bcalc_iter_h e2 (fun v2 ->
           let sum = match (v1, v2) with
             | (VBool(n1), VBool(n2)) -> VBool(n1 && n2)
             | _ -> raise Runtime in
           cont sum))


(********************************************************************************)
(* adding return *)
type rcalc =
    Num of int
  | Bool of bool
  | And of rcalc * rcalc
  | Add of rcalc * rcalc
  | Return of rcalc

exception Runtime

let rec interp_rcalc_h (e:rcalc) (cont : value -> value) =
  match e with
  | Num(n) -> cont (VNum n)
  | Bool(b) -> cont (VBool b)
  | Add(e1, e2) ->
    interp_rcalc_h e1
      (fun v1 -> interp_rcalc_h e2 (fun v2 ->
           let sum = match (v1, v2) with
             | (VNum(n1), VNum(n2)) -> VNum(n1 + n2)
             | _ -> raise Runtime in
           cont sum))
  | And(e1, e2) ->
    interp_rcalc_h e1
      (fun v1 -> interp_rcalc_h e2 (fun v2 ->
           let sum = match (v1, v2) with
             | (VBool(n1), VBool(n2)) -> VBool(n1 && n2)
             | _ -> raise Runtime in
           cont sum))
  | Return(e) ->
    (* evaluate and return e *)
    interp_rcalc_h e (fun x -> x)

let interp_rcalc e = interp_rcalc_h e (fun x -> x)

let () =
  let prog : rcalc = (Add(Num 10, Return(Num 40))) in
  assert (interp_rcalc (Add(Num 10, Num 10)) = VNum(20));
  assert (interp_rcalc prog = VNum(40));

(********************************************************************************)
(* adding exceptions *)

type ecalc =
    Num of int
  | Bool of bool
  | And of ecalc * ecalc
  | Add of ecalc * ecalc
  | Try_with of {try_e: ecalc; with_e: ecalc}
  | Raise

exception Uncaught_exception

let rec interp_ecalc_iter_h (c:ecalc) (cont : value -> value) (handler: unit -> value) =
  match c with
  | Num(n) -> cont (VNum n)
  | Bool(b) -> cont (VBool b)
  | Add(e1, e2) ->
    interp_ecalc_iter_h e1
      (fun v1 -> interp_ecalc_iter_h e2 (fun v2 ->
           let sum = match (v1, v2) with
             | (VNum(n1), VNum(n2)) -> VNum(n1 + n2)
             | _ -> raise Runtime in
           cont sum) handler) handler
  | And(e1, e2) ->
    interp_ecalc_iter_h e1
      (fun v1 -> interp_ecalc_iter_h e2 (fun v2 ->
           let sum = match (v1, v2) with
             | (VBool(n1), VBool(n2)) -> VBool(n1 && n2)
             | _ -> raise Runtime in
           cont sum) handler) handler
  | Raise -> handler ()
  | Try_with { try_e; with_e } ->
    (* run try_e with a new continuation that evaluates with_e *)
    let handler_cont = fun () -> interp_ecalc_iter_h with_e cont handler in
    interp_ecalc_iter_h try_e cont handler_cont

let interp_ecalc (e:ecalc) =
  interp_ecalc_iter_h e (fun x -> x) (fun x -> raise Uncaught_exception)

let p1 = Try_with { try_e = Raise;
                    with_e = Num(10)}

let p2 = Try_with { try_e = Add(Num(10), Num(20));
                    with_e = Num 20}

let p3 = Try_with { try_e = Add(Num(10), Raise);
                    with_e = Num 20}

let p4 = Try_with { try_e = Try_with { try_e = Raise;
                                       with_e = Num 30} ;
                    with_e = Num 20}

let () =
  assert ((interp_ecalc p1) = VNum(10));
  assert ((interp_ecalc p2) = VNum(30));
  assert ((interp_ecalc p3) = VNum(20));
  assert ((interp_ecalc p4) = VNum(30))

(* adding let-bindings *)