(*
 * fc_ir_ast.ml -- Produce IR for functional C
 * Also performs type checking.
 *
 * Ling Li, Xin Yu
 * Feb. 26, 2001, Caltech
 *)

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

(************************************************************************
 * ENVIRONMENTS
 *
 * Wrap the type and variable environment functions with
 * the normal error handling functions.
 *
 * A type environment (tenv) gives the definition of type
 * identifiers (typedefs).  The tenv_lookup function takes a type
 * identifier (a symbol) and returns the definition.  The tenv_expand
 * function will expand an outermost type identifier.
 *
 * A variable environment gives the type of each bound variable.
 * You will use the variable environment to keep the type of each
 * bound variable.  The venv_lookup function takes a variable
 * (a symbol) and returns the type of that variable.
 *
 * The check_types functions are used to compare types.  The
 * check_types function takes two types, and raises an exception
 * if the two types are not equal.
 ************************************************************************)

(*
 * Default exception.
 *)
let raise_exn pos v =
   SemException (pos, v)

let tenv_lookup = tenv_lookup raise_exn
let tenv_expand = tenv_expand raise_exn
let venv_lookup = venv_lookup raise_exn
let venv_lookup_expand = venv_lookup_expand raise_exn

let check_types = check_types raise_exn
let check_types_list = check_types_list raise_exn
let check_fields = check_fields raise_exn

(*
 * Atom pos labels the file.
 *)
let atom_pos pos =
   string_pos "Fc_ir_ast" (atom_pos pos)

(************************************************************************
 * UTILITIES
 *
 * These are utilities that you may find useful.
 ************************************************************************)

(*
 * Convert an AST binop to an IR binop.
 *)
let binop_of_ast_binop = function
   Fc_ast_type.PlusOp -> PlusOp
 | Fc_ast_type.MinusOp -> MinusOp
 | Fc_ast_type.TimesOp -> MulOp
 | Fc_ast_type.DivideOp -> DivOp
 | Fc_ast_type.ModOp -> RemOp
 | Fc_ast_type.BAndOp -> BAndOp
 | Fc_ast_type.BOrOp -> BOrOp
 | Fc_ast_type.BXorOp -> XorOp
 | Fc_ast_type.LShiftOp -> LslOp
 | Fc_ast_type.RShiftOp -> AsrOp
 | Fc_ast_type.EqOp
 | Fc_ast_type.NotEqOp
 | Fc_ast_type.LeOp
 | Fc_ast_type.LtOp
 | Fc_ast_type.GtOp
 | Fc_ast_type.GeOp ->
      raise (Invalid_argument "Fc_ir_ast.binop_of_ast_binop")

(*
 * Relations are only allowed in IfThenElse,
 * so they are not normal binary operators.
 *)
let relop_of_ast_binop = function
   Fc_ast_type.PlusOp
 | Fc_ast_type.MinusOp
 | Fc_ast_type.TimesOp
 | Fc_ast_type.DivideOp
 | Fc_ast_type.ModOp
 | Fc_ast_type.BAndOp
 | Fc_ast_type.BOrOp
 | Fc_ast_type.BXorOp
 | Fc_ast_type.LShiftOp
 | Fc_ast_type.RShiftOp ->
      raise (Invalid_argument "Fc_ir_ast.relop_of_ast_binop")
 | Fc_ast_type.EqOp -> EqOp
 | Fc_ast_type.NotEqOp -> NEqOp
 | Fc_ast_type.LeOp -> LeOp
 | Fc_ast_type.LtOp -> LtOp
 | Fc_ast_type.GtOp -> GtOp
 | Fc_ast_type.GeOp -> GeOp

let is_binop = function
   Fc_ast_type.PlusOp
 | Fc_ast_type.MinusOp
 | Fc_ast_type.TimesOp
 | Fc_ast_type.DivideOp
 | Fc_ast_type.ModOp
 | Fc_ast_type.BAndOp
 | Fc_ast_type.BOrOp
 | Fc_ast_type.BXorOp
 | Fc_ast_type.LShiftOp
 | Fc_ast_type.RShiftOp ->
      true
 | Fc_ast_type.EqOp
 | Fc_ast_type.NotEqOp
 | Fc_ast_type.LeOp
 | Fc_ast_type.LtOp
 | Fc_ast_type.GtOp
 | Fc_ast_type.GeOp ->
      false

(*
 * Get field offset and size.
 *)
let type_of_field tenv pos v ty =
   let ty = tenv_expand tenv pos ty in
      match ty with
         TyAny
       | TyChar
       | TyInt
       | TyFloat
       | TyArray _
       | TyFun _ ->
            raise (SemException (pos, NotAStruct ty))
       | TyStruct fields ->
            (try List.assoc v fields with
                Not_found ->
                   raise (SemException (pos, UnboundLabel v)))
       | TyId v ->
            raise (Invalid_argument "Fc_ir_ast.type_of_field")

(*
 * Dereferenced pointer type.
 *)
let type_of_deref tenv pos ty =
   match tenv_expand tenv pos ty with
    | TyArray (ty, _) -> ty
    | TyAny | TyChar | TyInt | TyFloat | TyStruct _ | TyFun _ ->
         raise (SemException (pos, NotAPointer ty))
    | TyId _ ->
         raise (Invalid_argument "Fc_ir_ast.type_of_deref")

let type_of_fun tenv venv pos f =
   match venv_lookup_expand tenv venv pos f with
    | TyFun (ar, ty) -> ar, ty
    | ty -> raise (SemException (pos, NotAFunction ty))

