(*
 * Escape analysis
 *
 * Ling Li, Xin Yu
 * Feb. 24, 2001, Caltech
 *
 * This step is called escape analysis.  Each global function
 * allocates a record, called a "frame."  Variables that escape
 * are stored in the frame.  A variable escapes if:
 *    1. it is an aggregate (a struct), or
 *    2. it is used in a nested global function, or
 *    3. it's address is taken.
 *
 * We do this in two passes.  In the first pass, we walk the
 * program, and compute which variables escape.  We assume all
 * variable binding occurrences bind different variables (the
 * program has been standardized).  This means we can use a table
 * to record tyhe frame for each escaping variable.  A frame stack
 * is passed during the analysis for the nesting level of global
 * functions.
 *
 * In the second pass, we transform the code.
 *    1. Add an explicit frame allocation to each global function
 *       definition using LetAlloc.
 *    2. All accesses to escaping variables are transformed
 *       into record operations on the frame.  For struct
 *       assignments and initializations, we use Memcpy to
 *       copy the date.  For scalar assignments, we use
 *       SetProject.
 *
 * Global variables are also stored in an outermost frame.
 * There is no need to distinguish between top-level expressions
 * and expressions inside a function definition.  Another way to think
 * about it: it is as if the entire program were nested inside an implicit
 * function definition.
 *
 * All global flags for nested functions are cleared.
 *)

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

(************************************************************************
 * TYPES
 ************************************************************************)

(*
 * Frames are just vars.
 *)
type frame = var

(*
 * A variable has:
 *    1. a type
 *    2. an optional frame where the variable is bound
 *)
type venv = (ty * frame) SymbolTable.t

(************************************************************************
 * ENVIRONMENTS
 ************************************************************************)

(*
 * Default exception.
 *)
let raise_exn code v =
   IRException (code, v)

let tenv_lookup      = tenv_lookup raise_exn
let tenv_expand      = tenv_expand raise_exn

(*
 * Empty table.
 *)
let empty_venv = SymbolTable.create ()

(*
 * Get the value.
 *)
let venv_lookup venv pos v =
   try SymbolTable.find venv v with
      Not_found ->
         raise (IRException (pos, UnboundVar v))

(*
 * Add a value.
 *)
let venv_add venv v ty frame =
   SymbolTable.add venv v (ty, frame)

(*
 * We adjust the calling convention here.
 * We add an extra argument to functions that
 * return an aggregate type.
 *)
let return_sym = Symbol.new_symbol_string "return"

(*
 * The data segment is also a frame.
 *)
let data_sym = new_symbol_string "data_seg"

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

(*
 * Expand a function type.
 *)
let dest_fun_type tenv pos ty =
   match tenv_expand tenv pos ty with
      TyFun (ty_vars, ty_res) -> ty_vars, ty_res
    | ty -> raise (IRException (pos, NotAFunction ty))

(*
 * Get the type of an atom.
 *)
let get_atom_type venv pos = function
   AtomChar _ -> TyChar
 | AtomInt _ -> TyInt
 | AtomFloat _ -> TyFloat
 | AtomNil ty -> TyArray (ty, None)
 | AtomVar v -> fst (venv_lookup venv pos v)

(*
 * Default value in a non-aggregate type.
 *)
let default_value tenv pos ty =
   let ty = tenv_expand tenv pos ty in
      match ty with
         TyChar -> AtomChar '\000'
       | TyInt -> AtomInt 0
       | TyFloat -> AtomFloat 0.0
       | TyArray (ty, None) -> AtomNil ty

       | TyAny
       | TyArray (_, Some _)
       | TyStruct _
       | TyFun _ ->
            raise (IRException (pos, CantInitialize ty))

       | TyId _ ->
            raise (Invalid_argument "Fc_ir_closure1.default_value")

(*
 * Check for an aggregate type.
 *)
let is_aggregate_type tenv pos ty =
   match tenv_expand tenv pos ty with
      TyArray (_, Some _)
    | TyStruct _ ->
         true
    | TyAny
    | TyChar
    | TyInt
    | TyFloat
    | TyArray (_, None)
    | TyFun _ ->
         false
    | TyId _ ->
         raise (Invalid_argument "Fc_ir_closure.is_aggregate_type")

(************************************************************************
 * FIRST PASS: COLLECT ESCAPING VARS
 ************************************************************************)

(*
 * Escape environment.
 * We keep three parts:
 *    esc_stack: a list of frames in scope
 *    esc_vars: a table of variables that escape, and the frame they belong to
 *    esc_funs: a table of frames and gflags for each fun
 *)
type esc =
   { esc_stack : frame list;
     esc_vars : (ty * frame) SymbolTable.t;
     esc_funs : (frame * bool) SymbolTable.t
   }

(*
 * Empty escape environment include the
 * outermost frame.
 *)
let empty_esc =
   { esc_stack = [data_sym];
     esc_vars = SymbolTable.create ();
     esc_funs = SymbolTable.create ()
   }

(*
 * Add a new frame to the escape table.
 *)
let push_frame esc f frame =
   let { esc_stack = stack; esc_funs = funs } = esc in
   let gflag =
      match stack with
         [_] -> true
       | _ -> false
   in
      { esc with esc_stack = frame :: stack;
                 esc_funs = SymbolTable.add funs f (frame, gflag)
      }

(*
 * Get the current frame, and the rest of the stack.
 *)
