open Syntax
open Value

(* defunctionalize continuation *)

(* initial continuation : c *)
let idc = C0

(* cons : (v -> t -> m -> v) -> t -> t *)
let rec cons k1 t2 = match t2 with
  | TNil -> Trail k1
  | Trail k2 -> Trail (fun v t m -> k1 v (cons k2 t) m)

(* apnd : t -> t -> t *)
let apnd t1 t2 = match t1 with
  | TNil -> t2
  | Trail k1 -> cons k1 t2

(* f1 : e -> string list -> v list -> c -> t -> m -> v *)
let rec f1 e xs vs c t m = match e with
  | Num n -> run_c c (VNum n) t m
  | Var x -> run_c c (List.nth vs (Env.offset x xs)) t m
  | Op (e1, op, e2) -> f1 e1 xs vs (COp1 (e2, xs, vs, op, c)) t m
  | Fun (x, e) ->
      run_c c (VFun (fun v c' t' m' ->
                       f1 e (x :: xs) (v :: vs) c' t' m')) t m
  | App (e1, e2) -> f1 e1 xs vs (CApp1 (e2, xs, vs, c)) t m
  | Try (e1, x, k, e2) ->
      let h = fun v vc c0 t0 m0 ->
                f1 e2 (x :: k :: xs) (v :: vc :: vs) c0 t0 m0
      in f1 e1 xs vs idc TNil (MCons ((c, t, h), m))
  | Perform e -> f1 e xs vs (CPerform (c)) t m
  | PerformShallow e -> f1 e xs vs (CPerformShallow (c)) t m

(* run_c : c -> v -> t -> m -> v *)
and run_c c v t m = match c with
  | C0 -> begin match t with
          | TNil -> begin match m with
                    | MNil -> v
                    | MCons ((c0, t0, h0), m0) -> run_c c0 v t0 m0
                    end
          | Trail k -> k v TNil m
          end
  | COp1 (e2, xs, vs, op, c) ->
      f1 e2 xs vs (COp2 (v, op, c)) t m
  | COp2 (v1, op, c) -> begin match (v1, v) with
      | (VNum n1, VNum n2) -> begin match op with
          | Plus -> run_c c (VNum (n1 + n2)) t m
          | Minus -> run_c c (VNum (n1 - n2)) t m
          | Times -> run_c c (VNum (n1 * n2)) t m
          | Divide -> if n2 = 0 then failwith "Division by zero"
                                else run_c c (VNum (n1 / n2)) t m
          end
      | _ -> failwith
               (to_string v1 ^ " or " ^ to_string v ^ " are not numbers")
      end
  | CApp1 (e2, xs, vs, c) ->
      f1 e2 xs vs (CApp2 (v, c)) t m
  | CApp2 (v1, c) -> apply1 v1 v c t m
  | CPerform (c) -> begin match m with
      | MNil -> failwith "empty metacontinuation in CPerform"
      | MCons ((c0, t0, h0), m0) ->
          let vc = VFun (fun v2 c2 t2 m2 ->
                           run_c c v2 t (MCons ((c2, t2, h0), m2))) in
          h0 v vc c0 t0 m0
    end
  | CPerformShallow (c) -> begin match m with
      | MNil -> failwith "empty metacontinuation in CPerformShallow"
      | MCons ((c0, t0, h0), m0) ->
          let vc = VFun (fun v2 c2 t2 m2 ->
                           let c2 = fun v t m -> run_c c2 v t m in
                           run_c c v2 (apnd t (cons c2 t2)) m2) in
          h0 v vc c0 t0 m0
    end

(* apply1 : v -> v -> c -> t -> m -> v *)
and apply1 v1 v2 c t m = match v1 with
  | VFun f -> f v2 c t m
  | _ -> failwith (to_string v1 ^ " is not a function; it can not be applied.")

(* f : e -> v *)
let f expr = f1 expr [] [] idc TNil MNil