(************************************************************************
 * TYPES
 *
 * The make_type function builds an IR type from the AST type.
 * The squash variable is a Boolean value.  If true, array dimensions
 * are omitted (as appropriate in the types of function parameters).
 * Otherwise, array dimensions are preserved (as they should be in
 * a struct type defintion).
 ************************************************************************)

(*
 * Make an IR type from an AST type.
 * If the squash flag is set, all array dimensions
 * are omitted.  It is escpecially important to omit
 * array demensions from function argument types.
 *)
let rec make_type squash tenv = function
   Fc_ast_type.TypeBool  _ ->
      TyInt
 | Fc_ast_type.TypeChar  _ ->
      TyChar
 | Fc_ast_type.TypeInt   _ ->
      TyInt
 | Fc_ast_type.TypeFloat _ ->
      TyFloat
 | Fc_ast_type.TypeId (v, pos) ->
      if squash then
         match tenv_lookup tenv (atom_pos pos) v with
            TyArray (ty, Some _) -> TyArray (ty, None)
          | _ -> TyId v
      else
         TyId v
 | Fc_ast_type.TypeArray (ty, len, pos) ->
      (* Length should be a constant *)
      let len =
         match len with
            Some expr ->
               let i =
                  try Fc_ast_eval.eval_int expr with
                     Fc_ast_eval.EvalException _ ->
                        raise (SemException (atom_pos pos, NotAConstant expr))
               in
                  Some i
          | None ->
               None
      in
      let len =
         if squash then
            None
         else
            len
      in
         TyArray (make_type squash tenv ty, len)
 | Fc_ast_type.TypeStruct (fields, _) ->
      TyStruct (make_type_fields tenv fields)
 | Fc_ast_type.TypeFun (ty_vars, ty_res, _) ->
      TyFun (List.map (make_type true tenv) ty_vars, make_type true tenv ty_res)

and make_type_fields tenv = function
   (v, ty, _) :: fields ->
      (v, make_type false tenv ty) :: make_type_fields tenv fields
 | [] ->
      []

(*
 * Default versions.
 * By default, the make_type function omits array dimensions.
 *)
let make_nosquash_type = make_type false
let make_type = make_type true

(*
 * Remove size from array type.
 *)
let squash_array_type tenv pos ty =
   let ty = tenv_expand tenv pos ty in
      match ty with
         TyArray (ty, Some _) -> TyArray (ty, None)
       | TyArray (_, None) -> ty
       | TyAny
       | TyChar
       | TyInt
       | TyFloat
       | TyStruct _
       | TyFun _ ->
            raise (SemException (pos, NotAPointer ty))
       | TyId _ ->
            raise (Invalid_argument "Fc_ir_ast.squash_array")

(************************************************************************
 * SIMPLE FIR
 *
 * These are some simple AST->IR conversion functions.
 ************************************************************************)

(*
 * Scalar operations.
 *)
let make_scalar op ty tenv venv e cont =
   let v = new_symbol_string "make_scalar" in
   let venv = venv_add venv v ty in
   let ty', exp = cont tenv venv v in
   let exp = LetAtom (v, ty, op e, exp) in
      ty', exp

let make_int = make_scalar (fun i -> AtomInt i) TyInt
let make_char = make_scalar (fun c -> AtomChar c) TyChar
let make_float = make_scalar (fun x -> AtomFloat x) TyFloat

(*
 * String allocation.
 *)
let make_string tenv venv s cont =
   let v = new_symbol_string "make_string" in
   let ty = TyArray (TyChar, None) in
   let venv = venv_add venv v ty in
   let ty, exp = cont tenv venv v in
   let exp = LetString (v, s, exp) in
      ty, exp

(*
 * COERCIONS
 *)