let pop_frame pos = function
   { esc_stack = frame :: funs } as esc ->
      frame, { esc with esc_stack = funs }
 | { esc_stack = [] } ->
      raise (IRException (pos, InternalError "pop_frame"))

(*
 * Add an escaping variable to the frame.
 *)
let esc_add esc v ty frame =
   { esc with esc_vars = SymbolTable.add esc.esc_vars v (ty, frame) }

(*
 * Add an atom.
 * If the atom is a var, and it is bound in
 * another frame, add it to the escape set.
 *)
let esc_var tenv venv esc pos v =
   let frame', _ = pop_frame pos esc in
   let ty, frame = venv_lookup venv pos v in
      if frame' <> frame || is_aggregate_type tenv pos ty then
         esc_add esc v ty frame
      else
         esc

let esc_atom tenv venv esc pos a =
   let pos = string_pos "esc_atom" pos in
      match a with
         AtomChar _
       | AtomInt _
       | AtomFloat _
       | AtomNil _ ->
            esc
       | AtomVar v ->
            esc_var tenv venv esc pos v

let esc_atoms tenv venv esc pos atoms =
   List.fold_left (fun esc a ->
         esc_atom tenv venv esc pos a) esc atoms

let esc_atom_opt tenv venv esc pos a =
   match a with
      Some a ->
         esc_atom tenv venv esc pos a
    | None ->
         esc

(*
 * Transform an expression.
 *)
let rec esc_expr tenv venv esc e =
   let pos = string_pos "esc_expr" (atom_pos e) in
      match e with
         LetTypes (types, e) ->
            esc_types_expr tenv venv esc pos types e
       | LetFuns (funs, e) ->
            esc_funs_expr tenv venv esc pos funs e
       | LetAtom (v, ty, a, e) ->
            esc_atom_expr tenv venv esc pos v ty a e
       | LetCopy (v, ty, a, e) ->
            esc_copy_expr tenv venv esc pos v ty a e
       | LetUnop (v, ty, op, a, e) ->
            esc_unop_expr tenv venv esc pos v ty op a e
       | LetBinop (v, ty, op, a1, a2, e) ->
            esc_binop_expr tenv venv esc pos v ty op a1 a2 e
       | TailCall (f, args) ->
            esc_tailcall_expr tenv venv esc pos f args
       | LetExtCall (v, f, ty, args, e) ->
            esc_extcall_expr tenv venv esc pos v f ty args e
       | LetString (v, s, e) ->
            esc_string_expr tenv venv esc pos v s e
       | IfThenElse (op, a1, a2, e1, e2) ->
            esc_if_expr tenv venv esc pos op a1 a2 e1 e2
       | SetVar (v, ty, a, e) ->
            esc_set_var_expr tenv venv esc pos v ty a e
       | LetAddrOfVar (v1, ty, v2, e) ->
            esc_addr_of_var_expr tenv venv esc pos v1 ty v2 e
       | SetSubscript (a1, ty, a2, a3, e) ->
            esc_set_subscript_expr tenv venv esc pos ty a1 a2 a3 e
       | LetSubscript (v, ty, a1, a2, e) ->
            esc_subscript_expr tenv venv esc pos v ty a1 a2 e
       | LetAddrOfSubscript (v, ty, a1, a2, e) ->
            esc_addr_of_subscript_expr tenv venv esc pos v ty a1 a2 e
       | SetProject (a1, l, ty, a2, e) ->
            esc_set_project_expr tenv venv esc pos ty a1 l a2 e
       | LetProject (v, ty, a, l, e) ->
            esc_project_expr tenv venv esc pos v ty a l e
       | LetAddrOfProject (v, ty, a, l, e) ->
            esc_addr_of_project_expr tenv venv esc pos v ty a l e
       | LetApply _
       | Return _
       | Memcpy _
       | LetAlloc _
       | LetClosure _
       | IntCall _ ->
            raise (IRException (pos, IRLevel 1))

(*
 * Add the type definitions.
 *)
and esc_types_expr tenv venv esc pos types e =
   let pos = string_pos "esc_types_expr" pos in
   let tenv =
      List.fold_left (fun tenv (v, ty) -> tenv_add tenv v ty) tenv types
   in
      esc_expr tenv venv esc e

(*
 * Function definitions.
 * Global functions should push a new frame
 * onto the frame stack.  Also, all aggregate arguments
 * escape.
 *)
and esc_fun_global tenv venv esc pos f ty vars body =
   (* Add a new frame for this function *)
   let pos = string_pos "esc_fun_global" pos in
   let ty_vars, _ = dest_fun_type tenv (int_pos 1 pos) ty in
   let frame = new_symbol_string "frame" in
   let esc = push_frame esc f frame in

   (* Add aggregate args to the frame *)
   let esc, venv =
      List.fold_left2 (fun (esc, venv) v ty ->
            let esc =
               if is_aggregate_type tenv (int_pos 2 pos) ty then
                  esc_add esc v ty frame
               else
                  esc
            in
               esc, venv_add venv v ty frame) (esc, venv) vars ty_vars
   in

   (* Convert the body *)
   let esc = esc_expr tenv venv esc body in
   let _, esc = pop_frame (int_pos 3 pos) esc in
      esc

