cps.ml

(* adding exceptions into the lambda calculus *)

type exn_expr =
  Lam of string * exn_expr
  | Var of string
  | App of exn_expr * exn_expr
  | Num of int
  | Add of exn_expr * exn_expr
  | Try_with of { try_e: exn_expr; handle_e: exn_expr }
  | Raise

type expr =
  Lam of string * expr
  | Var of string
  | App of expr * expr
  | Num of int
  | Add of expr * expr

module StringMap = Map.Make(String)

type value =
    VLam of string * expr * (value StringMap.t)
  | VNum of int

type env = value StringMap.t

exception Runtime

let rec interp_exp (e:expr) (env:env) : value =
  match e with
  | Num(n) -> VNum(n)
  | Var(s) ->
    (match StringMap.find_opt s env with
     | Some(v) -> v
     | None -> raise Runtime)
  | Lam(id, arg) -> VLam(id, arg, env)
  | App(e1, e2) ->
    let e2v = interp_exp e2 env in
    (match interp_exp e1 env with
     | VLam(arg, body, closure) ->
       let new_env = StringMap.add arg e2v closure in
       interp_exp body new_env
     | _ -> raise Runtime)
  | Add(e1, e2) ->
    (match ((interp_exp e1 env),(interp_exp e2 env)) with
     | (VNum(n1), VNum(n2)) -> VNum(n1 + n2)
     | _ -> raise Runtime )

let counter = ref 0

let fresh () : string =
  let cur = !counter in
  counter := !counter + 1;
  Format.sprintf "$%d" cur

let rec compile_k_h (e:exn_expr) (default_cont:expr) (handle_cont:expr) : expr =
  match e with
  | Num(n) -> App(default_cont, Num(n))
  | Lam(id, body) ->
    (* build a new lambda abstraction that takes the default and handle
       continuations as arguments *)
    let default_name = fresh () in
    let handle_name = fresh () in
    let compiled_body = compile_k_h body (Var(default_name)) (Var(handle_name)) in
    App(default_cont, Lam(default_name, Lam(handle_name, Lam(id, compiled_body))))
  | Var(s) -> App(default_cont, Var(s))
  | App(e1, e2) ->
    let fun_name = fresh () in
    let arg_name = fresh () in
    let call_function = compile_k_h e1
        (Lam(arg_name,
             App(App(App(Var(fun_name), default_cont), handle_cont), Var(arg_name)))) handle_cont in
    compile_k_h e2 (Lam(fun_name, call_function)) handle_cont
  | Add(e1, e2) ->
    let ret1 = fresh () in
    let ret2 = fresh () in
    let perform_sum = compile_k_h e2 (Lam(ret2, Add(Var(ret1), Var(ret2)))) handle_cont in
    compile_k_h e1 (Lam(ret1, perform_sum)) handle_cont
  | Try_with { try_e; handle_e } ->
    let compiled_handle = compile_k_h handle_e default_cont handle_cont in
    compile_k_h try_e default_cont (Lam("_", compiled_handle))
  | Raise -> App(handle_cont, Num(0))

let compile_k (e:exn_expr) : expr =
  compile_k_h e (Lam("x", Var("x"))) (Lam("x", Num(0)))

let compile_and_run (e:exn_expr) : value =
  let compiled = compile_k e in
  interp_exp compiled (StringMap.empty)

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

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

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

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

(* test a simple case with no exceptions *)
let p5 : exn_expr = App(Lam("x", Var("x")), Num(10))

(* test a simple ase with exceptions *)
let p6 = Try_with { try_e = App(Lam("x", Raise), Num(10));
                    handle_e = Num 20}

(* test to make sure no exceptions are raised inside lambda *)
let p7 : exn_expr = Lam("x", Raise)

(* test handling inside a lambda *)
let p8 : exn_expr = App(Lam("x", Try_with { try_e =  Raise; handle_e = Var("x")}), Num 40)

(* test some basic nesting *)
let p9 : exn_expr = Try_with { try_e = p8; handle_e =  Num 80 }