CPS Interpreter for STLC + shift/reset in OCaml

main.ml

type id = string

module Environment = Map.Make (
  struct
    type t = id
    let compare (x : id) y = compare x y
  end
)

type value =
  | IntV of int
  | BoolV of bool
  | FunV of (value -> cont -> value)
and cont = value -> value

let pp_value = function
  | IntV i -> string_of_int i
  | BoolV b -> string_of_bool b
  | FunV _ -> "<fun>"

type op = Plus | Mult | Lt

type term =
  | Var of id
  | IConst of int
  | BConst of bool
  | BinOp of op * term * term
  | FunExp of id * term
  | AppExp of term * term
  | ShiftExp of id * term
  | ResetExp of term

exception EvalError

let rec eval term env cont = match term with
  | Var x ->
      if Environment.mem x env then
        cont @@ Environment.find x env
      else
        raise EvalError
  | IConst i -> cont @@ IntV i
  | BConst b -> cont @@ BoolV b
  | BinOp (op, e1, e2) ->
    eval e1 env @@
      fun x1 -> eval e2 env @@
      fun x2 -> cont  begin
        match op, x1, x2 with
        | Plus, IntV x1, IntV x2 -> IntV (x1 + x2)
        | Mult, IntV x1, IntV x2 -> IntV (x1 * x2)
        | Lt, IntV x1, IntV x2 -> BoolV (x1 < x2)
        | _ -> raise EvalError
      end
  | FunExp (x, e) ->
      cont @@ FunV (fun v -> fun c -> eval e (Environment.add x v env) c)
  | AppExp (e1, e2) ->
      eval e1 env @@
        fun v1 -> eval e2 env @@
        fun v2 -> begin
          match v1 with
          | FunV f -> f v2 cont
          | _ -> raise EvalError
        end
  | ShiftExp (k, e) ->
      let env' = Environment.add k (FunV (fun v -> fun c -> c (cont v))) env in
      eval e env' @@ fun x -> x
  | ResetExp e ->
      cont @@ eval e env (fun x -> x)

let _ =
  let env = Environment.singleton "x" (IntV 1) in
  let test t =
    print_endline @@ pp_value @@ eval t env (fun x -> x) in
  test @@ IConst 3;
  test @@ BinOp (Plus, IConst 3, IConst 4);
  test @@ BinOp (Mult, IConst 3, IConst 4);
  test @@ BinOp (Lt, IConst 3, IConst 4);
  test @@ Var "x";
  test @@ FunExp ("x", Var "x");
  test @@ AppExp (FunExp ("x", Var "x"), IConst 111);
  test @@ ResetExp (AppExp (FunExp ("x", Var "x"), IConst 111));
  test @@ ResetExp (BinOp (Plus, IConst 3, ShiftExp ("k", IConst 4)));
  test @@ ResetExp (BinOp (Plus, IConst 3, ShiftExp ("k", AppExp (Var "k", AppExp (Var "k", IConst 4)))))