open Format
open Ast
open Codevm

(* Déclaration de l'exception à lever quand une variable (locale ou globale) 
   est mal utilisée *)
exception VarUndef of string

(* Fonction de pretty printing utilisée pour écrire dans le fichier de sortie
   le code machine généré à la compilation. *)
let emit_code fmt = function
    PUSHI n -> fprintf fmt "\tpushi %d\n" n
  | STOREL n -> fprintf fmt "\tstorel %d\n" n
  | PUSHL n -> fprintf fmt "\tpushl %d\n" n
  | STOREG n -> fprintf fmt "\tstoreg %d\n" n
  | PUSHG n -> fprintf fmt "\tpushg %d\n" n
  | PUSHN n -> fprintf fmt "\tpushn %d\n" n
  | START -> fprintf fmt "\tstart\n"
  | ADD -> fprintf fmt "\tadd\n"
  | SUB -> fprintf fmt "\tsub\n"
  | MUL -> fprintf fmt "\tmul\n"
  | DIV -> fprintf fmt "\tdiv\n"
  | WRITEI -> fprintf fmt "\twritei\n"
  | STOP -> fprintf fmt "\tstop\n"


(* gp_max (resp. fp_max) contient le nombre maximal de variables
   globales (resp. locales) utilisées par le programme. Ces deux
   références sont mises à jour pendant la compilation et utilisées
   pour générer le code d'initialisation de la machine virtuelle
   (cf. fonction preambule). *)

let gp = ref 0
let fp_max = ref 0

(* Cette fonction génère le code machine pour initialiser les
   registres de la machine virtuelle avant chaque programme
   (réservation de l'espace nécessaire pour les variables globales et
   locales). *)
let preambule gp fp = [(PUSHN gp) ; START ; (PUSHN fp)]

(* La table des symboles du compilateur est implémentée à l'aide 
   d'une table de hachage. Chaque entrée lie une variable globale à son
   indice dans la pile par rapport au registre gp. *)
let ( genv : (string,int) Hashtbl.t ) = Hashtbl.create 17 

(* On utilise une table d'association dont les clés (les variables locales) 
   sont des chaînes de caractères *)
module StrMap = Map.Make(String)

(* Compilation d'une expression *)	
let compile_expr = 

  (* Fonction récursive locale à compile_expr utilisée pour générer le code 
     machine de l'arbre de syntaxe abstraite associé à une valeur de type 
     Ast.expr *)
  let rec comprec env fp = function
      Cst i -> [PUSHI i]
    | Var x -> 
	begin
	  try [PUSHL(StrMap.find x env)]
	  with Not_found ->
	    try [PUSHG(Hashtbl.find genv x)]
	    with Not_found -> raise(VarUndef x)
	end
    | Op(o,e1,e2)-> 
	let op = 
	  match o with
	    Sum -> [ADD]
	  | Diff -> [SUB]
	  | Prod -> [MUL]
	  | Quot -> [DIV]
	in
	let code_e1 = comprec env fp e1 in
	let code_e2 = comprec env fp e2 in
	  code_e1@code_e2@op
    | Letin(x,e1,e2) -> 
	begin
	  let code_e1 = comprec env fp e1 in
	  if !fp_max=fp then incr fp_max;
	  let code_e2 = comprec (StrMap.add x fp env) (fp+1) e2 in
	    code_e1@(STOREL(fp)::code_e2)
	end
  in 
    comprec StrMap.empty 0

(* Compilation d'une instruction *)            
let compile_instr = function
    Set(x,e) -> 
      let code = (compile_expr e)@[STOREG(!gp)] in
      begin
	Hashtbl.add genv x !gp;
	incr gp;
	code
      end
  | Print(e) -> (compile_expr e)@[WRITEI]


(* Enregistre dans le fichier ofile le code machine correspondant à la 
   compilation du programme p *)
let compile_prg p ofile =
  let code = List.map compile_instr p in

  (* On ajoute le préambule d'initialisation au code du programme p *)
  let code_final = (preambule !gp !fp_max)::code in
  let f = open_out ofile in 

  (* On utilise les fonctions de la librairie Format pour écrire le code dans 
     le fichier de sortie *)
  let fmt = formatter_of_out_channel f in
  List.iter (List.iter (emit_code fmt)) code_final;

  (* Il ne faut pas oublier de générer le code pour arrêter
     correctement la machine virtuelle à la fin d'un programme. *)
  emit_code fmt STOP;

  (* On "flush" le buffer afin de s'assurer que tout y a été écrit
  avant de le fermer *)
  fprintf fmt "@?"; 
  close_out f