(* The source calculus. *)
module S = Lambda
(* The target calculus. *)
module T = Tail

exception NotValue of S.term
(** ^ Raised when trying to use a non-value term as such *)

let freshId =
    (** Generates a fresh variable name string *)
    let cId = ref 0 in
    (fun () ->
        incr cId ;
        (string_of_int !cId))

let freshWithPrefix pre =
    Atom.fresh (pre ^ (freshId ()) ^ "_")
let prefixHint prefix hint = match hint with
    | Some h -> prefix ^ h
    | None -> prefix
let freshBlockVarHinted hint = freshWithPrefix (prefixHint "bl_" hint)
let freshBlockVar () = freshBlockVarHinted None
let freshVarHinted hint = freshWithPrefix (prefixHint "v_" hint)
let freshVar () = freshVarHinted None

let letCont name varName body next =
    (** Allocates a block for a continuation, then runs [next] *)
    T.LetBlo(name, T.Lam(T.NoSelf, [varName], body), next)

let rec cps_value (t: S.term) : T.value = match t with
| S.Var v -> T.VVar v
| S.Lit v -> T.VLit v
| S.BinOp (l, op, r) -> T.VBinOp (cps_value l, op, cps_value r)
| S.Let _ | S.Lam _ | S.App _ | S.Print _ | S.IfZero _ -> raise (NotValue t)

let cps_value_as_term (t: S.term) (cont: T.variable): T.term =
    T.TailCall(T.vvar cont, [cps_value t])

let rec cps_term_inner (t: S.term) (cont: T.variable) (nameHint: string option)
: T.term = match t with
| S.Var _ -> cps_value_as_term t cont
| S.Lit _ -> cps_value_as_term t cont
| S.BinOp _ -> cps_value_as_term t cont
| S.Lam (self, var, term) ->
        let fName = freshBlockVarHinted nameHint
        and innerCont = freshBlockVar () in
        T.LetBlo(fName,
            T.Lam(self, [var; innerCont], cps_term_inner term innerCont None),
            T.TailCall(T.vvar cont, T.vvars [fName]))
| S.App (f, x) ->
        let xCont = freshBlockVar ()
        and fCont = freshBlockVar () in
        let xVal = freshVarHinted nameHint
        and fVal = freshVar () in

        letCont xCont xVal (
            letCont fCont fVal
                (T.TailCall(T.vvar fVal, T.vvars [xVal; cont])) @@
            (cps_term_inner f fCont None)) @@
            cps_term_inner x xCont None

| S.Print term ->
        let curCont = freshBlockVar ()
        and termVal = freshVar () in
        letCont curCont termVal (
            T.Print(T.vvar termVal,
            T.TailCall(T.vvar cont, T.vvars [termVal])))
            (cps_term_inner term curCont None)
| S.Let (var, value, next) ->
        let curCont = freshBlockVar () in
        letCont curCont var (cps_term_inner next cont None) @@
        cps_term_inner value curCont (Some (Atom.hint var))
| S.IfZero (expr, tIf, tElse) ->
        let curCont = freshBlockVar ()
        and exprVal = freshVar () in
        letCont curCont exprVal (T.IfZero(T.vvar exprVal,
                cps_term_inner tIf cont None,
                cps_term_inner tElse cont None)) @@
            cps_term_inner expr curCont None

let cps_term (t: S.term): T.term =
    (** Entry point. Transforms a [Lambda] term into a [Tail] term, applying a
     * continuation-passing-style transformation. *)
    let exitBlock = freshBlockVarHinted (Some "exit") in
    letCont exitBlock (freshVar ()) T.Exit @@
        cps_term_inner t exitBlock (Some "main_entry")