and esc_fun_local tenv venv esc pos f ty vars body =
   (* Get the frame *)
   let pos = string_pos "esc_fun_local" pos in
   let ty_vars, _ = dest_fun_type tenv pos ty in
   let frame, _ = pop_frame pos esc in

   (* Add aggregate args to the frame *)
   let esc, venv =
      List.fold_left2 (fun (esc, venv) v ty ->
            let esc =
               if is_aggregate_type tenv pos ty then
                  esc_add esc v ty frame
               else
                  esc
            in
               esc, venv_add venv v ty frame) (esc, venv) vars ty_vars
   in
      (* Convert the body *)
      esc_expr tenv venv esc body

and esc_funs_expr tenv venv esc pos funs e =
   (* Add all the functions to the variable environment *)
   let pos = string_pos "esc_funs_expr" pos in
   let frame, _ = pop_frame (int_pos 0 pos) esc in
   let venv =
      List.fold_left (fun venv (f, _, ty, _, _) ->
            venv_add venv f ty frame) venv funs
   in

   (* Convert all the funs *)
   let esc =
      List.fold_left (fun esc (f, gflag, ty, vars, body) ->
            let esc_fun =
               if gflag then
                  esc_fun_global
               else
                  esc_fun_local
            in
               esc_fun tenv venv esc (int_pos 2 pos) f ty vars body) esc funs
   in
      (* Convert the body *)
      esc_expr tenv venv esc e

(*
 * Atom assignment.  The atom assignment is
 * _not_ a copy.
 *)
and esc_atom_expr tenv venv esc pos v ty a e =
   (* Compute the rest using the current type *)
   let pos = string_pos "esc_atom_expr" pos in
   let frame, _ = pop_frame pos esc in
   let venv = venv_add venv v ty frame in
   let esc = esc_atom tenv venv esc pos a in
      esc_expr tenv venv esc e

(*
 * Copy expression is an allocation point.
 * If the value is an aggregate, it escapes.
 *)
and esc_copy_expr tenv venv esc pos v ty a e =
   (* Compute the rest using the current type *)
   let pos = string_pos "esc_copy_expr" pos in
   let frame, _ = pop_frame pos esc in
   let venv = venv_add venv v ty frame in
   let esc = esc_atom_opt tenv venv esc pos a in
   let esc = esc_expr tenv venv esc e in
      if is_aggregate_type tenv pos ty then
         esc_add esc v ty frame
      else
         esc

(*
 * Unary operation.
 * Arguments are always of scalar type, so
 * we don't worry about aggregates.
 *)
and esc_unop_expr tenv venv esc pos v ty op a e =
   let pos = string_pos "esc_unop_expr" pos in
   let frame, _ = pop_frame pos esc in
   let venv = venv_add venv v ty frame in
   let esc = esc_atom tenv venv esc pos a in
      esc_expr tenv venv esc e

(*
 * Binary operation.
 * Arguments are always of scalar type.
 *)
and esc_binop_expr tenv venv esc pos v ty op a1 a2 e =
   let pos = string_pos "esc_binop_expr" pos in
   let frame, _ = pop_frame pos esc in
   let venv = venv_add venv v ty frame in
   let esc = esc_atom tenv venv esc pos a1 in
   let esc = esc_atom tenv venv esc pos a2 in
      esc_expr tenv venv esc e

(*
 * Tailcall.
 * Add all the vars that escape.
 * The functions do not escape.
 *)
and esc_tailcall_expr tenv venv esc pos f args =
   let pos = string_pos "esc_tailcall_expr" pos in
      esc_atoms tenv venv esc pos args

(*
 * External function call.
 *)
and esc_extcall_expr tenv venv esc pos v f ty args e =
   let pos = string_pos "esc_extcall_expr" pos in
   let frame, _ = pop_frame pos esc in
   let _, ty_res = dest_fun_type tenv pos ty in
   let venv = venv_add venv v ty_res frame in
   let esc = esc_atoms tenv venv esc pos args in
   let esc = esc_expr tenv venv esc e in
      if is_aggregate_type tenv pos ty_res then
         esc_add esc v ty_res frame
      else
         esc

(*
 * String allocation.
 *)
and esc_string_expr tenv venv esc pos v s e =
   let pos = string_pos "esc_string_expr" pos in
   let frame, _ = pop_frame pos esc in
   let ty = TyArray (TyChar, None) in
   let venv = venv_add venv v ty frame in
      esc_expr tenv venv esc e

(*
 * Conditional.
 *)
and esc_if_expr tenv venv esc pos op a1 a2 e1 e2 =
   let pos = string_pos "esc_if_expr" pos in
   let esc = esc_atom tenv venv esc pos a1 in
   let esc = esc_atom tenv venv esc pos a2 in
   let esc = esc_expr tenv venv esc e1 in
   let esc = esc_expr tenv venv esc e2 in
      esc

(*
 * Set a variable.
 *)
and esc_set_var_expr tenv venv esc pos v ty a e =
   let pos = string_pos "esc_set_var_expr" pos in
   let esc = esc_var tenv venv esc pos v in
   let esc = esc_atom tenv venv esc pos a in
      esc_expr tenv venv esc e

(*
 * Address of a variable is eliminated.
 * If the variable is an aggregate, the AddressOf operation
 * is removed.  Otherwise, the variable is added to the escape
 * set.
 *)
and esc_addr_of_var_expr tenv venv esc pos v1 ty v2 e =
   (* Add the variable to the escape set *)
   let pos = string_pos "esc_addr_of_var_expr" pos in
   let frame, _ = pop_frame pos esc in
   let ty', frame' = venv_lookup venv pos v2 in
   let esc = esc_add esc v2 ty' frame' in
   let venv = venv_add venv v1 ty frame in
      esc_expr tenv venv esc e

