(*
 * Convert to continuation-passing style.
 *
 * Ling Li, Xin Yu
 * Feb. 24, 2001, Caltech
 *)

open Symbol
open Fc_ir_type
open Fc_ir_env
open Fc_ir_exn
open Fc_ir_exn_type
open Fc_ir_standardize

(************************************************************************
 * UTILITIES
 ************************************************************************)

(*
 * Default name of the continuation argument.
 * The name is significant: we want to call
 * the "exit" function as the very last thing in
 * the program.  This symbol will get shadowed
 * inside global functions.
 *)
let cont_sym = Symbol.add "exit"

(* empty type -- TyVoid *)
let tyVoid = TyInt

(* Recursively insert cont_sym to function type *)
let rec insert_cont ty =
   match ty with (* do not use tenv_expand *)
    | TyFun (args_ty, ret_ty) ->
         let ret_ty = insert_cont ret_ty in
         let cont_ty = TyFun ([ret_ty], tyVoid) in
         let args_ty = List.map insert_cont args_ty in
            TyFun (cont_ty :: args_ty, tyVoid)
    | TyArray (ty, dim) ->
         TyArray (insert_cont ty, dim)
    | TyStruct fields ->
         TyStruct (List.map (fun (v, ty) -> v, insert_cont ty) fields)
    | _ -> ty

(************************************************************************
 * EXPRESSIONS
 ************************************************************************)

(*
 * Continuation-passing transformation.
 *)
let rec cont_expr tenv e =
   let pos = string_pos "cont_expr" (atom_pos e) in
      match e with
         LetTypes (types, e) ->
            cont_types_expr tenv pos types e
       | LetFuns (funs, e) ->
            cont_funs_expr tenv pos funs e
       | LetAtom (v, ty, a, e) ->
            LetAtom (v, insert_cont ty, a, cont_expr tenv e)
       | LetCopy (v, ty, a, e) ->
            LetCopy (v, insert_cont ty, a, cont_expr tenv e)
       | LetUnop (v, ty, op, a, e) ->
            LetUnop (v, insert_cont ty, op, a, cont_expr tenv e)
       | LetBinop (v, ty, op, a1, a2, e) ->
            LetBinop (v, insert_cont ty, op, a1, a2, cont_expr tenv e)
       | LetApply (v, ty, f, args, e) ->
            cont_apply_expr tenv pos v ty f args e
       | TailCall (f, args) ->
            TailCall (f, args)
       | LetExtCall (v, f, ty, args, e) ->
            LetExtCall (v, f, ty, args, cont_expr tenv e)
       | Return a ->
            cont_return_expr tenv pos a
       | LetString (v, s, e) ->
            LetString (v, s, cont_expr tenv e)
       | IfThenElse (op, a1, a2, e1, e2) ->
            IfThenElse (op, a1, a2, cont_expr tenv e1, cont_expr tenv e2)
       | SetVar (v, ty, a, e) ->
            SetVar (v, insert_cont ty, a, cont_expr tenv e)
       | LetAddrOfVar (v1, ty, v2, e) ->
            LetAddrOfVar (v1, insert_cont ty, v2, cont_expr tenv e)
       | SetSubscript (a1, ty, a2, a3, e) ->
            SetSubscript (a1, insert_cont ty, a2, a3, cont_expr tenv e)
       | LetSubscript (v, ty, a1, a2, e) ->
            LetSubscript (v, insert_cont ty, a1, a2, cont_expr tenv e)
       | LetAddrOfSubscript (v, ty, a1, a2, e) ->
            LetAddrOfSubscript (v, insert_cont ty, a1, a2, cont_expr tenv e)
       | SetProject (a1, l, ty, a2, e) ->
            SetProject (a1, l, insert_cont ty, a2, cont_expr tenv e)
       | LetProject (v, ty, a, l, e) ->
            LetProject (v, insert_cont ty, a, l, cont_expr tenv e)
       | LetAddrOfProject (v, ty, a, l, e) ->
            LetAddrOfProject (v, insert_cont ty, a, l, cont_expr tenv e)
       | Memcpy _
       | LetAlloc _
       | LetClosure _
       | IntCall _ ->
            raise (IRException (pos, IRLevel 1))

(*
 * Add the types to the environment.
 *)
and cont_types_expr tenv pos types e =
   let pos = string_pos "cont_types_expr" pos in
   let types = List.map (fun (v, ty) -> (v, insert_cont ty)) types in
      LetTypes (types, cont_expr tenv e)

(*
 * On global functions, add a continuation argument.
 * If a function is global, and it is defined as a LetExtCall,
 * clear the global flag.
 *)
and cont_funs_expr tenv pos funs e =
   let pos = string_pos "cont_funs_expr" pos in
   let cont_fun_expr (f, gflag, ty, args, body) =
      let body = cont_expr tenv body in
         if gflag then
            (* Global fun: add continuation *)
            let gflag =
               match body with
                | LetExtCall _ -> false (* ext fun can also escape. So??? *)
                | _ -> true
            in
               f, gflag, insert_cont ty, cont_sym :: args, body
         else
            (* Local fun: change return type to TyVoid *)
            let ty =
               match ty with
                | TyFun (args_ty, _) ->
                     let args_ty = List.map insert_cont args_ty in
                        TyFun (args_ty, tyVoid)
                | _ -> raise (IRException (pos, NotAFunction ty))
            in
               f, false, ty, args, body
   in
   let funs = List.map cont_fun_expr funs in
      LetFuns (funs, cont_expr tenv e)

(*
 * Convert a function application to a tail call.
 *)
and cont_apply_expr tenv pos v ty f args e =
   let pos = string_pos "cont_apply_expr" pos in
   let g = new_symbol_string "cont_apply_expr_fun" in
   let fun_ty = TyFun ([ty], tyVoid) in
   let cont = TailCall (f, (AtomVar g) :: args) in
      LetFuns ([g, true, fun_ty, [v], cont_expr tenv e], cont)
      (* continuation is global since it is escaped *)

(*
 * A return becomes a call-current-continuation.
 *)
and cont_return_expr tenv pos a =
   TailCall (cont_sym, [a])

(*
 * The CPS result may introduce duplicate variables with
 * the name "cont",. so standardize the code.
 *)
let cont_expr tenv e =
   standardize_expr (cont_expr tenv e)