open Instr

type t = continuation

let halt = Halt

let return_to cont = Return_to cont

let apply_to func cont = Apply_to { func; cont }

let eval_to expr cont = Eval_to { expr; cont }

let pop_to n cont = Pop_to { n; cont }

let exception_to handler cont = Exception_to { handler; cont }

let rec build_cont instructions cont =
  match instructions with
  | [] -> cont
  | instr :: rest ->
    eval_to instr (build_cont rest cont)

let apply_cont k v : (t * value option) option =
  match k with
  | Halt -> None
  | Return_to k' -> Some (k', Some v)
  | Apply_to { func; cont } ->
    (match func with
     | Closure { code; env } ->
       Some (Eval_to { expr = Instr.LOAD code; cont }, Some v)
     | _ -> None)
  | Eval_to { expr; cont } ->
    Some (cont, Some v)
  | Pop_to { n; cont } ->
    Some (cont, Some v)
  | Exception_to { handler; cont } ->
    Some (handler, Some v)

let rec depth = function
  | Halt -> 0
  | Return_to k -> 1 + depth k
  | Apply_to { cont; _ } -> 1 + depth cont
  | Eval_to { cont; _ } -> 1 + depth cont
  | Pop_to { cont; _ } -> 1 + depth cont
  | Exception_to { cont; _ } -> 1 + depth cont

let rec show_cont k =
  match k with
  | Halt -> "Halt"
  | Return_to k' -> Printf.sprintf "Return_to(%s)" (show_cont k')
  | Apply_to { func; cont } -> Printf.sprintf "Apply_to(%s, %s)" (show_value func) (show_cont cont)
  | Eval_to { expr; cont } -> Printf.sprintf "Eval_to(%s, %s)" (show expr) (show_cont cont)
  | Pop_to { n; cont } -> Printf.sprintf "Pop_to(%d, %s)" n (show_cont cont)
  | Exception_to { handler; cont } -> Printf.sprintf "Exception_to(%s, %s)" (show_cont handler) (show_cont cont)