(*
 * Set an array entry.
 *)
and esc_subscript_expr tenv venv esc pos v ty a1 a2 e =
   let pos = string_pos "esc_subscript_expr" pos in
   let esc = esc_atom tenv venv esc pos a1 in
   let esc = esc_atom tenv venv esc pos a2 in
   let frame, _ = pop_frame pos esc in
   let venv = venv_add venv v ty frame in
      esc_expr tenv venv esc e

and esc_set_subscript_expr tenv venv esc pos ty a1 a2 a3 e =
   let pos = string_pos "esc_set_subscript_expr" pos in
   let esc = esc_atom tenv venv esc pos a1 in
   let esc = esc_atom tenv venv esc pos a2 in
   let esc = esc_atom tenv venv esc pos a3 in
      esc_expr tenv venv esc e

and esc_addr_of_subscript_expr tenv venv esc pos v ty a1 a2 e =
   let pos = string_pos "esc_addr_of_subscript_expr" pos in
   let esc = esc_atom tenv venv esc pos a1 in
   let esc = esc_atom tenv venv esc pos a2 in
   let frame, _ = pop_frame pos esc in
   let venv = venv_add venv v ty frame in
      esc_expr tenv venv esc e

(*
 * Struct projections.
 *)
and esc_project_expr tenv venv esc pos v ty a l e =
   let pos = string_pos "esc_project_expr" pos in
   let esc = esc_atom tenv venv esc pos a in
   let frame, _ = pop_frame pos esc in
   let venv = venv_add venv v ty frame in
      esc_expr tenv venv esc e

and esc_set_project_expr tenv venv esc pos ty a1 l a2 e =
   let pos = string_pos "esc_set_project_expr" pos in
   let esc = esc_atom tenv venv esc pos a1 in
   let esc = esc_atom tenv venv esc pos a2 in
      esc_expr tenv venv esc e

and esc_addr_of_project_expr tenv venv esc pos v ty a l e =
   let pos = string_pos "esc_addr_of_project_expr" pos in
   let esc = esc_atom tenv venv esc pos a in
   let frame, _ = pop_frame pos esc in
   let venv = venv_add venv v ty frame in
      esc_expr tenv venv esc e

(************************************************************************
 * SECOND PASS: ALLOCATE FRAMES
 ************************************************************************)

(*
 * Conversion requires a set of variables that
 * escape, as well as the types of the frames.
 *)
type aelim =
   { aelim_esc : (ty * frame) SymbolTable.t;
     aelim_frames : ty SymbolTable.t;
     aelim_funs : (frame * bool) SymbolTable.t
   }

(*
 * Convert the escaping vars.
 *)
let aelim_of_esc { esc_vars = vars; esc_funs = esc_funs } =
   (* Collect all the escaping vars and frame fields *)
   let esc, frames =
      SymbolTable.fold (fun (esc, frames) v entry ->
            let ty, frame = List.hd entry in
            let esc = SymbolTable.add esc v (ty, frame) in
            let frames =
               try
                  let fields = SymbolTable.find frames frame in
                  let fields = (v, ty) :: fields in
                     SymbolTable.add frames frame fields
               with
                  Not_found ->
                     SymbolTable.add frames frame [v, ty]
            in
               esc, frames) (SymbolTable.create (), SymbolTable.create ()) vars
   in

   (* Build the frame types *)
   let frames =
      SymbolTable.fold (fun frames frame fields ->
            let ty = TyStruct (List.hd fields) in
               SymbolTable.add frames frame ty) (SymbolTable.create ()) frames
   in
      { aelim_esc = esc; aelim_frames = frames; aelim_funs = esc_funs }

(*
 * Elim types.
 * Functions never take aggregate arguments.
 *)
let rec aelim_type tenv pos ty =
   let pos = string_pos "aelim_type" pos in
   match ty with
    | TyAny | TyChar | TyInt | TyFloat
    | TyId _ -> ty
    | TyArray (ty, dim) ->
         TyArray (aelim_type tenv pos ty, dim)
    | TyStruct fields ->
         let fields =
            List.map (fun (v, ty) -> v, aelim_type tenv pos ty) fields
         in
            TyStruct fields
    | TyFun (ty_vars, ty_ret) ->
         let to_pointer ty =
            let ty = aelim_type tenv pos ty in
               if is_aggregate_type tenv pos ty then
                  TyArray (ty, None)
               else ty
         in
         let ty_vars = List.map to_pointer ty_vars in
            TyFun (ty_vars, aelim_type tenv pos ty_ret)

(*
 * Type of the aggregate.
 *)
let type_of_aggregate tenv pos ty =
   let ty = aelim_type tenv pos (tenv_expand tenv pos ty) in
      match ty with
         TyArray (ty, Some _) -> TyArray (ty, None)
       | _ -> TyArray (ty, None)

(*
 * Lookup a variable.
 * Expand the type if the variable escapes.
 *)
type aelim_type =
   AggrType of ty * frame
 | ScalarType of ty * frame
 | NoType

let venv_lookup_var tenv esc pos v =
   try
      let ty, frame = SymbolTable.find esc.aelim_esc v in
      let ty = aelim_type tenv pos ty in
         if is_aggregate_type tenv pos ty then
            AggrType (type_of_aggregate tenv pos ty, frame)
         else
            ScalarType (ty, frame)
   with
      Not_found ->
         NoType

