open Format
open Mips
open Ast
open AstO

(* Labels pour les tests *)
module Label = struct
  let i = ref 1
  let fresh () = let j = !i in incr i; "label_"^(string_of_int j)
end

(* Taille de la frame, en octets (chaque variable locale occupe sint octets) *)
let frame_size = ref 0

(* Les variables globales sont stockées dans une table de hachage *)
let (genv : (string, unit) Hashtbl.t) = Hashtbl.create 17 

(* On utilise une table d'association dont les clés sont les variables locales
   (des chaînes de caractères) et où la valeur associée est la position
   par rapport à $fp (en octets) *)
module Smap = Map.Make(String)

let addSP n = Arith (Mips.Add,SP,SP,Oimm n)
let push = Arith (Mips.Sub, SP, SP, Oimm sint)
let pop = Arith (Mips.Add, SP, SP, Oimm sint)


let arith_to_mips = function
  | Add -> Mips.Add
  | Sub -> Mips.Sub
  | Mul -> Mips.Mul
  | Div -> Mips.Div
  | Mod -> Mips.Rem
  | _ -> failwith "ERREUR INTERNE 12 : fichier compile.ml - fonction arith_to_mips"

let comp_to_mips = function
  | Lt -> Mips.Lt
  | Le -> Mips.Le
  | Gt -> Mips.Gt
  | Ge -> Mips.Ge
  | Eq -> Mips.Eq
  | Neq -> Mips.Neq


let cst_to_real = function
  | Int i -> i
  | Bool true -> 1
  | Bool false -> 0
  | Unit -> 0
  | EmptyList -> 0

(** Copie les nb*sintoctets au sommet de la pile dans les nb*sintoctets en partant de
 * FP - utile pour comprec_argsT - ATTENTION utilise V0 *)
let comprec_move nb = 
  let rec aux fp = function
    | 0 -> []
    | nb -> 
        [Lw (V0, Areg(nb*sint-sint,SP));
         Sw (V0, Areg(-fp,FP))] @
        (aux (fp+sint) (nb-1))
  in aux 0 nb

(* Fonction récursive utilisée pour générer le code 
   machine de l'arbre de syntaxe abstraite associé à une valeur de type 
   AstO.expr ; à l'issue de l'exécution de ce code, la valeur doit se trouver
   dans le registre V0 *)
