(* Code generation: translate takes a semantically checked AST and
produces LLVM IR

LLVM tutorial: Make sure to read the OCaml version of the tutorial

http://llvm.org/docs/tutorial/index.html

Detailed documentation on the OCaml LLVM library:

http://llvm.moe/
http://llvm.moe/ocaml/

This code was original from MicroC and modified by:
Keir Lauritzen, Leon Song, and Guy Yardeni

*)

module L = Llvm
module S = Sast

module StringMap = Map.Make(String)

let context = L.global_context () ;;
let the_module = L.create_module context "Carl"
and i32_t   = L.i32_type  context
and i8_t    = L.i8_type   context
and i1_t    = L.i1_type   context
and f_t     = L.double_type context
and void_t  = L.void_type context;;

let void_p  = L.pointer_type i8_t;;
let str_t   = L.pointer_type i8_t ;;
let charpp_t = L.pointer_type str_t;;
let intp_t = L.pointer_type i32_t;;
let arr_t   = L.i64_type context ;;
let action_t = L.pointer_type (L.function_type void_t [| i32_t; charpp_t; intp_t |]);;
(* let actions_t = L.pointer_type (L.pointer_type (L.function_type void_t [| i32_t; |]));; *)

(*L.dump_type str_t;;*)

let translate program =

    let ltype_of_typ = function
    | S.Int -> i32_t
    | S.Bool -> i1_t
    | S.String_t -> str_t
    | S.Void -> void_t 
    | S.Action_t -> action_t 
    | S.Charpp -> charpp_t
    | S.Intp -> intp_t
    | S.Char -> i8_t
    | S.Float -> f_t
    | S.Array_f -> void_p
    | S.Array_s -> void_p in

   (* | S.Float -> i32_t *)
   (* | S.Array -> arr_t *) 

    let rec global_init = function    
    | S.Int -> L.const_int i32_t 0
    | S.Bool -> L.const_int i1_t 0
    | S.Char -> L.const_int i8_t 0
    | S.String_t -> L.const_pointer_null str_t
    | S.Void -> L.const_int void_t 0
    | S.Float -> L.const_float f_t 0.0
    | S.Action_t -> L.const_pointer_null action_t
    | S.S_arr_t (t, i) -> L.const_array (ltype_of_typ t) (Array.make i (global_init t))
    | S.Array_f -> L.const_pointer_null void_p
    | S.Array_s -> L.const_pointer_null void_p in


    (* Declare each global variable; remember its value in a map *)
    let global_vars =
        let global_var m (t, n) =
            let init = global_init t in 
        StringMap.add n (L.define_global n init the_module) m in
    List.fold_left global_var StringMap.empty program.Sast.var_decls in

    (* Declare printf(), which the print built-in function will call *)
    let printf_t = L.var_arg_function_type i32_t [| L.pointer_type i8_t |] in
    let printf_func = L.declare_function "printf" printf_t the_module in

    let strlen_t = L.function_type f_t [| str_t |] in
    let strlen_func = L.declare_function "strlen2" strlen_t the_module in

    let create_t = L.function_type void_p [||] in
    let create_func = L.declare_function "create" create_t the_module in

    let array_add_string_t = L.function_type i32_t [|void_p ; str_t ; str_t|] in 
    let array_add_string_func = L.declare_function "array_add_string" array_add_string_t the_module in

    let array_retrieve_string_t = L.function_type str_t [|void_p ; str_t|] in 
    let array_retrieve_string_func = L.declare_function "array_retrieve_string" array_retrieve_string_t the_module in

    let array_add_float_t = L.function_type i32_t [|void_p ; str_t ; f_t|] in 
    let array_add_float_func = L.declare_function "array_add_float" array_add_float_t the_module in

    let array_retrieve_float_t = L.function_type f_t [|void_p ; str_t|] in 
    let array_retrieve_float_func = L.declare_function "array_retrieve_float" array_retrieve_float_t the_module in

    (* Define each function (arguments and return type) so we can call it *)
    let function_decls =
        let function_decl m fdecl =
            let name = fdecl.S.fname
            and formal_types =
	        Array.of_list (List.map (fun (t,_) -> ltype_of_typ t) fdecl.S.formals) in 
            let ftype = L.function_type (ltype_of_typ fdecl.S.typ) formal_types in
        StringMap.add name (L.define_function name ftype the_module, fdecl) m in
    List.fold_left function_decl StringMap.empty program.S.func_decls in
  
    (* Fill in the body of the given function *)
    let build_function_body fdecl =
        let (the_function, _) = StringMap.find fdecl.S.fname function_decls in
        let builder = L.builder_at_end context (L.entry_block the_function) in

        let int_format_str = L.build_global_stringptr "%d\n" "fmt" builder in
        let str_format_str = L.build_global_stringptr "%s\n" "fmt2" builder in
        let flt_format_str = L.build_global_stringptr "%f\n" "fmt3" builder in
    
        (* Construct the function's "locals": formal arguments and locally
        declared variables.  Allocate each on the stack, initialize their
        value, if appropriate, and remember their values in the "locals" map *)
        let local_vars =
            let add_formal m (t, n) p = L.set_value_name n p;
            let local = L.build_alloca (ltype_of_typ t) n builder in
	    ignore (L.build_store p local builder);
	    StringMap.add n local m in

        let add_local m (t, n) =
	        let local_var = L.build_alloca (ltype_of_typ t) n builder in 
        StringMap.add n local_var m in

        let formals = List.fold_left2 add_formal StringMap.empty fdecl.S.formals
            (Array.to_list (L.params the_function)) in
        List.fold_left add_local formals fdecl.S.locals in

        (* Return the value for a variable or formal argument *)
        let lookup n =  try StringMap.find n local_vars
                        with Not_found -> try StringMap.find n global_vars
                         with Not_found -> (*ignore(print_endline "in lookup");*) StringMap.find n local_vars
        (*with Not_found -> fst (StringMap.find n function_decls) *)in

        let lookup_func n = fst (StringMap.find n function_decls) in

        let build_string e builder = 
            let str = L.build_global_stringptr e "str" builder in
            let null = L.const_int i32_t 0 in
        L.build_in_bounds_gep str [| null |] "str" builder in

        
        (* Construct code for an expression; return its value *)
        let rec expr builder = function
        | S.Int_Lit i -> L.const_int i32_t i
        | S.Bool_Lit b -> L.const_int i1_t (if b then 1 else 0)
        | S.String_Lit s -> build_string s builder
        | S.Float_Lit f -> L.const_float f_t f
        | S.Char_Lit c -> L.const_int i8_t (Char.code c)
        | S.Noexpr -> L.const_int i32_t 0
        | S.Id s -> L.build_load (lookup s) s builder
        | S.Binop (e1, op, e2) -> (*(try
	        (let e1' = expr builder e1
	        and e2' = expr builder e2 in
	        (match op with
	        | S.Add     -> L.build_add
	        | S.Sub     -> L.build_sub
	        | S.Mult    -> L.build_mul
            | S.Div     -> L.build_sdiv
	        | S.And     -> L.build_and
	        | S.Or      -> L.build_or
	        | S.Equal   -> L.build_icmp L.Icmp.Eq
	        | S.Neq     -> L.build_icmp L.Icmp.Ne
	        | S.Less    -> L.build_icmp L.Icmp.Slt
	        | S.Leq     -> L.build_icmp L.Icmp.Sle
	        | S.Greater -> L.build_icmp L.Icmp.Sgt
	        | S.Geq     -> L.build_icmp L.Icmp.Sge
	        ) e1' e2' "tmp" builder) with _ ->
            ignore(print_endline "INEXPR2"); *)
            let e1' = expr builder e1
            and e2' = expr builder e2 in
            (match op with
            | S.Add     -> L.build_fadd
            | S.Sub     -> L.build_fsub
            | S.Mult    -> L.build_fmul
            | S.Div     -> L.build_fdiv
            | S.And     -> L.build_and
            | S.Or      -> L.build_or
            | S.Equal   -> L.build_fcmp L.Fcmp.Oeq
            | S.Neq     -> L.build_fcmp L.Fcmp.One
            | S.Less    -> L.build_fcmp L.Fcmp.Olt
            | S.Leq     -> L.build_fcmp L.Fcmp.Ole
            | S.Greater -> L.build_fcmp L.Fcmp.Ogt
            | S.Geq     -> L.build_fcmp L.Fcmp.Oge
            ) e1' e2' "tmp" builder  (* )     *)
        | S.Unop(op, e) ->
    	    let e' = expr builder e in
	        (match op with
	        | S.Neg     -> L.build_neg
            | S.Not     -> L.build_not) e' "tmp" builder
        | S.Assign (s, e) -> let e' = expr builder e in (*ignore(print_endline s);*)
	                   ignore (L.build_store e' (lookup s) builder); e'
        | S.Assign_array (s, i, e) -> let e' = expr builder e in
            let index =L.build_in_bounds_gep (lookup s) [| L.const_int i32_t 0; L.const_int i32_t i |] "str" builder in                    
            ignore (L.build_store e' index builder); e'
        | S.Call ("print_int", [e]) ->
	        L.build_call printf_func [| int_format_str ; (expr builder e) |]
	        "printf" builder
        | S.Call ("print_string", [e]) ->   
            L.build_call printf_func [| str_format_str ; (expr builder e) |]
           "printf" builder
        | S.Call ("print_float", [e]) ->   
            L.build_call printf_func [| flt_format_str ; (expr builder e) |]
           "printf" builder
        | S.Call ("strlen", [a]) ->
            L.build_call strlen_func [|expr builder a|] "strlen" builder
        | S.Call ("create", act) -> (*ignore(print_endline "create");*)
            L.build_call create_func [||] "create" builder
        | S.Call ("array_add_string", [a; b; c]) ->
            L.build_call array_add_string_func [|(expr builder a) ; (expr builder b) ; (expr builder c)|]
            "array_add_string" builder
        | S.Call ("array_retrieve_string", [a; b]) ->
            L.build_call array_retrieve_string_func [|(expr builder a) ; (expr builder b)|]
            "array_retrieve_string" builder
        | S.Call ("array_add_float", [a; b; c]) ->
            L.build_call array_add_float_func [|(expr builder a) ; (expr builder b) ; (expr builder c)|]
            "array_add_float" builder
        | S.Array_Assign_Float (name, key, value) -> 
            L.build_call array_add_float_func [|(L.build_load (lookup name) name builder) ; (expr builder key) ; (expr builder value)|]
            "array_add_float" builder
        | S.Array_Assign_String (name, key, value) -> 
            L.build_call array_add_string_func [|(L.build_load (lookup name) name builder) ; (expr builder key) ; (expr builder value)|]
            "array_add_string" builder
        | S.Call ("array_retrieve_float", [a; b]) ->
            L.build_call array_retrieve_float_func [|(expr builder a) ; (expr builder b)|]
            "array_retrieve_float" builder
        | S.Retrieve_Float (name, key) ->
            L.build_call array_retrieve_float_func [|(L.build_load (lookup name) name builder) ; (expr builder key)|]
            "array_retrieve_float" builder
        | S.Retrieve_String (name, key) ->
            L.build_call array_retrieve_string_func [|(L.build_load (lookup name) name builder) ; (expr builder key)|]
            "array_retrieve_string" builder
        | S.Call (f, act) -> (*ignore(print_endline ("Function:"^f));*)
                let (fdef, fdecl) = StringMap.find f function_decls in
	            let actuals = List.rev (List.map (expr builder) (List.rev act)) in
	            let result = (match fdecl.S.typ with S.Void -> ""
                                            | _ -> f ^ "_result") in
            L.build_call fdef (Array.of_list actuals) result builder  
        | S.Assign_func (s, f) -> ignore(L.build_store (lookup_func f) (lookup s) builder); 
                    L.const_int i32_t 0
        | S.Assign_func_array (s, i, f) -> 
            let index =L.build_in_bounds_gep (lookup s) [| L.const_int i32_t 0; L.const_int i32_t i |] "str" builder in            
            ignore(L.build_store (lookup_func f) index builder); 
                    L.const_int i32_t 0
        in

        (* Invoke "f builder" if the current block doesn't already
        have a terminal (e.g., a branch). *)
        let add_terminal builder f =
            match L.block_terminator (L.insertion_block builder) with
        	| Some _ -> ()
            | None -> ignore (f builder) 
        in
	
        (* Build the code for the given statement; return the builder for
           the statement's successor *)
        let rec stmt builder = function
    	    | S.Block sl -> List.fold_left stmt builder sl
            | S.Expr e -> ignore (expr builder e); builder
            | S.Return e -> ignore (match fdecl.S.typ with
    	        | S.Void -> L.build_ret_void builder
    	        | _ -> L.build_ret (expr builder e) builder); builder
            | S.If (predicate, then_stmt, else_stmt) ->
                let bool_val = expr builder predicate in
    	        let merge_bb = L.append_block context "merge" the_function in
    
    	        let then_bb = L.append_block context "then" the_function in
    	        add_terminal (stmt (L.builder_at_end context then_bb) then_stmt)
    	            (L.build_br merge_bb);
    
    	        let else_bb = L.append_block context "else" the_function in
    	        add_terminal (stmt (L.builder_at_end context else_bb) else_stmt)
    	            (L.build_br merge_bb);
    
    	        ignore (L.build_cond_br bool_val then_bb else_bb builder);
    	        L.builder_at_end context merge_bb
    
            | S.While (predicate, body) ->
    	        let pred_bb = L.append_block context "while" the_function in
    	        ignore (L.build_br pred_bb builder);
    
    	        let body_bb = L.append_block context "while_body" the_function in
    	        add_terminal (stmt (L.builder_at_end context body_bb) body)
    	            (L.build_br pred_bb);

    	        let pred_builder = L.builder_at_end context pred_bb in
    	        let bool_val = expr pred_builder predicate in

    	        let merge_bb = L.append_block context "merge" the_function in
    	        ignore (L.build_cond_br bool_val body_bb merge_bb pred_builder);
    	        L.builder_at_end context merge_bb

            | S.For (e1, e2, e3, body) -> stmt builder
        	    ( S.Block [S.Expr e1 ; S.While (e2, S.Block [body ; S.Expr e3]) ] )  
        in

        (* Build the code for each statement in the function *)
        let builder = stmt builder (S.Block fdecl.S.body) in

        (* Add a return if the last block falls off the end *)
    add_terminal builder (match fdecl.S.typ with
    | S.Void -> L.build_ret_void
    | t -> L.build_ret (L.const_int (ltype_of_typ t) 0)) 
    in

List.iter build_function_body program.Sast.func_decls;
the_module