(*
 * Look up the frame for a function.
 *)
let venv_lookup_fun esc pos f =
   try SymbolTable.find esc.aelim_funs f with
      Not_found ->
         raise (IRException (pos, UnboundVar f))

(*
 * Look up the type for a frame.
 *)
let venv_lookup_frame tenv esc pos frame =
   try aelim_type tenv pos (SymbolTable.find esc.aelim_frames frame) with
      Not_found ->
         TyStruct []

(*
 * Atom conversion.
 *)
let aelim_atom tenv esc pos a cont =
   let pos = string_pos "aelim_atom" pos in
      match a with
       | AtomChar _ | AtomInt _ | AtomFloat _ | AtomNil _ ->
            cont a
       | AtomVar v ->
            match venv_lookup_var tenv esc pos v with
             | AggrType (ty, frame) ->
                  LetAddrOfProject (v, ty, AtomVar frame, v, cont a)
             | ScalarType (ty, frame) ->
                  LetProject (v, ty, AtomVar frame, v, cont a)
             | NoType ->
                  cont a

let aelim_atoms tenv esc pos atoms cont =
   let rec loop args = function
      a :: atoms ->
         aelim_atom tenv esc pos a (fun a ->
               loop (a :: args) atoms)
    | [] ->
         cont (List.rev args)
   in
      loop [] atoms

let aelim_atom_opt tenv esc pos a cont =
   match a with
      Some a ->
         aelim_atom tenv esc pos a (fun a -> cont (Some a))
    | None ->
         cont None

(*
 * Transform an expression.
 *)
let rec aelim_expr tenv esc e =
   let pos = atom_pos e in
      match e with
         LetTypes (types, e) ->
            aelim_types_expr tenv esc pos types e
       | LetFuns (funs, e) ->
            aelim_funs_expr tenv esc pos funs e
       | LetAtom (v, ty, a, e) ->
            aelim_atom_expr tenv esc pos v ty a e
       | LetCopy (v, ty, a, e) ->
            aelim_copy_expr tenv esc pos v ty a e
       | LetUnop (v, ty, op, a, e) ->
            aelim_unop_expr tenv esc pos v ty op a e
       | LetBinop (v, ty, op, a1, a2, e) ->
            aelim_binop_expr tenv esc pos v ty op a1 a2 e
       | TailCall (f, args) ->
            aelim_tailcall_expr tenv esc pos f args
       | LetExtCall (v, f, ty, args, e) ->
            aelim_extcall_expr tenv esc pos v f ty args e
       | LetString (v, s, e) ->
            aelim_string_expr tenv esc pos v s e
       | IfThenElse (op, a1, a2, e1, e2) ->
            aelim_if_expr tenv esc pos op a1 a2 e1 e2
       | SetVar (v, ty, a, e) ->
            aelim_set_var_expr tenv esc pos v ty a e
       | LetAddrOfVar (v1, ty, v2, e) ->
            aelim_addr_of_var_expr tenv esc pos v1 ty v2 e
       | SetSubscript (a1, ty, a2, a3, e) ->
            aelim_set_subscript_expr tenv esc pos ty a1 a2 a3 e
       | LetSubscript (v, ty, a1, a2, e) ->
            aelim_subscript_expr tenv esc pos v ty a1 a2 e
       | LetAddrOfSubscript (v, ty, a1, a2, e) ->
            aelim_addr_of_subscript_expr tenv esc pos v ty a1 a2 e
       | SetProject (a1, l, ty, a2, e) ->
            aelim_set_project_expr tenv esc pos ty a1 l a2 e
       | LetProject (v, ty, a, l, e) ->
            aelim_project_expr tenv esc pos v ty a l e
       | LetAddrOfProject (v, ty, a, l, e) ->
            aelim_addr_of_project_expr tenv esc pos v ty a l e
       | LetApply _
       | Return _
       | Memcpy _
       | LetAlloc _
       | LetClosure _
       | IntCall _ ->
            raise (IRException (pos, IRLevel 1))

(*
 * Add the type definitions.
 *)
and aelim_types_expr tenv esc pos types e =
   let pos = string_pos "aelim_types_expr" pos in
   let types = List.map (fun (v, ty) -> v, aelim_type tenv pos ty) types in
   let tenv =
      List.fold_left (fun tenv (v, ty) -> tenv_add tenv v ty) tenv types
   in
   let e = aelim_expr tenv esc e in
      LetTypes (types, e)

(*
 * Function definitions.
 * For global functions, allocate the frame,
 * copy the aggregate arguments to the frame, and
 * convert the body.
 *)