let rec comprec env fp = function
  | Cst Unit ->
      (* on met une valeur quelconque !!! *)
      []
  | Cst i -> let i = cst_to_real i in 
      [Li (V0, i)]
  | Var x -> 
      begin
        try
          let fp = Smap.find x env in
          [Lw (V0, Areg (-fp, FP))]
        with Not_found ->
          (*if not (Hashtbl.mem genv x) then raise (VarUndef x);*)
          [Lw (V0, Alab ("var_"^x))]
      end
  | App (x,args) ->
      let code_args = comprec_args env fp 0 args in
      code_args@[Jal ("fun_"^x)]
  | AppT (x,args) ->
      let code_args = comprec_argsT env fp args in
      code_args@[J ("funT_"^x)]
  | Ref e ->
      let code_e = comprec env fp e in
      code_e @
      [Move (V1,V0);
       Li (V0,9);
       Li (A0,sint);
       Syscall;
       Sw (V1, Areg (0,V0))]
  | Binop (Comp o,e1,e2) ->
      let op = comp_to_mips o in
      let code_e1 = comprec env fp e1 in
      let code_e2 = comprec env fp e2 in
      code_e1 @ 
      [push;
       Sw (V0, Areg (0, SP))] @
      code_e2 @
      [Lw (V1, Areg (0,SP));
       pop;
       SComp (op, V0, V1, V0)]
  | Binop (Semi,e1,e2) ->
      let code_e1 = comprec env fp e1 in
      let code_e2 = comprec env fp e2 in
      code_e1 @
      code_e2
  | Binop (Concat,e1,e2) ->
      let code_e1 = comprec env fp e1 in
      let code_e2 = comprec env fp e2 in
      code_e1 @
      [Move (V1,V0);
       Li (V0,9);
       Li (A0,2*sint);
       Syscall;
       Sw (V1, Areg (0,V0));
       push;
       Sw (V0, Areg (0,SP))] @
      code_e2 @
      [Move (V1,V0);
       Lw (V0, Areg (0,SP));
       pop;
       Sw (V1, Areg (sint,V0))]
  | Binop (Assert,e1,e2) ->
      let code_e1 = comprec env fp e1 in
      let code_e2 = comprec env fp e2 in
      code_e1 @ 
      [push;
       Sw (V0, Areg (0, SP))] @
      code_e2 @
      [Lw (V1, Areg (0,SP));
       pop;
       Sw (V0, Areg (0,V1))]

  (* Pendant l'optimisation, quand les opérations commutent,
   * on a mis les constantes à droite pour cela ! *)
  | Binop (o,e1,Cst (Int i)) -> 
      let op = arith_to_mips o in
      let code_e1 = comprec env fp e1 in
      code_e1 @ 
      [Arith (op, V0, V0, Oimm i)]
  | Binop (o,e1,e2)-> 
      let op = arith_to_mips o in
      let code_e1 = comprec env fp e1 in
      let code_e2 = comprec env fp e2 in
      code_e1 @ 
      [push;
       Sw (V0, Areg (0, SP))] @
      code_e2 @
      [Lw (V1, Areg (0,SP));
       pop;
       Arith (op, V0, V1, Oreg V0)]
  | Unop (Addi i,e) ->
      let code_e = comprec env fp e in
      code_e @
      [Arith (Mips.Add, V0, V0, Oimm i)]

  | Unop (Minus,e) ->
      let code_e = comprec env fp e in
      code_e @
      [Neg (V0, V0)]
  | Unop (Excl,e) ->
      let code_e = comprec env fp e in
      code_e @
      [Lw (V0, Areg (0, V0))]

  | IfThenElse ((Binop ((Comp comp),e1c,e2c)), e1, e2) ->
      let op = comp_to_mips comp
      and code_e1c = comprec env fp e1c
      and code_e2c = comprec env fp e2c
      and code_e1 = comprec env fp e1
      and label_e2d = Label.fresh ()
      and code_e2 = comprec env fp e2 
      and label_e2f = Label.fresh () in
      code_e1c @ 
      [push;
       Sw (V0, Areg (0, SP))] @
      code_e2c @
      [Lw (V1, Areg (0,SP));
       pop;
       BComp (op, V1, V0, label_e2d)] @
      code_e2 @
      [B label_e2f;
       Label label_e2d] @
      code_e1 @
      [Label label_e2f]
  | IfThenElse (e,e1,e2) ->
      let code_e = comprec env fp e
      and code_e1 = comprec env fp e1
      and label_e2d = Label.fresh ()
      and code_e2 = comprec env fp e2 
      and label_e2f = Label.fresh () in
      code_e @
      [B0Comp (Mips.Eq,V0,label_e2d)] @
      code_e1 @
      [B label_e2f;
       Label label_e2d] @
      code_e2 @
      [Label label_e2f]
  | LetIn (x,e1,e2) -> 
      if !frame_size = fp then frame_size := sint + !frame_size;
      let code_e1 = comprec env fp e1 in
      let code_e2 = comprec (Smap.add x fp env) (fp+sint) e2 in
      code_e1 @ 
      [Sw (V0, Areg (-fp, FP))] @ 
      code_e2
  | MatchWith (e,e1,x1,x2,e2) ->
      if !frame_size = fp then frame_size := sint + !frame_size;
      if !frame_size = fp+sint then frame_size := sint + !frame_size;
      let code_e = comprec env fp e
      and code_e1 = comprec env fp e1
      and label_e2d = Label.fresh ()
      and code_e2 = comprec (Smap.add x1 fp (Smap.add x2 (fp+sint) env))
      (fp+2*sint) e2
      and label_e2f = Label.fresh () in
      code_e @
      [B0Comp (Mips.Eq,V0,label_e2d);
       Lw (V1, Areg (0,V0)); (* V1 est x1 *)
       Sw (V1, Areg (-fp,FP));
       Lw (V1, Areg (sint,V0)); (* V1 est x2 *)
       Sw (V1, Areg (-fp-sint,FP))] @
      code_e2 @
      [B label_e2f;
       Label label_e2d] @
      code_e1 @
      [Label label_e2f]