let coerce_cif _str _ty _unop tenv venv pos v cont =
   let pos = string_pos _str pos in
   let ty' = venv_lookup_expand tenv venv pos v in
      (* We do not need to coerce if it has already been that type *)
      if ty' <> _ty then
         let v' = new_symbol_string _str in
         let venv = venv_add venv v' _ty in
         let ty, exp = cont tenv venv v' in
         let exp =
            match ty' with
             | TyAny | TyChar | TyInt | TyFloat ->
                  LetUnop (v', _ty, _unop, AtomVar v, exp)
             | TyArray _ | TyStruct _ | TyFun _ ->
                  raise (SemException (pos, NotAScalar ty'))
             | TyId _ ->
                  raise (Invalid_argument ("Fc_ir_ast." ^ _str))
         in
            ty, exp
      else
         cont tenv venv v

let coerce_char = coerce_cif "coerce_char" TyChar UToChar
let coerce_int = coerce_cif "coerce_int" TyInt UToInt
let coerce_float = coerce_cif "coerce_float" TyFloat UToFloat
let coerce_bool = coerce_cif "coerce_bool" TyInt UToInt

let coerce_pointer tenv venv pos v cont =
   let pos = string_pos "coerce_pointer" pos in
   let ty' = venv_lookup_expand tenv venv pos v in
      match ty' with
       | TyArray _ ->
            cont tenv venv v
       | TyFun _ ->
            let f' = new_symbol_string "coerce_fun_pointer" in
            let ty' = TyArray (ty', None) in
            let venv = venv_add venv f' ty' in
            let ty, exp = cont tenv venv f' in
               ty, LetUnop (f', ty', UToFunPointer, AtomVar v, exp)
       | TyAny | TyChar | TyInt | TyFloat | TyStruct _ ->
            raise (SemException (pos, NotAPointer ty'))
       | TyId _ ->
            raise (Invalid_argument "Fc_ir_ast.coerce_pointer")

let coerce_fun tenv venv pos v cont =
   let pos = string_pos "coerce_fun" pos in
   let ty = venv_lookup_expand tenv venv pos v in
      match ty with
       | TyFun _ -> cont tenv venv v
       | TyArray ((TyFun _) as ty', _) ->
            let v' = new_symbol v in
            let venv = venv_add venv v' ty' in
            let ty, exp = cont tenv venv v' in
               ty, LetUnop (v', ty', UToFun, AtomVar v, exp)
       | TyAny | TyChar | TyInt | TyFloat | TyArray _ | TyStruct _ ->
            raise (SemException (pos, NotAFunction ty))
       | TyId _ ->
            raise (Invalid_argument "Fc_ir_ast.coerce_fun")

(* coerce v to v1 *)
let coerce_to_type tenv venv pos ty v cont =
   match (tenv_expand tenv pos ty) with
    | TyChar  -> coerce_char  tenv venv pos v cont
    | TyInt   -> coerce_int   tenv venv pos v cont
    | TyFloat -> coerce_float tenv venv pos v cont
    | TyArray _ ->
         coerce_pointer tenv venv pos v cont (*??? how to coerce char* to int*? *)
    | TyStruct _ ->
         cont tenv venv v (*??? in fact, not changed *)
    | TyFun _ ->
         coerce_fun tenv venv pos v cont
    | TyAny
    | TyId _  -> raise (Invalid_argument "Fc_ir_ast.coerce_to_type")

let coerce_args tenv venv pos tyl vl cont =
   let pos = string_pos "coerce_args" pos in
   let len1 = List.length tyl in
   let len2 = List.length vl in
   let rec loop vl2 tenv venv tyl vl =
      match tyl, vl with
       | ty :: tyl, v :: vl ->
            coerce_to_type tenv venv pos ty v (fun tenv venv v ->
                  loop (v :: vl2) tenv venv tyl vl)
       | [], [] ->
            cont tenv venv (List.rev vl2)
       | _ -> raise (SemException (pos, ArityMismatch (len1, len2)))
   in
      loop [] tenv venv tyl vl

(* return highest ty *)
let binop_type tenv pos ty1 ty2 =
   match ty1, ty2 with
    | TyArray (ty1', _), TyArray (ty2', _) ->
         check_types tenv pos ty1' ty2';
         ty1, ty2, TyInt
    | TyArray (ty, _), _ -> ty1, TyInt, ty1
    | _, TyArray (_, ty) -> TyInt, ty2, ty2
    | TyFloat, _ | _, TyFloat -> TyFloat, TyFloat, TyFloat
    | TyChar, _ | _, TyChar (* convert to TyInt *)
    | TyInt, _ | _, TyInt -> TyInt, TyInt, TyInt
    | _ -> raise (SemException (pos, BinopTypeError (ty1, ty2)))

let check_binop_type op ty1 ty2 =
   match op, ty1, ty2 with
    | PlusOp, TyArray _, TyArray _
    | MinusOp, TyInt, TyArray _
    | (MulOp | DivOp), TyArray _, _
    | (MulOp | DivOp), _, TyArray _ ->
         false
    | (RemOp | BAndOp | BOrOp | LslOp | LsrOp | AsrOp | XorOp), TyInt, TyInt
    | (RemOp | BAndOp | BOrOp | XorOp), TyChar, TyChar -> (* ??? this should be allowed *)
         true
    | (RemOp | BAndOp | BOrOp | LslOp | LsrOp | AsrOp | XorOp), _, _ ->
         false
    | _, _, _ -> true

(* Return type ??? How to deal with nested function? *)
let return_sym = new_symbol_string "return_type"
let break_sym = new_symbol_string "break_sym"

(************************************************************************
 * IR PRODUCTION
 *
 * Here is the hard part.  You have to implement each of the
 * conversion functions for building the IR from the AST.  The
 * "cont" function represents the "rest" of the expression.
 * See the lab description for details.
 ************************************************************************)

(*
 * Address-of.
 *)
let rec make_addr_of_expr tenv venv pos e cont =
   match e with
    | Fc_ast_type.VarExpr (v, _) ->
         let pos = string_pos "make_addr_of_var" pos in
         let v' = new_symbol_string "addr_of_var" in
         let ty' = TyArray (venv_lookup venv pos v, None) in
         let venv = venv_add venv v' ty' in
         let ty, exp = cont tenv venv v' in
            ty, LetAddrOfVar (v', ty', v, exp)              (* let v':ty' = &v in exp *)
    | Fc_ast_type.SubscriptExpr (e1, e2, _) ->
         let pos = string_pos "make_addr_of_subscript_expr" pos in
         make_expr tenv venv e1 (fun tenv venv v1 ->
         coerce_pointer tenv venv pos v1 (fun tenv venv v1 ->
            make_expr tenv venv e2 (fun tenv venv v2 ->
            coerce_int tenv venv pos v2 (fun tenv venv v2 ->
               let v' = new_symbol_string "addr_of_subscript" in
               let ty' = squash_array_type tenv pos (venv_lookup venv pos v1) in
               let venv = venv_add venv v' ty' in
               let ty, exp = cont tenv venv v' in
               let exp = LetAddrOfSubscript (v', ty', AtomVar v1, AtomVar v2, exp) in
                  ty, exp))))                            (* let v':ty' = &v1[v2] in exp *)
    | Fc_ast_type.ProjectExpr (e, v, _) ->
         let pos = string_pos "make_addr_of_project_expr" pos in
            make_addr_of_expr tenv venv pos e (fun tenv venv v1 ->
               let v' = new_symbol_string "addr_of_project" in
               let ty' = type_of_deref tenv pos (venv_lookup venv pos v1) in
               let ty' = type_of_field tenv pos v ty' in
               let ty' = TyArray (ty', None) in
               let venv = venv_add venv v' ty' in
               let ty, exp = cont tenv venv v' in
               let exp = LetAddrOfProject (v', ty', AtomVar v1, v, exp) in
                  ty, exp)                            (* let v':ty' = &v1.v in exp *)
    | Fc_ast_type.UnOpExpr (Fc_ast_type.UStarOp, e, _) ->
         let pos = string_pos "make_addr_of_deref_expr" pos in
            make_expr tenv venv e (fun tenv venv v ->
                  cont tenv venv v)
    | _ -> raise (SemException (pos, NotAddressible e))

(*
 * Unary operations. ??? Checking needed
 *)
and make_unop_expr tenv venv pos op e cont =
   make_expr tenv venv e (fun tenv venv v ->
      let pos = string_pos "make_unop_expr" pos in
      let v' = new_symbol_string "unop_expr" in
      let ty = venv_lookup_expand tenv venv pos v in
      let let_unop op ty' =
         let venv = venv_add venv v' ty' in
         let ty, exp = cont tenv venv v' in
            ty, LetUnop (v', ty', op, AtomVar v, exp)
      in
         match op, ty with
          | Fc_ast_type.UMinusOp, TyFloat ->
               let_unop UMinusOp TyFloat
          | Fc_ast_type.UMinusOp, (TyChar | TyInt) ->
               let_unop UMinusOp TyInt
          | Fc_ast_type.UNotOp, _ ->
               coerce_int tenv venv pos v (fun tenv venv v ->
                  let ty, exp = cont tenv venv v in
                  let f = new_symbol_string "unot_expr_fun" in
                  let fun_ty = TyFun ([TyInt], ty) in
                  let exp =
                     LetFuns ([f, false, fun_ty, [v], exp],
                              IfThenElse (EqOp, AtomVar v, AtomInt 0,
                                          TailCall (f, [AtomInt 1]), TailCall (f, [AtomInt 0])))
                  in
                     ty, exp)
          | Fc_ast_type.UStarOp, TyArray (ty', _) ->
               let venv = venv_add venv v' ty' in
               let ty, exp = cont tenv venv v' in
                  ty, LetSubscript (v', ty', AtomVar v, AtomInt 0, exp)
          | _ -> raise (SemException (pos, NotAPrefixValue ty))
   )

(*
 * Pre- and postfix operations.
 *)
and make_uarith_expr tenv venv pos op e cont =
   let pos = string_pos "make_uarith_expr" pos in
      match e with
       | Fc_ast_type.VarExpr (v, _) ->
            make_uarith_var tenv venv pos op v cont
       | _ ->
            make_addr_of_expr tenv venv pos e (fun tenv venv v ->
                  make_uarith_pointer tenv venv pos op v cont)

and make_uarith_type venv pos op ty' cont =
   let v' = new_symbol_string "uarith_expr" in
   let venv = venv_add venv v' ty' in
   let a =
      match ty' with
       | TyInt | TyChar | TyArray _ -> AtomInt 1
       | TyFloat -> AtomFloat 1.0
       | _ -> raise (SemException (pos, NotAScalar ty'))
   in
   let post, op =
      match op with
       | Fc_ast_type.PreIncrOp -> false, PlusOp
       | Fc_ast_type.PreDecrOp -> false, MinusOp
       | Fc_ast_type.PostIncrOp -> true, PlusOp
       | Fc_ast_type.PostDecrOp -> true, MinusOp
   in
      cont venv post op v' a

and make_uarith_var tenv venv pos op v cont =
   let pos = string_pos "make_uarith_var" pos in
   let a1 = AtomVar v in
   let ty' = venv_lookup venv pos v in
      make_uarith_type venv pos op ty' (fun venv post op v' a2 ->
            let ty, exp = cont tenv venv v' in
               if post then
                  let v'' = new_symbol_string "uarith_var_temp" in
                     ty, LetAtom (v', ty', a1,
                         LetBinop (v'', ty', op, a1, a2,
                         SetVar (v, ty', AtomVar v'', exp)))
               else
                  ty, LetBinop (v', ty', op, a1, a2,
                      SetVar (v, ty', AtomVar v', exp))
      )

(* other than a variable *)
and make_uarith_pointer tenv venv pos op v cont =
   let pos = string_pos "make_uarith_pointer" pos in
   let a1 = AtomVar v in
   let ty' = type_of_deref tenv pos (venv_lookup venv pos v) in
      make_uarith_type venv pos op ty' (fun venv post op v' a2 ->
            let ty, exp = cont tenv venv v' in
               if post then
                  let v'' = new_symbol_string "uarith_pointer_temp" in
                     ty, LetSubscript (v', ty', a1, AtomInt 0,
                         LetBinop (v'', ty', op, AtomVar v', a2,
                         SetSubscript (a1, ty', AtomInt 0, AtomVar v'', exp)))
               else
                  ty, LetSubscript (v', ty', a1, AtomInt 0,
                      LetBinop (v', ty', op, AtomVar v', a2,
                      SetSubscript (a1, ty', AtomInt 0, AtomVar v', exp)))

      )

(*
 * Binary operations.
 *)
and make_binop_expr tenv venv pos op e1 e2 cont =
   make_expr tenv venv e1 (fun tenv venv v1 ->
      make_expr tenv venv e2 (fun tenv venv v2 ->
         let pos = string_pos "make_binop_expr" pos in
         let ty1 = venv_lookup_expand tenv venv (int_pos 1 pos) v1 in
         let ty2 = venv_lookup_expand tenv venv (int_pos 2 pos) v2 in
         if is_binop op then
            let op = binop_of_ast_binop op in
            let ty1', ty2', ty' = binop_type tenv (int_pos 3 pos) ty1 ty2 in
            let v' = new_symbol_string "binop_expr" in
            let venv = venv_add venv v' ty' in
            let ty, exp = cont tenv venv v' in
               coerce_to_type tenv venv (int_pos 4 pos) ty1' v1 (fun tenv venv v1 ->
                     coerce_to_type tenv venv (int_pos 5 pos) ty2' v2 (fun tenv venv v2 ->
                           if check_binop_type op ty1' ty2' then
                              ty', LetBinop (v', ty', op, AtomVar v1, AtomVar v2, exp)
                           else
                              raise (SemException (pos, BinopTypeError (ty1, ty2)))
         ))
         else
            let op = relop_of_ast_binop op in
            let v' = new_symbol_string "relop_expr" in
            let venv = venv_add venv v' TyInt in
            let ty, exp = cont tenv venv v' in
            let f = new_symbol_string "relop_expr_fun" in
            let fun_ty = TyFun ([TyInt], ty) in
            let exp =
               LetFuns ([f, false, fun_ty, [v'], exp],
                        IfThenElse (op, AtomVar v1, AtomVar v2,
                                    TailCall (f, [AtomInt 1]), TailCall (f, [AtomInt 0])))
            in
                ty, exp
   ))

(*
 * Boolean relations are short-circuit.
 *)
and make_boolop_expr tenv venv pos op e1 e2 cont =
   let t = Fc_ast_type.IntExpr (1, pos) in
   let f = Some (Fc_ast_type.IntExpr (0, pos)) in
   let e =
      match op with
       | Fc_ast_type.LAndOp ->
            Fc_ast_type.IfExpr (e1, Fc_ast_type.IfExpr (e2, t, f, pos), f, pos)
       | Fc_ast_type.LOrOp ->
            Fc_ast_type.IfExpr (e1, t, Some (Fc_ast_type.IfExpr (e2, t, f, pos)), pos)
   in
      make_expr tenv venv e cont

(*
 * Subscript.
 *)
and make_subscript_expr tenv venv pos e1 e2 cont =
   make_expr tenv venv e1 (fun tenv venv v1 ->
      make_expr tenv venv e2 (fun tenv venv v2 ->
         let pos = string_pos "make_subscript_expr" pos in
            coerce_pointer tenv venv pos v1 (fun tenv venv v1 ->
               coerce_int tenv venv pos v2 (fun tenv venv v2 ->
                  let v' = new_symbol_string "subscript_expr" in
                  let ty' = venv_lookup venv pos v1 in
                  let ty' = type_of_deref tenv pos ty' in
                  let venv = venv_add venv v' ty' in
                  let ty, exp = cont tenv venv v' in (* let v':ty' = v1[v2] in exp *)
                     ty, LetSubscript (v', ty', AtomVar v1, AtomVar v2, exp)))))

(*
 * Struct projection.
 *)
and make_project_expr tenv venv pos e1 label cont =
   let pos = string_pos "make_project_expr" pos in
      make_addr_of_expr tenv venv pos e1 (fun tenv venv v1 ->
            let v' = new_symbol_string "project_expr" in
            let ty' = type_of_deref tenv pos (venv_lookup venv pos v1) in
            let ty' = type_of_field tenv pos label ty' in
            let venv = venv_add venv v' ty' in
            let ty, exp = cont tenv venv v' in (* let v':ty' = v1->label in exp *)
               ty, LetProject (v', ty', AtomVar v1, label, exp))

(*
 * Function application.
 *)
and make_apply_expr tenv venv pos f args cont =
   let pos = string_pos "make_apply_expr" pos in
      make_expr tenv venv f (fun tenv venv f ->
      coerce_fun tenv venv pos f (fun tenv venv f ->
            make_list_expr tenv venv args (fun tenv venv args ->
                  let tyarg, ty' = type_of_fun tenv venv pos f in
                  let v' = new_symbol_string "apply_expr" in
                  let venv = venv_add venv v' ty' in
                  let ty, exp = cont tenv venv v' in
                     coerce_args tenv venv pos tyarg args (fun tenv venv args ->
                        ty, LetApply (v', ty', f, List.map (fun v -> AtomVar v) args, exp))
      )))

(*
 * Assignment.
 *)
and make_assign_expr tenv venv pos op e1 e2 cont =
   make_expr tenv venv e2 (fun tenv venv v2 ->
      match e1 with
       | Fc_ast_type.VarExpr (v1, _) ->
            let pos = string_pos "make_set_var_expr" pos in
            let ty' = venv_lookup venv pos v1 in
            let v' = new_symbol_string "assign_var" in
            let venv = venv_add venv v' ty' in
            let ty, exp = cont tenv venv v' in
               coerce_to_type tenv venv pos ty' v2 (fun tenv venv v2 ->
                  let exp =
                     match op with
                      | Some op ->
                           LetBinop (v', ty', (binop_of_ast_binop op),
                              AtomVar v1, AtomVar v2, SetVar (v1, ty', AtomVar v', exp))
                      | None ->
                           LetAtom (v', ty', AtomVar v2, SetVar (v1, ty', AtomVar v', exp))
                  in
                     ty, exp)
       | Fc_ast_type.SubscriptExpr (e3, e2, _) ->
            let pos = string_pos "make_set_subscript_expr" pos in
               make_expr tenv venv e3 (fun tenv venv a1 ->
               make_expr tenv venv e2 (fun tenv venv a2 ->
               coerce_int tenv venv pos a2 (fun tenv venv a2 ->
                  let ty' = type_of_deref tenv pos (venv_lookup venv pos a1) in
                  let v' = new_symbol_string "assign_subscript" in
                  let venv = venv_add venv v' ty' in
                     coerce_to_type tenv venv pos ty' v2 (fun tenv venv v2 ->
                        let ty, exp = cont tenv venv v2 in
                        let exp = SetSubscript (AtomVar a1, ty', AtomVar a2, AtomVar v', exp) in
                           match op with
                            | Some op ->
                                 make_expr tenv venv e1 (fun tenv venv v1 ->
                                    ty, LetBinop (v', ty', (binop_of_ast_binop op),
                                                  AtomVar v1, AtomVar v2, exp))
                            | None ->
                                 ty, LetAtom (v', ty', AtomVar v2, exp)
               ))))
       | Fc_ast_type.ProjectExpr (e, label, _) ->
            let pos = string_pos "make_set_project_expr" pos in
               make_addr_of_expr tenv venv pos e (fun tenv venv v ->
                  let ty' = type_of_deref tenv pos (venv_lookup venv pos v) in
                  let ty' = type_of_field tenv pos label ty' in
                  let v' = new_symbol_string "assign_project" in
                  let venv = venv_add venv v' ty' in
                     coerce_to_type tenv venv pos ty' v2 (fun tenv venv v2 ->
                        let ty, exp = cont tenv venv v2 in
                        let exp = SetProject (AtomVar v, label, ty', AtomVar v2, exp) in
                           match op with
                            | Some op ->
                                 make_expr tenv venv e1 (fun tenv venv v1 ->
                                    ty, LetBinop (v', ty', (binop_of_ast_binop op),
                                                  AtomVar v1, AtomVar v2, exp))
                            | None ->
                                 ty, LetAtom (v', ty', AtomVar v2, exp)
               ))
       | Fc_ast_type.UnOpExpr (Fc_ast_type.UStarOp, e, pos) ->
            let e1 = Fc_ast_type.SubscriptExpr (e, Fc_ast_type.IntExpr (0, pos), pos) in
               make_assign_expr tenv venv (atom_pos pos) op e1 e2 cont

       | _ -> raise (SemException (pos, NotAddressible e1))
      )

(*
 * Conditional.
 *)
and make_if_expr tenv venv pos e1 e2 e3 cont =
   (* Wrap the continuation in a function call *)
   let pos = string_pos "make_if_expr" pos in
   let f = new_symbol_string "make_if_cont_fun" in
   let v = new_symbol_string "make_if_cont_arg" in
   (* Get type and expression for true case *)
   let true_ty, true_exp =
      make_expr tenv venv e2 (fun tenv venv v2 ->
            let ty = venv_lookup venv (int_pos 1 pos) v2 in
               ty, TailCall (f, [AtomVar v2]))
   in
   (* Get type and expression for false case *)
   let false_ty, false_exp =
      match e3 with
       | Some e3 ->
            make_expr tenv venv e3 (fun tenv venv v3 ->
                  let ty = venv_lookup venv (int_pos 2 pos) v3 in
                     ty, TailCall (f, [AtomVar v3]))
       | None ->
            TyInt, TailCall (f, [AtomInt 0])
   in
   (* Make sure types of both cases are the same ??? *)
   let _ = check_types tenv (int_pos 3 pos) true_ty false_ty in

   (* Build the function body *)
   let ty, body =
      let venv = venv_add venv v true_ty in
         cont tenv venv v
   in
   let fun_ty = TyFun ([true_ty], ty) in
      (* Build the expression *)
      make_expr tenv venv e1 (fun tenv venv v1 ->
            coerce_bool tenv venv (int_pos 4 pos) v1 (fun tenv venv v1 ->
                  let exp =
                     LetFuns ([f, false, fun_ty, [v], body],
                              IfThenElse (EqOp, AtomVar v1, AtomInt 0,
                                          false_exp, true_exp))
                  in
                     ty, exp))

(*
 * For loop.  Just transform it into a while loop.
 *)
and make_for_expr tenv venv pos e1 e2 e3 e4 cont =
   let e =
      Fc_ast_type.SeqExpr (**)
         ([e1; Fc_ast_type.WhileExpr (**)
              (e2, Fc_ast_type.SeqExpr ([e4; e3], pos), pos)], pos)
   in
      make_expr tenv venv e cont

(*
 * While loop.
 * This becomes a recursive function.
 *)
and make_while_expr tenv venv pos e1 e2 cont =
   (* Put the rest of the program in a function called "break" *)
   let pos = string_pos "make_while_expr" pos in
   let whl = new_symbol_string "make_while_body" in
   (* Get type and expression for while break *)
   let brk_ty, brk_exp =
       let v = new_symbol_string "make_while_expr" in
       let venv = venv_add venv v TyInt in
       let ty, exp = cont tenv venv v in
          ty, LetAtom (v, TyInt, AtomInt 0, exp) (* use 0 as a 'unuse' var ??? *)
   in
   let brkfun_ty = TyFun ([], brk_ty) in
   (* Get expression for while body *)
   let _, whl_body_exp =
      let venv = venv_add venv break_sym brkfun_ty in
         make_expr tenv venv e2 (fun tenv venv v ->
               brk_ty, TailCall (whl, []))
   in
   let _, whl_exp =
      make_expr tenv venv e1 (fun tenv venv v1 ->
            coerce_bool tenv venv (int_pos 4 pos) v1 (fun tenv venv v1 ->
                  brk_ty, IfThenElse (EqOp, AtomVar v1, AtomInt 0,
                          TailCall (break_sym, []), whl_body_exp)))
   in
   let whlfun_ty = TyFun ([], brk_ty) in
   let whlfun =
      LetFuns ([whl, false, whlfun_ty, [], whl_exp], TailCall (whl, []))
   in
      brk_ty, LetFuns ([break_sym, false, brkfun_ty, [], brk_exp], whlfun)

(*
 * Sequential execution.
 * Have to wrap the continuation in a function
 * so that environment will be correct.
 *)
and make_seq_expr tenv venv pos el cont =
   let pos = string_pos "make_seq_expr" pos in
   let rest = new_symbol_string "make_seq_expr_fun" in
   let v = new_symbol_string "make_seq_expr_arg" in
   match el with
    | [] ->
         let venv = venv_add venv v TyInt in
         let rest_ty, rest_exp = cont tenv venv v in
            rest_ty, LetAtom (v, TyInt, AtomInt 0, rest_exp)
    | _ ->
   let ty, _ = (* to get the type of the last exp. Really ugly ??? *)
      make_list_expr tenv venv el (fun tenv venv v ->
            let vlast = List.hd (List.rev v) in
            let ty = venv_lookup venv pos vlast in
               ty, TailCall (rest, [AtomVar vlast]))
   in
   let venv = venv_add venv v ty in
   let rest_ty, rest_exp =
       cont tenv venv v
   in
   let _, exp = (* to generate the real code... since here we add TailCall, we need to give a correct type ???*)
      make_list_expr tenv venv el (fun tenv venv v ->
            let vlast = List.hd (List.rev v) in
               rest_ty, TailCall (rest, [AtomVar vlast]))
   in
   let fun_ty = TyFun ([ty], rest_ty) in
      rest_ty, LetFuns ([rest, false, fun_ty, [v], rest_exp], exp)
(* what's the type of LetFuns? ty? or fun_ty? *)

(*
 * Return from the current function.
 *)
and make_return_expr tenv venv pos e =
   make_expr tenv venv e (fun tenv venv v ->
         let pos = string_pos "make_return_expr" pos in
         let ret_ty = venv_lookup venv pos return_sym in
            coerce_to_type tenv venv pos ret_ty v (fun tenv venv v ->
                  ret_ty, Return (AtomVar v)))

(*
 * Break from the current loop.
 * Just passes control to the "break" function.
 *)
and make_break_expr tenv venv pos cont =
   let pos = string_pos "make_break_expr" pos in
      TyAny, TailCall (break_sym, []) (*??? in fact I should lookup the type of brk *)

(*
 * Variable definitions.
 *)
and make_var_def_expr tenv venv (v, ty, e_opt, pos) cont =
   let pos = string_pos "make_var_def_expr" (atom_pos pos) in
   let ty' = make_nosquash_type tenv ty in
   let ty, exp =
      let venv = venv_add venv v ty' in
         cont tenv venv v
   in
      (* If this is a fun decl, build an external call *)
      match ty' with
         TyFun (ty_vars, ty_res) ->
            (* Build the parameter list *)
            let vars = List.map (fun _ -> new_symbol_string "v") ty_vars in
            let args = List.map (fun v -> AtomVar v) vars in
            let v' = new_symbol v in

            (* External function name is prefixed by % *)
            let f = Symbol.add ("%" ^ Symbol.to_string v) in
               ty, LetFuns ([v, true, ty', vars,
                   LetExtCall (v', f, ty', args, Return (AtomVar v'))], exp)
       | _ ->
            (* If its not a fun decl, use LetCopy to initialize the value *)
            make_expr_opt tenv venv pos ty' e_opt (fun tenv venv a_opt ->
                   ty, LetCopy (v, ty', a_opt, exp))

and make_var_defs_expr tenv venv pos defs cont =
   match defs with
    | def :: defs ->
         make_var_def_expr tenv venv def (fun tenv venv _ ->
               make_var_defs_expr tenv venv pos defs cont)
    | [] ->
         let v = new_symbol_string "make_var_defs" in
         let venv = venv_add venv v TyInt in
         let ty, exp = cont tenv venv v in
            ty, LetAtom(v, TyInt, AtomInt 0, exp)

(*
 * Function declaration just modifies the var environment.
 *)
and make_fun_decl_expr tenv venv pos f ty_vars ty_res cont =
   let pos = string_pos "make_fun_decl_expr" pos in
      raise (SemException (pos, NotImplemented))

(*
 * Function definition.
 *)
and make_fun_def_expr tenv venv pos f vars ty_res body cont =
   (* Function type *)
   let pos = string_pos "make_fun_def_expr" pos in
   let ret_ty = make_type tenv ty_res in
   let arg_list = List.map (fun (sym, _, _) -> sym) vars in
   let arg_ty = List.map (fun (_, ty, _) -> make_type tenv ty) vars in
   let fun_ty = TyFun (arg_ty, ret_ty) in
   let venv = venv_add venv f fun_ty in
   let ty, exp = cont tenv venv f in
   let _, body =
      let venv = venv_add venv return_sym ret_ty in
      let venv = List.fold_left2 venv_add venv arg_list arg_ty in
         make_return_expr tenv venv pos body
   in
      ty, LetFuns ([f, true, fun_ty, arg_list, body], exp)

(*
 * Type definitions.
 *)
and make_type_defs_expr tenv venv pos types cont =
   let pos = string_pos "make_type_defs_expr" pos in
   let typ = List.map (fun (s, t, _) -> s, make_type tenv t) types in
   let tenv =
      List.fold_left (fun tenv (v, ty) -> tenv_add tenv v ty) tenv typ
   in
   let ty, exp = cont tenv venv (new_symbol_string "make_type_defs_expr_nouse") in
      ty, LetTypes (typ, exp)

(*
 * Produce the intermediate code.
 *)
and make_expr tenv venv e cont =
   match e with
      Fc_ast_type.CharExpr (c, _) ->
         make_char tenv venv c cont
    | Fc_ast_type.IntExpr (i, _) ->
         make_int tenv venv i cont
    | Fc_ast_type.FloatExpr (x, _) ->
         make_float tenv venv x cont
    | Fc_ast_type.StringExpr (s, _) ->
         make_string tenv venv s cont
    | Fc_ast_type.VarExpr (v, _) ->
         cont tenv venv v
    | Fc_ast_type.AddrOfExpr (e, pos) ->
         make_addr_of_expr tenv venv (atom_pos pos) e cont
    | Fc_ast_type.UnOpExpr (op, e, pos) ->
         make_unop_expr tenv venv (atom_pos pos) op e cont
    | Fc_ast_type.UArithExpr (op, e, pos) ->
         make_uarith_expr tenv venv (atom_pos pos) op e cont
    | Fc_ast_type.BoolOpExpr (op, e1, e2, pos) ->
         make_boolop_expr tenv venv pos op e1 e2 cont
    | Fc_ast_type.BinOpExpr (op, e1, e2, pos) ->
         make_binop_expr tenv venv (atom_pos pos) op e1 e2 cont
    | Fc_ast_type.SubscriptExpr (e1, e2, pos) ->
         make_subscript_expr tenv venv (atom_pos pos) e1 e2 cont
    | Fc_ast_type.ProjectExpr (e, v, pos) ->
         make_project_expr tenv venv (atom_pos pos) e v cont
    | Fc_ast_type.ApplyExpr (f, args, pos) ->
         make_apply_expr tenv venv (atom_pos pos) f args cont
    | Fc_ast_type.AssignExpr (op, e1, e2, pos) ->
         make_assign_expr tenv venv (atom_pos pos) op e1 e2 cont
    | Fc_ast_type.IfExpr (e1, e2, e3, pos) ->
         make_if_expr tenv venv (atom_pos pos) e1 e2 e3 cont
    | Fc_ast_type.ForExpr (e1, e2, e3, e4, pos) ->
         make_for_expr tenv venv pos e1 e2 e3 e4 cont
    | Fc_ast_type.WhileExpr (e1, e2, pos) ->
         make_while_expr tenv venv (atom_pos pos) e1 e2 cont
    | Fc_ast_type.SeqExpr (el, pos) ->
         make_seq_expr tenv venv (atom_pos pos) el cont
    | Fc_ast_type.ReturnExpr (e, pos) ->
         make_return_expr tenv venv (atom_pos pos) e
    | Fc_ast_type.BreakExpr pos ->
         make_break_expr tenv venv (atom_pos pos) cont
    | Fc_ast_type.VarDefs (defs, pos) ->
         make_var_defs_expr tenv venv (atom_pos pos) defs cont
    | Fc_ast_type.FunDecl (f, ty_vars, ty_res, pos) ->
         make_fun_decl_expr tenv venv (atom_pos pos) f ty_vars ty_res cont
    | Fc_ast_type.FunDef (f, vars, ty_res, body, pos) ->
         make_fun_def_expr tenv venv (atom_pos pos) f vars ty_res body cont
    | Fc_ast_type.TypeDefs (types, pos) ->
         make_type_defs_expr tenv venv (atom_pos pos) types cont

(*
 * Convert a list of expressions.
 *)
and make_list_expr tenv venv el cont =
   let rec loop tenv venv vars = function
      e :: el ->
         make_expr tenv venv e (fun tenv venv v ->
               loop tenv venv (v :: vars) el)
    | [] ->
         cont tenv venv (List.rev vars)
   in
      loop tenv venv [] el

(*
 * Convert an optional expression.
 *)
and make_expr_opt tenv venv pos ty e_opt cont =
   match e_opt with
      None ->
         cont tenv venv None
    | Some e ->
         make_expr tenv venv e (fun tenv venv v ->
               coerce_to_type tenv venv pos ty v (fun tenv venv v ->
                     cont tenv venv (Some (AtomVar v))))

(************************************************************************
 * GLOBAL FUNCTIONS
 *
 * The make_prog function converts a program (an expr list).
 * The make_expr function converts a single expression.
 * Both functions return the innermost type and variable
 * environment.
 ************************************************************************)

(*
 * Apply to a program.
 *)
let make_prog tenv venv el =
   let tenv_ref = ref tenv in
   let venv_ref = ref venv in
   let rec loop tenv venv vars = function
      e :: el ->
         make_expr tenv venv e (fun tenv venv v ->
               loop tenv venv (v :: vars) el)
    | [] ->
         tenv_ref := tenv;
         venv_ref := venv;
         TyInt, Return (AtomInt 0)
   in
   let ty, exp = loop tenv venv [] el in
   let exp = standardize_expr exp in
      !tenv_ref, !venv_ref, ty, exp

(*
 * Real make_expr with default continuation.
 *)
let make_expr tenv venv e =
   let ref_tenv = ref tenv in
   let ref_venv = ref venv in
   let ty, exp =
      make_expr tenv venv e (fun tenv venv v ->
            let pos = Fc_ast_util.pos_of_expr e in
            let pos = string_pos "make_expr_global" (atom_pos pos) in
            let ty = venv_lookup venv pos v in
               ref_tenv := tenv;
               ref_venv := venv;
               ty, Return (AtomVar v))
   in
   let exp = standardize_expr exp in
      !ref_tenv, !ref_venv, ty, exp