and aelim_fun_global tenv esc pos f ty vars body =
   let pos = string_pos "aelim_fun_global" pos in

   (* Convert the function body *)
   let body = aelim_expr tenv esc body in

   (* Copy all the aggregate arguments *)
   let ty_vars, ty_res = dest_fun_type tenv pos ty in
   let body =
      List.fold_left2 (fun body v ty ->
            match venv_lookup_var tenv esc pos v with
             | AggrType (ty', frame) ->
                  let v' = new_symbol v in
                     LetAddrOfProject (v', ty', AtomVar frame, v,
                                       Memcpy (AtomVar v', ty, AtomVar v, body))
                                               (*LetAtom (v, ty', AtomVar v', body)))???*)
             | ScalarType (ty, frame) ->
                  SetProject (AtomVar frame, v, ty, AtomVar v, body)
             | NoType ->
                  body) body vars ty_vars
   in

   (* Allocate the frame *)
   let frame, gflag = venv_lookup_fun esc pos f in
   let ty_frame = venv_lookup_frame tenv esc pos frame in
   let body =
      LetTypes ([frame, ty_frame], LetAlloc (frame, TyId frame, body))
   in
   let ty = aelim_type tenv pos ty in
      f, gflag, ty, vars, body

and aelim_fun_local tenv esc pos f ty vars body =
   (* Get the frame *)
   let pos = string_pos "aelim_fun_local" pos in

   (* Convert the body *)
   let body = aelim_expr tenv esc body in

   (* Copy all the aggregate arguments *)
   let ty_vars, ty_res = dest_fun_type tenv pos ty in
   let body =
      List.fold_left2 (fun body v ty ->
            match venv_lookup_var tenv esc pos v with
               AggrType (ty', frame) ->
                  let v' = new_symbol v in
                  let ty' = type_of_aggregate tenv pos ty in
                     LetAddrOfProject (v', ty', AtomVar frame, v,
                                       Memcpy (AtomVar v', ty, AtomVar v,
                                               LetAtom (v, ty', AtomVar v', body)))
             | ScalarType (ty, frame) ->
                  SetProject (AtomVar frame, v, ty, AtomVar v, body)

             | NoType ->
                  body) body vars ty_vars
   in

   (* Convert the function type *)
   let ty = aelim_type tenv pos ty in
      f, false, ty, vars, body

and aelim_funs_expr tenv esc pos funs e =
   (* Add add the functions to the variable environment *)
   let pos = string_pos "aelim_funs_expr" pos in

   (* Convert all the funs *)
   let funs =
      List.map (fun (f, gflag, ty, vars, body) ->
            let aelim_fun =
               if gflag then
                  aelim_fun_global
               else
                  aelim_fun_local
            in
               aelim_fun tenv esc (int_pos 2 pos) f ty vars body) funs
   in

   (* Convert the body *)
   let e = aelim_expr tenv esc e in
      LetFuns (funs, e)

(*
 * Atom assignment.  If the value is an aggregate,
 * it becomes a pointer.
 *)
and aelim_atom_expr tenv esc pos v ty a e =
   let pos = string_pos "aelim_atom_expr" pos in
   let ty = aelim_type tenv pos ty in
   let e = aelim_expr tenv esc e in
      aelim_atom tenv esc pos a (fun a ->
            if is_aggregate_type tenv pos ty then
               LetAtom (v, type_of_aggregate tenv pos ty, a, e)
            else
               match venv_lookup_var tenv esc pos v with
                  AggrType (ty', frame) ->
                     let v' = new_symbol v in
                        LetAddrOfProject (v', ty', AtomVar frame, v,
                                          Memcpy (AtomVar v', ty, a, e))
                | ScalarType (ty, frame) ->
                     SetProject (AtomVar frame, v, ty, a, e)
                | NoType ->
                     LetAtom (v, ty, a, e))

(*
 * Copy expression is an allocation point.
 * If the value is an aggregate, copy it with Memcpy.
 *)
and aelim_copy_expr tenv esc pos v ty a e =
   (* Compute the rest using the current type *)
   let pos = string_pos "aelim_copy_expr" pos in
   let ty = aelim_type tenv pos ty in
   let e = aelim_expr tenv esc e in
      aelim_atom_opt tenv esc pos a (fun a ->
         try
            let a =
               match a with
                | None -> default_value tenv pos ty
                | Some a -> a
            in
               match venv_lookup_var tenv esc pos v with
                | AggrType (ty', frame) ->
                     let v' = new_symbol v in
                        LetAddrOfProject (v', ty', AtomVar frame, v,
                                          Memcpy (AtomVar v', ty, a, e))
                | ScalarType (ty, frame) -> (*??? If a is None, need it be init? *)
                     SetProject (AtomVar frame, v, ty, a, e)
                | NoType ->
                     LetAtom (v, ty, a, e)
         with
            IRException (_, CantInitialize _) -> e
      )

(*
 * Unary operation.
 * Arguments are always of scalar type, so
 * we don't worry about aggregates.
 *)
and aelim_unop_expr tenv esc pos v ty op a e =
   let pos = string_pos "aelim_unop_expr" pos in
   let ty = aelim_type tenv pos ty in
   let e = aelim_expr tenv esc e in
      aelim_atom tenv esc pos a (fun a ->
            match venv_lookup_var tenv esc pos v with
               AggrType _ ->
                  raise (IRException (pos, InternalError "aelim_unop_expr"))
             | ScalarType (ty, frame) ->
                  let v' = new_symbol v in
                     LetUnop (v', ty, op, a,
                              SetProject (AtomVar frame, v, ty, AtomVar v', e))
             | NoType ->
                  LetUnop (v, ty, op, a, e))

(*
 * Binary operation.
 * Arguments are always of scalar type.
 *)
and aelim_binop_expr tenv esc pos v ty op a1 a2 e =
   let pos = string_pos "aelim_binop_expr" pos in
   let ty = aelim_type tenv pos ty in
   let e = aelim_expr tenv esc e in
      aelim_atom tenv esc pos a1 (fun a1 ->
            aelim_atom tenv esc pos a2 (fun a2 ->
                  match venv_lookup_var tenv esc pos v with
                     AggrType _ ->
                        raise (IRException (pos, InternalError "aelim_binop_expr"))
                   | ScalarType (ty, frame) ->
                        let v' = new_symbol v in
                           LetBinop (v', ty, op, a1, a2,
                                     SetProject (AtomVar frame, v, ty, AtomVar v', e))
                   | NoType ->
                        LetBinop (v, ty, op, a1, a2, e)))

(*
 * Tailcall.
 * Check the function type.
 * If it returns an aggregate, we need to
 * pass the return argument.
 *)
and aelim_tailcall_expr tenv esc pos f args =
   let pos = string_pos "aelim_tailcall_expr" pos in
      aelim_atoms tenv esc pos args (fun args ->
            TailCall (f, args))

(*
 * External function call.
 *)
and aelim_extcall_expr tenv esc pos v f ty args e =
   let pos = string_pos "aelim_extcall_expr" pos in
   let _, ty_res = dest_fun_type tenv pos ty in
   let ty_res = aelim_type tenv pos ty_res in
   let ty = aelim_type tenv pos ty in
   let e = aelim_expr tenv esc e in
      aelim_atoms tenv esc pos args (fun args ->
            match venv_lookup_var tenv esc pos v with
               AggrType (ty', frame) ->
                  let v' = new_symbol v in
                  let v'' = new_symbol v in
                  let ty_vars, _ = dest_fun_type tenv pos ty in
                  let ty = TyFun (ty_vars, ty') in
                     LetExtCall (v', f, ty, args,
                                 LetAddrOfProject (v'', ty', AtomVar frame, v,
                                                   Memcpy (AtomVar v'', ty_res, AtomVar v', e)))
             | ScalarType (ty', frame) ->
                  let v' = new_symbol v in
                     LetExtCall (v', f, ty, args,
                                 SetProject (AtomVar frame, v, ty', AtomVar v', e))
             | NoType ->
                  LetExtCall (v, f, ty, args, e))

(*
 * String allocation.
 *)
and aelim_string_expr tenv esc pos v s e =
   let pos = string_pos "aelim_string_expr" pos in
   let ty = TyArray (TyChar, None) in
   let e = aelim_expr tenv esc e in
      match venv_lookup_var tenv esc pos v with
         AggrType _ ->
            raise (IRException (pos, InternalError "aelim_string_expr"))
       | ScalarType (ty, frame) ->
            let v' = new_symbol v in
               LetString (v', s, SetProject (AtomVar frame, v, ty, AtomVar v', e))
       | NoType ->
            LetString (v, s, e)

(*
 * Conditional.
 *)
and aelim_if_expr tenv esc pos op a1 a2 e1 e2 =
   let pos = string_pos "aelim_if_expr" pos in
   let e1 = aelim_expr tenv esc e1 in
   let e2 = aelim_expr tenv esc e2 in
      aelim_atom tenv esc pos a1 (fun a1 ->
            aelim_atom tenv esc pos a2 (fun a2 ->
                  IfThenElse (op, a1, a2, e1, e2)))

(*
 * Set a variable.
 *)
and aelim_set_var_expr tenv esc pos v ty a e =
   let pos = string_pos "aelim_set_var_expr" pos in
   let e = aelim_expr tenv esc e in
   let ty = aelim_type tenv pos ty in
      aelim_atom tenv esc pos a (fun a ->
            match venv_lookup_var tenv esc pos v with
               AggrType (ty', frame) ->
                  let v' = new_symbol v in
                     LetAddrOfProject (v', ty', AtomVar frame, v,
                                       Memcpy (AtomVar v', ty, a, e))
             | ScalarType (ty, frame) ->
                  SetProject (AtomVar frame, v, ty, a, e)
             | NoType ->
                  SetVar (v, ty, a, e))

(*
 * Address of a variable is eliminated.
 * The variable better have escaped.
 *)
and aelim_addr_of_var_expr tenv esc pos v1 ty v2 e =
   (* Add the variable to the escape set *)
   let pos = string_pos "aelim_addr_of_var_expr" pos in
   let ty = aelim_type tenv pos ty in
   let e = aelim_expr tenv esc e in
   let inner ty' v' =
      match venv_lookup_var tenv esc pos v1 with
         AggrType _ ->
            raise (IRException (pos, InternalError "aelim_addr_of_var_expr1"))
       | ScalarType (ty, frame1) ->
            SetProject (AtomVar frame1, v1, ty, AtomVar v', e)
       | NoType ->
            LetAtom (v1, ty', AtomVar v', e)
   in
   let v' = new_symbol v2 in
      match venv_lookup_var tenv esc pos v2 with
         AggrType (ty', frame2) ->
            LetAddrOfProject (v', ty', AtomVar frame2, v2, inner ty' v')
       | ScalarType (_, frame2) ->
            LetAddrOfProject (v', ty, AtomVar frame2, v2, inner ty v')
       | NoType ->
            raise (IRException (pos, InternalError "aelim_addr_of_var_expr2"))

(*
 * Set an array entry.
 *)
and aelim_set_subscript_expr tenv esc pos ty a1 a2 a3 e =
   let pos = string_pos "aelim_set_subscript_expr" pos in
   let ty = aelim_type tenv pos ty in
   let e = aelim_expr tenv esc e in
      aelim_atom tenv esc pos a1 (fun a1 ->
            aelim_atom tenv esc pos a2 (fun a2 ->
                  aelim_atom tenv esc pos a3 (fun a3 ->
                        if is_aggregate_type tenv pos ty then
                           let v' = new_symbol_string "aelim_set_subscript_expr" in
                           let ty' = type_of_aggregate tenv pos ty in
                              LetAddrOfSubscript (v', ty', a1, a2,
                                                  Memcpy (AtomVar v', ty, a3, e))
                        else
                           SetSubscript (a1, ty, a2, a3, e))))

and aelim_subscript_expr tenv esc pos v ty a1 a2 e =
   let pos = string_pos "aelim_subscript_expr" pos in
   let ty = aelim_type tenv pos ty in
   let e = aelim_expr tenv esc e in
      aelim_atom tenv esc pos a1 (fun a1 ->
            aelim_atom tenv esc pos a2 (fun a2 ->
                  match venv_lookup_var tenv esc pos v with
                     AggrType (ty', frame) ->
                        let v' = new_symbol v in
                        let v'' = new_symbol v in
                           LetAddrOfSubscript (v', ty', a1, a2,
                                               LetAddrOfProject (v'', ty', AtomVar frame, v,
                                                                 Memcpy (AtomVar v'', ty, AtomVar v', e)))
                   | ScalarType (ty, frame) ->
                        let v' = new_symbol v in
                           LetSubscript (v', ty, a1, a2,
                                         SetProject (AtomVar frame, v, ty, AtomVar v', e))
                   | NoType ->
                        LetSubscript (v, ty, a1, a2, e)))

and aelim_addr_of_subscript_expr tenv esc pos v ty a1 a2 e =
   let pos = string_pos "aelim_addr_of_subscript_expr" pos in
   let ty = aelim_type tenv pos ty in
   let v' = new_symbol v in
   let e = aelim_expr tenv esc e in
   let e =
      match venv_lookup_var tenv esc pos v with
         AggrType _ ->
            raise (IRException (pos, InternalError "aelim_addr_of_subscript_expr"))
       | ScalarType (ty, frame) ->
            SetProject (AtomVar frame, v, ty, AtomVar v', e)
       | NoType ->
            LetAtom (v, ty, AtomVar v', e)
   in
      aelim_atom tenv esc pos a1 (fun a1 ->
            aelim_atom tenv esc pos a2 (fun a2 ->
                  LetAddrOfSubscript (v', ty, a1, a2, e)))

(*
 * Struct projections.
 *)
and aelim_set_project_expr tenv esc pos ty a1 l a2 e =
   let pos = string_pos "aelim_set_project_expr" pos in
   let ty = aelim_type tenv pos ty in
   let e = aelim_expr tenv esc e in
      aelim_atom tenv esc pos a1 (fun a1 ->
            aelim_atom tenv esc pos a2 (fun a2 ->
                  if is_aggregate_type tenv pos ty then
                     let v' = new_symbol_string "aelim_set_project_expr" in
                     let ty' = type_of_aggregate tenv pos ty in
                        LetAddrOfProject (v', ty', a1, l,
                                          Memcpy (AtomVar v', ty, a2, e))
                  else
                     SetProject (a1, l, ty, a2, e)))

and aelim_project_expr tenv esc pos v ty a l e =
   let pos = string_pos "aelim_project_expr" pos in
   let ty = aelim_type tenv pos ty in
   let e = aelim_expr tenv esc e in
      aelim_atom tenv esc pos a (fun a ->
            match venv_lookup_var tenv esc pos v with
               AggrType (ty', frame) ->
                  let v' = new_symbol v in
                  let v'' = new_symbol v in
                     LetAddrOfProject (v', ty', a, l,
                                       LetAddrOfProject (v'', ty', AtomVar frame, v,
                                                         Memcpy (AtomVar v'', ty, AtomVar v', e)))
             | ScalarType (ty, frame) ->
                  let v' = new_symbol v in
                     LetProject (v', ty, a, l,
                                 SetProject (AtomVar frame, v, ty, AtomVar v', e))
             | NoType ->
                  LetProject (v, ty, a, l, e))

and aelim_addr_of_project_expr tenv esc pos v ty a l e =
   let pos = string_pos "aelim_addr_of_project_expr" pos in
   let ty = aelim_type tenv pos ty in
   let v' = new_symbol v in
   let e = aelim_expr tenv esc e in
   let e =
      match venv_lookup_var tenv esc pos v with
         AggrType _ ->
            raise (IRException (pos, InternalError "aelim_addr_of_project_expr"))
       | ScalarType (ty, frame) ->
            SetProject (AtomVar frame, v, ty, AtomVar v', e)
       | NoType ->
            LetAtom (v, ty, AtomVar v', e)
   in
      aelim_atom tenv esc pos a (fun a ->
            LetAddrOfProject (v', ty, a, l, e))

(************************************************************************
 * GLOBAL FUNCTIONS
 ************************************************************************)

(*
 * For the global function:
 *    1. Compute the escape set
 *    2. Rewrite the code
 *)
let aelim_expr tenv e =
   (* Compute the set of escaping variables *)
   let esc = esc_expr tenv empty_venv empty_esc e in

   (* Rewrite the code *)
   let aelim = aelim_of_esc esc in
   let e = aelim_expr tenv aelim e in

   (* Get the outermost data frame *)
   let pos = atom_pos e in
   let frame, _ = pop_frame pos esc in
   let ty_frame = venv_lookup_frame tenv aelim pos frame in
      frame, ty_frame, standardize_expr_notype e