(** Prépare les arguments (sur la pile) pour l'appel d'une fonction
 * i est le numéro de l'argument en cours de traitement *)
and comprec_args env fp i = function
  | [] -> []
  | e::s ->
      let code_e = comprec env fp e in
      code_e @ (* beaucoup d'améliorations possibles ici : on push  à chaque
      fois ! *)
      [push;
       Sw (V0, Areg (0,SP))] @
      (comprec_args env fp (i+1) s)

(** Prépare les arguments pour un appel terminal de fonction *)
and comprec_argsT env fp el = 
  (* Calcule les arguments et les met sur la pile *)
  let code_args = comprec_args env fp 0 el in
  (* Puis remplace les anciens arguments de la fonction par ceux qui viennent
   * d'être calculés *)
  let code_move = comprec_move (List.length el) in
  code_args @
  code_move @
  [addSP (sint*(List.length el))] (* on supprime les arguments calculés *)

(* Compilation d'une expression *)        
let compile_expr e =
  comprec Smap.empty 0 e

(* Compilation de toutes les fonctions *)
let compile_fun = function
  | LetFun (x,args,e) ->
      let (nb,env) = List.fold_left 
          (fun (nb,env) x -> (nb+1,Smap.add x (sint*nb) env)) (0,Smap.empty) args
      in
      let old_frame_size = !frame_size in
      frame_size := sint*(nb+2);
      let code = comprec env !frame_size e in
      let new_frame_size = !frame_size in
      frame_size := old_frame_size;
      [Label ("fun_"^x);
       Arith (Mips.Sub, SP, SP, Oimm (new_frame_size-sint*nb));
       Sw (RA, Areg (-(2*sint)+new_frame_size-sint*nb, SP));
       Sw (FP, Areg (-sint+new_frame_size-sint*nb, SP));
       Arith (Mips.Add, FP, SP, Oimm (new_frame_size-sint));
       Label ("funT_"^x); (* pour les appels terminaux *)
      ] @
      code @
      [Lw (RA, Areg (-(2*sint)+new_frame_size-sint*nb, SP));
       Lw (FP, Areg (-sint+new_frame_size-sint*nb, SP));
       Arith (Mips.Add, SP, SP, Oimm new_frame_size);
       Jr RA]
  | _ ->
      []

(* Compilation de toutes les instructions non fonction *)            
let compile_instr = function
  | LetVar (x,e) ->
      let x = "var_"^x in 
      let code = 
        compile_expr e @ 
        [Sw (V0, Alab x)] 
      in
      Hashtbl.replace genv x ();
      code
  | LetIgnore (e) ->
      let code =
        compile_expr e @
        [pop]
      in
      code
  | _ ->
      []

(* Compile le programme p et enregistre le code dans le fichier ofile *)
let compile_program p ofile =
  let code_fun = List.map compile_fun p in
  let code_fun = List.concat code_fun in
  let code_instr = List.map compile_instr p in
  let code_instr = List.concat code_instr in
  (* On ajoute la sauvegarde de $ra et le code de print *)
  (* On ajoute la sauvegarde de $ra et le code de print *)
  let p = 
    { text = 
        [Label "main";
         Move (S0, RA); (* sauvegarde de $ra *)
         Arith (Mips.Sub, SP, SP, Oimm !frame_size); (* alloue la frame *)
         Arith (Mips.Add, FP, SP, Oimm (!frame_size - sint)); (* $fp = ... *)
        ] @ 
        code_instr @
        [Arith (Mips.Add, SP, SP, Oimm !frame_size); (* désalloue la frame *)
         Move (RA, S0); (* restaure $ra *)
         Jr RA;
         (* print_int *)
         Label "fun_print_int_0";
         Lw (A0, Areg (0,SP));
         Li (V0, 1);
         Syscall;
         pop;
         Jr RA;
         (* print_newline *)
         Label "fun_print_newline_0";
         Li (V0, 4);
         La (A0, "newline");
         Syscall;
         Jr RA;
         (* read_int *)
         Label "fun_read_int_0";
         Li (V0, 5);
         Syscall;
         Jr RA;
         (* not *)
         Label "fun_not_0";
         Lw (V0, Areg (0,SP));
         SComp (Mips.Eq,V0,V0,Zero);
         pop;
         Jr RA
        ]@
        code_fun;
      data = 
        Hashtbl.fold (fun x _ l -> Word (x, 1) :: l) genv 
          [Asciiz ("newline", "\n")]
    }
  in
  let f = open_out ofile in 
  let fmt = formatter_of_out_channel f in
  Mips.print_program fmt p;
  (* on "flush" le buffer afin de s'assurer que tout y a été écrit
     avant de le fermer *)
  fprintf fmt "@?"; 
  close_out f
