open Callcc;;

exception Prompt;;
exception Bug;;  (* should never be raised *)

type 'a control = 'a cont * exn list;;

module E = Extensible.Make (struct end);;

(* To be compatible with the exceptio mecanism, we freeze exception as
values at prompts borderies *)

type 'a result = Value of 'a | Exception of exn;;
(* This restarts frozen  computation or reraises uncaught exceptions *)
let freeze f x =  try Value (f x) with z -> Exception  z;;
(* to freezes the current computation state catching all exceptions *)
let unfreeze = function Value x -> x | Exception z -> raise z;;

type 'a prompt = (bool * 'a result cont) E.constructor;;

let new_prompt = E.create;;

let stack = ref ([]: E.t list);;
let push pc = stack := pc :: !stack;;
let pop () =  match !stack with [] -> raise Prompt
                  | pc :: rest -> (stack := rest; pc);;

let set p e = 
  unfreeze              
    (callcc (fun normal_continuation ->
      let z = push (E.inject p (true, normal_continuation)) in
      let v = freeze e() in         
      let (effective_continuation,_) = 
           E.matches (pop())
             p (fun (b,c) -> (c,[]))
             (fun sc -> 
                match v with
                    Value _ -> raise Bug
                  | Exception z -> raise z)
      in
          (throw effective_continuation v)
      )) 
;;

let pop_control (p:'a prompt) =
  let rec pop_more control =
     E.matches (pop())
        p (fun (b, c) ->
              if b then (c, control)
              else pop_more (E.inject p (b,c):: control))
        (fun pc -> pop_more (pc :: control))
  in pop_more [];;

let rec push_control =  function
      (pc :: control) -> (push pc; push_control control)
    | [] -> ();;

let eating_cupto inside p f =
  let (abort, control) = pop_control p in
  let reified x v =
     unfreeze (callcc (fun after ->
       let z = push  (E.inject p (inside, after)) in
       let z = push_control control in
       throw x v)) in
    callcc (fun x -> throw abort (freeze f (reified x)));;

let cupto p f = eating_cupto true p f

let vomish p f = (fun k -> set p (fun () -> f k))
let cupto_minus p f = eating_cupto false p (vomish p f)
let cupto_plus p f = eating_cupto true p (vomish p f)