open Cparse
open Genlab

let number_str s strings =
    let rec aux n l = match l with
    | [] -> n, [s]
    | s' :: l' when s' = s -> n, l
    | s' :: l' -> let n', l = aux (n + 1) l' in n', s' :: l
    in aux 0 strings;;
    

(* Un dictionnaire est soit vide, soit un couple (clé, valeur) suivi de la suite du dictionnaire *)
type ('a, 'b) dict = NIL | D of 'a * 'b * ('a, 'b) dict;; 

let new_dict = NIL;;

type env = Env of int * string list * string list * bool * bool * int * (string, int) dict * string list;; (* jumpcount, strings, string count, global variables, declaring parameters, in_function, variables dictionary, functions that return a 64bit output *)

(* On recherche si une clé est présente dans le dictionnaire. Renvoie un booléen. *)
let rec contains d key = match d with
    | NIL -> false
    | D(k, _, _) when k = key -> true
    | D(_, _, d') -> contains d' key;;

(* On récupère, si elle existe, la valeur associée à une clé dans un dictionnaire. Si ce n'est pas possible, on renvoie une erreur. Cas d'erreur : utilisation d'une variable locale non déclarée. *)
let rec search d key = match d with
    | NIL -> failwith ("Var " ^ key ^ " not found") 
    | D(k, v, _) when k = key -> v
    | D(_, _, d') -> search d' key;;

(* On ajoute un couple (clé, valeur) dans un dictionnaire, et si la clé existe déjà, on se contente de remplacer la valeur. *)
let rec append d key value = match d with
    | NIL -> D(key, value, NIL)
    | D(k, v, d') when k = key -> D(key, value, d')
    | D(k, v, d') -> D(k, v, (append d' key value));;

(* Tableau des registres d'arguments utiles lors d'un appel de fonction, trié dans l'ordre. *)
let args = [|"%rdi"; "%rsi"; "%rdx"; "%rcx"; "%r8"; "%r9"|];;

(* Tableau des fonctions systèmes disposant d'un retour sur 64 bits. *)
let system_funs_64 = [|"malloc"; "calloc"; "realloc"; "exit"; "fopen"|];;

(* Récupère le nom en assembleur d'une variable. Si la variable est déclarée, alors on renvoie -8n(%rbp) où n est le numéro de la variable, sinon on renvoie s(%rip) où s est le nom de la variable, en supposant que s est une variable globale préalablement déclarée (peut être une variable sytème comme stdout ou stderr *)
let get_name_var vars key =
    if contains vars key then (string_of_int (8 * (search vars key)) ^ "(%rbp)") else key ^ "(%rip)";;

(* Écris le code assembleur d'une fonction. Se référer au README pour plus de détails. *)
let compile out decl_list =
    Printf.fprintf out ".section\t.text\n\t\t.global main\n\n";
	(* Évalue une expression dans un certain environement. Le résultat de l'évaluation est placé dans %rax. On renvoie le nouvel environnement. *)
    let rec evaluate e env = match env with Env(loopcount, strings, globals, decl_params, in_function, var_count, vars, funs) -> match e with
    | VAR(s) -> begin
                    Printf.fprintf out "\t\tMOVQ\t\t%s,\t%%rax\t\t# %s\n" (get_name_var vars s) s;
                    env
                end
    | CST(i) -> begin
                    Printf.fprintf out "\t\tMOVQ\t\t$%d,\t\t%%rax\n" i;
                    env
                end
    | STRING(s) -> begin
                       let n, l = number_str s strings in
                       Printf.fprintf out "\t\tMOVQ\t\t$.str%d,\t\t%%rax\n" n;
                       Env(loopcount, l, globals, decl_params, in_function, var_count, vars, funs)
                   end;
    | SET_VAR(s, e) -> begin
                         let new_env = evaluate (e_of_expr e) env in
                         Printf.fprintf out "\t\tMOVQ\t\t%%rax,\t\t%s\n" (get_name_var vars s);
                         new_env
                       end
    | SET_ARRAY(s, e1, e2) -> begin
                                let new_env = evaluate (e_of_expr e1) env in
                                Printf.fprintf out "\t\tPUSHQ\t\t%%rax\n";
                                let new_new_env = evaluate (e_of_expr e2) new_env in
                                Printf.fprintf out "\t\tPOPQ\t\t%%rbx\n";
                                Printf.fprintf out "\t\tMOVQ\t\t%s,\t%%rdx\n" (get_name_var vars s);
                                Printf.fprintf out "\t\tLEAQ\t\t0(, %%rbx, 8),\t%%rbx\n";
                                Printf.fprintf out "\t\tADDQ\t\t%%rbx,\t\t%%rdx\n";
                                Printf.fprintf out "\t\tMOVQ\t\t%%rax,\t\t(%%rdx)\n";
                                new_new_env
                              end
    | CALL(s, l) -> begin
                      let rec empile_args env i = function
                      | [] -> env
                      | e :: l' -> begin
                                     let new_env = empile_args env (i + 1) l' in
                                     let new_new_env = evaluate (e_of_expr e) new_env in
                                     Printf.fprintf out "\t\tPUSHQ\t\t%%rax\n";
                                     new_new_env
                                   end
                      in let rec depile_args env i = function
                      | [] -> env
                      | e :: l' -> if i >= 6 then env else begin
                                       Printf.fprintf out "\t\tPOPQ\t\t%s\n" args.(i);
                                       depile_args env (i + 1) l'
                                   end
                      in let newenv = empile_args env 0 l in
                      let new_env = depile_args newenv 0 l in
                      Printf.fprintf out "\t\tMOVQ\t\t$0,\t\t%%rax\n";
                      Printf.fprintf out "\t\tCALLQ\t\t%s\n" s;
                      for i = 0 to (List.length l) - 7 do
                          Printf.fprintf out "\t\tPOPQ\t\t%%rbx\n";
                      done;
                      if not (List.mem s funs) then
                          Printf.fprintf out "\t\tCLTQ\n";
                      new_env
                    end
    | OP1(mop, e) -> (match mop with
                     | M_MINUS -> begin
                                      let new_env = evaluate (e_of_expr e) env in
                                      Printf.fprintf out "\t\tNEGQ\t\t%%rax\n";
                                      new_env
                                  end
                     | M_NOT -> begin
                                    let new_env = evaluate (e_of_expr e) env in
                                    Printf.fprintf out "\t\tNOTQ\t\t%%rax\n";
                                    new_env 
                                end
                     | M_POST_INC -> begin
                                       let new_env = evaluate (e_of_expr e) env in
                                       (match (e_of_expr e) with
                                       | VAR(s) -> begin
                                                       Printf.fprintf out "\t\tINCQ\t\t%s\n" (get_name_var vars s);
                                                       new_env
                                                   end
                                       | OP2(S_INDEX, e1, e2) -> (match (e_of_expr e1) with
                                                                 | VAR(s) -> begin
                                                                               Printf.fprintf out "\t\tPUSHQ\t\t%%rax\n";
                                                                               let new_new_env = evaluate (e_of_expr e2) env in
                                                                               Printf.fprintf out "\t\tMOVQ\t\t%s,\t%%rdx\n" (get_name_var vars s);
                                                                               Printf.fprintf out "\t\tLEAQ\t\t0(, %%rax, 8),\t%%rax\n";
                                                                               Printf.fprintf out "\t\tADDQ\t\t%%rax,\t\t%%rdx\n";
                                                                               Printf.fprintf out "\t\tINCQ\t\t(%%rdx)\n";
                                                                               Printf.fprintf out "\t\tPOPQ\t\t%%rax\n";
                                                                               new_new_env
                                                                             end
                                                                 | _ -> new_env)
                                       | _ -> new_env)
                                     end
                     | M_POST_DEC -> begin
                                         let new_env = evaluate (e_of_expr e) env in
                                         match (e_of_expr e) with
                                         | VAR(s) -> begin
                                                         Printf.fprintf out "\t\tDECQ\t\t%s\n" (get_name_var vars s);
                                                         new_env
                                                     end
                                         | OP2(S_INDEX, e1, e2) -> (match (e_of_expr e1) with
                                                                   | VAR(s) -> begin
                                                                                 Printf.fprintf out "\t\tPUSHQ\t\t%%rax\n";
                                                                                 let new_new_env = evaluate (e_of_expr e2) env in
                                                                                 Printf.fprintf out "\t\tMOVQ\t\t%s,\t%%rdx\n" (get_name_var vars s);
                                                                                 Printf.fprintf out "\t\tLEAQ\t\t0(, %%rax, 8),\t%%rax\n";
                                                                                 Printf.fprintf out "\t\tADDQ\t\t%%rax,\t\t%%rdx\n";
                                                                                 Printf.fprintf out "\t\tDECQ\t\t(%%rdx)\n";
                                                                                 Printf.fprintf out "\t\tPOPQ\t\t%%rax\n";
                                                                                 new_new_env
                                                                               end
                                                                   | _ -> new_env)
                                         | _ -> new_env
                                     end
                     | M_PRE_INC -> begin
                                        let new_env = (match (e_of_expr e) with
                                        | VAR(s) -> begin
                                                        Printf.fprintf out "\t\tINCQ\t\t%s\n" (get_name_var vars s);
                                                        env
                                                      end
                                        | OP2(S_INDEX, e1, e2) -> (match (e_of_expr e1) with
                                                                 | VAR(s) -> begin
                                                                                let new_new_env = evaluate (e_of_expr e2) env in
                                                                                Printf.fprintf out "\t\tMOVQ\t\t%s,\t%%rdx\n" (get_name_var vars s);
                                                                                Printf.fprintf out "\t\tLEAQ\t\t0(, %%rax, 8),\t%%rax\n";
                                                                                Printf.fprintf out "\t\tADDQ\t\t%%rax,\t\t%%rdx\n";
                                                                                Printf.fprintf out "\t\tINCQ\t\t(%%rdx)\n";
                                                                                Printf.fprintf out "\t\tMOVQ\t\t(%%rdx),\t\t%%rax\n";
                                                                                new_new_env
                                                                             end
                                                                 | _ -> env);
                                        | _ -> env) in
                                        evaluate (e_of_expr e) new_env
                                      end
                     | M_PRE_DEC -> begin
                                        let new_env = (match (e_of_expr e) with
                                        | VAR(s) -> begin
                                                        Printf.fprintf out "\t\tDECQ\t\t%s\n" (get_name_var vars s);
                                                        env
                                                    end
                                        | OP2(S_INDEX, e1, e2) -> (match (e_of_expr e1) with
                                                                 | VAR(s) -> begin
                                                                                 let new_new_env = evaluate (e_of_expr e2) env in
                                                                                 Printf.fprintf out "\t\tMOVQ\t\t%s,\t%%rdx\n" (get_name_var vars s);
                                                                                 Printf.fprintf out "\t\tLEAQ\t\t0(, %%rax, 8),\t%%rax\n";
                                                                                 Printf.fprintf out "\t\tADDQ\t\t%%rax,\t\t%%rdx\n";
                                                                                 Printf.fprintf out "\t\tDECQ\t\t(%%rdx)\n";
                                                                                 Printf.fprintf out "\t\tMOVQ\t\t(%%rdx),\t\t%%rax\n";
                                                                                 new_new_env
                                                                             end
                                                                 | _ -> env);
                                        | _ -> env) in
                                        evaluate (e_of_expr e) new_env
                                    end)
    | OP2(bop, e1, e2) -> begin
                            let new_env = evaluate (e_of_expr e2) env in
                            Printf.fprintf out "\t\tPUSHQ\t\t%%rax\n";
                            let new_new_env = evaluate (e_of_expr e1) new_env in match new_new_env with Env(loopcount, strings, _, _, _, _, _, _) -> 
                            Printf.fprintf out "\t\tPOPQ\t\t%%rbx\n";
                            (match bop with
                            | S_MUL -> begin
                                           Printf.fprintf out "\t\tIMULQ\t\t%%rbx,\t\t%%rax\n";
                                           new_new_env
                                       end
                            | S_DIV -> begin
                                           Printf.fprintf out "\t\tCQO\n";
                                           Printf.fprintf out "\t\tIDIVQ\t\t%%rbx\n";
                                           Printf.fprintf out "\t\tMOVQ\t\t$0,\t\t%%rdx\n";
                                           Env(loopcount + 2, strings, globals, decl_params, in_function, var_count, vars, funs)
                                       end
                            | S_MOD -> begin
                                           Printf.fprintf out "\t\tCQO\n";
                                           Printf.fprintf out "\t\tIDIVQ\t\t%%rbx\n";
                                           Printf.fprintf out "\t\tMOVQ\t\t%%rdx,\t\t%%rax\n";
                                           Printf.fprintf out "\t\tMOVQ\t\t$0,\t\t%%rdx\n";
                                           Env(loopcount + 2, strings, globals, decl_params, in_function, var_count, vars, funs)
                                       end
                            | S_ADD -> begin
                                           Printf.fprintf out "\t\tADDQ\t\t%%rbx,\t\t%%rax\n";
                                           new_new_env
                                       end
                            | S_SUB -> begin
                                           Printf.fprintf out "\t\tSUBQ\t\t%%rbx,\t\t%%rax\n";
                                           new_new_env
                                       end
                            | S_INDEX -> begin
                                             Printf.fprintf out "\t\tLEAQ\t\t0(,%%rbx,8),\t%%rbx\n";
                                             Printf.fprintf out "\t\tADDQ\t\t%%rbx,\t\t%%rax\n";
                                             Printf.fprintf out "\t\tMOVQ\t\t(%%rax),\t\t%%rax\n";
                                             new_new_env
                                         end)
                        end
    | CMP(cop, e1, e2) -> begin
                            let new_env = evaluate (e_of_expr e1) env in
                            Printf.fprintf out "\t\tPUSHQ\t\t%%rax\n";
                            let new_new_env = evaluate (e_of_expr e2) new_env in match new_new_env with Env(loopcount, strings, _, _, _, _, _, _) ->
                            Printf.fprintf out "\t\tPOPQ\t\t%%rbx\n";
                            Printf.fprintf out "\t\tCMPQ\t\t%%rax,\t\t%%rbx\n";
                            (match cop with
                             | C_LT -> Printf.fprintf out "\t\tJL\t\t";
                             | C_LE -> Printf.fprintf out "\t\tJLE\t\t";
                             | C_EQ -> Printf.fprintf out "\t\tJE\t\t";);
                            Printf.fprintf out ".destjump%d\n" (loopcount + 1);
                            Printf.fprintf out "\t\tMOVQ\t\t$0,\t\t%%rax\n";
                            Printf.fprintf out "\t\tJMP\t\t.destjump%d\n" (loopcount + 2);
                            Printf.fprintf out "\t.destjump%d:\n" (loopcount + 1);
                            Printf.fprintf out "\t\tMOVQ\t\t$1,\t\t%%rax\n";
                            Printf.fprintf out "\t.destjump%d:\n" (loopcount + 2);
                            Env(loopcount + 2, strings, globals, decl_params, in_function, var_count, vars, funs)
                          end
    | EIF(e1, e2, e3) -> begin
                           let new_env = evaluate (e_of_expr e1) env in match new_env with Env(loopcount, strings, _, _, _, _, _, _) ->
                           Printf.fprintf out "\t\tCMPQ\t\t$0,\t\t%%rax\n";
                           let x = (loopcount + 1) in
                           Printf.fprintf out "\t\tJE\t\t.destjump%d\n" x;
                           let new_new_env = evaluate (e_of_expr e2) (Env(loopcount + 1, strings, globals, decl_params, in_function, var_count, vars, funs)) in match new_new_env with Env(loopcount2, strings2, _, _, _, _, _, _) ->
                           let y = (loopcount2 + 1) in
                           Printf.fprintf out "\t\tJMP\t\t.destjump%d\n" y;
                           Printf.fprintf out "\t.destjump%d:\n" x;
                           let new_new_env = evaluate (e_of_expr e3) (Env(loopcount2 + 1, strings2, globals, decl_params, in_function, var_count, vars, funs)) in
                           Printf.fprintf out "\t.destjump%d:\n" y;
                           new_new_env
                         end
    | ESEQ(l) -> begin
                     let rec aux env l = match l with
                     | [] -> env
                     | e :: l' -> aux (evaluate (e_of_expr e) env) l'
                     in aux env l
                 end
	(* Déclare une liste de variables ou de fonctions, en mettant bien à jour l'environnement. *)
    in let rec compile_decl_list env = function 
        | [] -> env
        | h :: t -> let new_env = compile_decl env h in compile_decl_list new_env t
	(* Compte la place nécessaire pour les variables au sein d'un bloc de code. *)
    and count_vars (_, c) = match c with
        | CBLOCK(vdl, lcl) -> let vars = (List.length vdl) in
                              let rec aux lcl = match lcl with
                              | [] -> 0, 0
                              | c :: l -> let v, p = count_vars c in let v2, p2 = aux l in (v + v2), (max p p2)
                              in let v, p = aux lcl in (v + vars), p
        | CIF(_, c1, c2) -> let i1, j1 = (count_vars c1) in let i2, j2 = (count_vars c2) in (max i1 i2), (max j1 j2)
        | CEXPR(_) -> 0, 0
        | CWHILE(_, c) -> count_vars c
        | CRETURN(_) -> 0, 0
	(* Déclare une variable locale en distingant si c'est un paramètre ou non, ou une fonction. *)
    and compile_decl env decl = match env with Env(loopcount, strings, globals, decl_params, in_function, var_count, vars, funs) -> match decl with 
        | CDECL(l, s) -> begin
                             if decl_params then begin
                                 if var_count < 6 then begin
                                     Printf.fprintf out "\t\tMOVQ\t\t%s,\t\t%d(%%rbp)\t# %s\n" args.(var_count) (-8 * (var_count + 1)) s;
                                     Env(loopcount, strings, globals, decl_params, in_function, var_count + 1, (append vars s (-(var_count + 1))), funs)
                                 end
                                 else begin
                                     Printf.fprintf out "\t\tMOVQ\t\t%d(%%rbp),\t%%rax\n" (8 * (var_count - 4));
                                     Printf.fprintf out "\t\tMOVQ\t\t%%rax,\t\t%d(%%rbp)\t# %s\n" (-8 * (var_count + 1)) s;
                                     Env(loopcount, strings, globals, decl_params, in_function, var_count + 1, (append vars s (-(var_count + 1))), funs)
                                 end   
                             end
                             else begin
                                 if in_function then begin
                                     Env(loopcount, strings, globals, decl_params, in_function, var_count + 1, (append vars s (-(var_count + 1))), funs)
                                 end
                                 else begin
                                     Env(loopcount, strings, (s :: globals), decl_params, in_function, var_count, vars, funs)
                                 end
                             end
                         end
        | CFUN (l, s, vdl, lc) -> begin
                                      Printf.fprintf out "\n%s:\n" s;
                                      let nb_decl = List.length vdl in
                                      let total_vars, max_params = count_vars lc in
                                      let size = total_vars + nb_decl + max_params in
                                      let real_size = size + (size mod 2) in
                                      Printf.fprintf out "\t\tENTERQ\t\t$%d,\t\t$0\n" (8 * real_size);
                                      let new_env = compile_decl_list (Env(loopcount, strings, globals, true, true, 0, new_dict, funs)) vdl in match new_env with Env(_, _, _, _, _, var_count, vars, _) ->
                                      let new_new_env = compile_code lc (Env(loopcount, strings, globals, false, true, var_count, vars, funs)) in match new_new_env with Env(loopcount2, strings2, globals2, _,  _, var_count2, vars2, _) ->
                                      Printf.fprintf out "\t\tLEAVEQ\n";
                                      Printf.fprintf out "\t\tRETQ\n";
                                      Env(loopcount2, strings2, globals2, false, false, var_count2, vars2, funs)
                                  end
	(* Écris le code assembleur d'un bout de code. *)
    and compile_code (_, c) env = match env with Env(loopcount, strings, globals, decl_params, in_function, var_count, vars, funs) -> match c with
        | CBLOCK(vdl, lcl) -> begin
                                  let new_env = compile_decl_list env vdl in
                                  let rec aux lcl env = match lcl with
                                  | [] -> env
                                  | c :: l -> let new_env = (compile_code c env) in aux l new_env
                                  in let new_env = aux lcl new_env in match new_env with Env(loopcount2, strings2, globals2, _, _, _, _, _) ->
                                  Env(loopcount2, strings2, globals2, decl_params, in_function, var_count, vars, funs)
                              end
        | CEXPR(expr) -> begin
                             evaluate (e_of_expr expr) env
                         end
        | CIF(expr, c1, c2) -> begin
                                 let new_env = evaluate (e_of_expr expr) env in
                                 match new_env with Env(loopcount2, strings2, _, _, _, _, _, _) ->
                                 let x = (loopcount2 + 1) in
                                 Printf.fprintf out "\t\tCMPQ\t\t$0,\t\t%%rax\n";
                                 Printf.fprintf out "\t\tJE\t\t.destjump%d\n" x;
                                 let new_new_env = compile_code c1 (Env(loopcount2 + 1, strings2, globals, decl_params, in_function, var_count, vars, funs)) in
                                 match new_new_env with Env(loopcount3, strings3, _, _, _, _, _, _) ->
                                 let y = (loopcount3 + 1) in
                                 Printf.fprintf out "\t\tJMP\t\t.destjump%d\n" y;
                                 Printf.fprintf out "\t.destjump%d:\n" x;
                                 let new_new_new_env = compile_code c2 (Env(loopcount3 + 1, strings3, globals, decl_params, in_function, var_count, vars, funs)) in
                                 Printf.fprintf out "\t.destjump%d:\n" y;
                                 new_new_new_env
                               end
        | CWHILE(expr, c) -> begin
                              let x = (loopcount + 1) in
                              Printf.fprintf out "\t.whileloop%d:\n" x;
                              let new_env = evaluate (e_of_expr expr) (Env(loopcount + 1, strings, globals, decl_params, in_function, var_count, vars, funs)) in
                              Printf.fprintf out "\t\tCMPQ\t\t$0,\t\t%%rax\n";
                              Printf.fprintf out "\t\tJE\t\t.endloop%d\n" x;
                              let new_new_env = compile_code c new_env in
                              Printf.fprintf out "\t\tJMP\t\t.whileloop%d\n" x;
                              Printf.fprintf out "\t.endloop%d:\n" x;
                              new_new_env
                          end
        | CRETURN(o) -> begin
                          let new_env = match o with
                          | Some(e) -> evaluate (e_of_expr e) env
                          | None -> env
                          in Printf.fprintf out "\t\tLEAVEQ\n";
                          Printf.fprintf out "\t\tRETQ\n";
                          new_env
                        end
	(* Récupère la liste des fonctions déclarées, afin de gérer quelles fonctions renvoient un entier sur 64 bits ou 32 bits. *)
    in let rec get_functions = function
    | [] -> []
    | CFUN(_, s, _, _) :: l -> s :: get_functions l
    | _ :: l -> get_functions l
	(* Compile le code. *)
    in let final_env = compile_decl_list (Env(0, [], [], false, false, 0, new_dict, (Array.to_list system_funs_64) @ (get_functions decl_list))) decl_list in match final_env with Env(_, strings, globals, _, _, _, _, _) ->
	(* Déclare les chaînes de caractères présentes. *)
    Printf.fprintf out "\n.section\t.data\n";
    let rec add_string_globals n = function
    | [] -> ();
    | s :: l' -> begin
                     Printf.fprintf out ".str%d:\n" n;
                     Printf.fprintf out "\t\t.string\t\t\"%s\"\n" (String.escaped s);
                     Printf.fprintf out "\t\t.text\n";
                     add_string_globals (n + 1) l';
                 end
    in add_string_globals 0 strings;
	(* Déclare les variables globales présentes. *)
    let rec add_int_globals = function
    | [] -> ()
    | s :: l -> begin
                    add_int_globals l;
                    Printf.fprintf out ".comm\t\t%s,\t\t8,\t\t8\n" s;
                end
    in add_int_globals globals;;