open Ast
open AstO

module Sset = Set.Make(String)

let delete_prime s = 
  let n = String.length s in
  for i=0 to (n-1) do
    if s.[i]='\'' then (print_int i;s.[i]<-'_')
  done

let convert_unop = function
  | Ast.Minus -> Minus
  | Ast.Excl -> Excl

let convert_binop = function
  | Ast.Add -> Add
  | Ast.Sub -> Sub
  | Ast.Mul -> Mul
  | Ast.Div -> Div
  | Ast.Mod -> Mod
  | Ast.Comp comp -> Comp comp
  | Ast.And -> failwith "ERREUR INTERNE 10 - fichier opt.ml - convert_binop"
  | Ast.Or -> failwith "ERREUR INTERNE 11 - fichier opt.ml -convert_binop"
  | Ast.Concat -> Concat
  | Ast.Assert -> Assert
  | Ast.Semi -> Semi

(** var_unit est l'ensemble Sset des variable à transformer en unit () *)
let rec delete_unit var_unit = function
  | AstD.Cst cst ->
      AstD.Cst cst
  | AstD.Var p when Sset.mem p var_unit ->
      AstD.Cst Unit
  | AstD.Var p ->
      AstD.Var p
  | AstD.App (p,el) ->
      let el = List.map (fun (b,e)->(b,delete_unit var_unit e)) el in
      AstD.App (p,el)
  | AstD.Ref e ->
      let e = delete_unit var_unit e in
      AstD.Ref e
  | AstD.Binop (op,e1,e2) ->
      let e1 = delete_unit var_unit e1 in
      let e2 = delete_unit var_unit e2 in
      AstD.Binop (op,e1,e2)
  | AstD.Unop (op,e) ->
      let e = delete_unit var_unit e in
      AstD.Unop (op,e)
  | AstD.IfThenElse (e,e1,e2) ->
      let e = delete_unit var_unit e in
      let e1 = delete_unit var_unit e1 in
      let e2 = delete_unit var_unit e2 in
      AstD.IfThenElse (e,e1,e2)
  | AstD.LetIn (p,e1,e2) ->
      let e1 = delete_unit var_unit e1 in
      let e2 = delete_unit var_unit e2 in
      AstD.LetIn (p,e1,e2)
  | AstD.MatchWith (e,e1,p1,p2,e2) ->
      let e = delete_unit var_unit e in
      let e1 = delete_unit var_unit e1 in
      let e2 = delete_unit var_unit e2 in
      AstD.MatchWith (e,e1,p1,p2,e2) 

let is16op n = -32768 <= n && n <= 32767

(** Supprime les arguments de type unit et les met avant avec un ; 
 * cstr est le constructeur à utiliser pour l'application*)
let rec opt_args cstr p el =
  (* l2 contient les arguments non unit *)
  let rec aux l2 = function
    | [] -> cstr p l2
    | (true,AstD.Cst Unit)::s ->
        aux l2 s
    | (true,e)::s ->
        let e = opt_expr "" e in
        Binop (Semi,e,aux l2 s)
    | (false,e)::s ->
        let e = opt_expr "" e in
        aux (l2@[e]) s
  in aux [] el

and mk_add e1 e2 = match e1,e2 with
  | Cst (Int n1), Cst (Int n2) ->
      Cst (Int (n1 + n2))
  | e, Cst (Int 0) | Cst (Int 0), e ->
      e
  | Unop (Addi n1, e), Cst (Int n2)
  | Cst (Int n2), Unop (Addi n1, e) ->
      mk_add (Cst (Int (n1+n2))) e
  | e, Cst (Int n) | Cst (Int n), e when is16op n ->
      Unop (Addi n,e)
  | _ ->
      Binop (Add, e1, e2)

and mk_sub e1 e2 = match e1,e2 with
  | Cst (Int n1), Cst (Int n2) ->
     Cst (Int (n1 - n2))
  | e, Cst (Int 0) ->
     e
  | e, Cst (Int n) ->
     mk_add e (Cst (Int (-n)))
  | _ ->
     Binop (Sub, e1, e2)

and mk_mul e1 e2 = match e1,e2 with
  | Cst (Int n1), Cst (Int n2) ->
     Cst (Int (n1 * n2))
  | e, Cst (Int 1) | Cst (Int 1), e ->
     e
  | Cst (Int n), e ->
     Binop (Mul, e, Cst (Int n))
     (* les constantes sont à droites pour une compilation optimale *)
  | _ ->
     Binop (Mul, e1, e2)

and mk_div e1 e2 = match e1,e2 with
  | Cst (Int n1), Cst (Int n2) ->
     Cst (Int (n1/n2))
  | e, Cst (Int 1) ->
     e
  | _ ->
     Binop (Div, e1, e2)

(** Optimise une expression *)
and opt_expr func (e:AstD.expr) = match e with
  | AstD.Cst cst ->
      Cst cst
  | AstD.Var p ->
      delete_prime p;
      Var p
  | AstD.App (p,el) when p=func ->
      delete_prime p;
      opt_args (fun x y -> AppT (x,y)) p el
  | AstD.App (p,el) ->
      delete_prime p;
      opt_args (fun x y -> App (x,y)) p el
  | AstD.Ref e ->
      let e = opt_expr "" e in
      Ref e
  | AstD.Binop (Ast.Semi,e1,e2) ->
      let e1 = opt_expr "" e1 in
      let e2 = opt_expr func e2 in
      Binop (Semi,e1,e2)
  | AstD.Binop (Ast.Or,e1,e2) ->
      opt_expr func
      (AstD.IfThenElse 
        (e1,
         AstD.Cst (Ast.Bool true),
         (AstD.IfThenElse (e2,AstD.Cst (Ast.Bool true),AstD.Cst (Ast.Bool false)))
        )
      )
  | AstD.Binop (Ast.And,e1,e2) ->
      opt_expr func 
      (AstD.IfThenElse 
        (e1,
         (AstD.IfThenElse (e2,AstD.Cst (Ast.Bool true),AstD.Cst (Ast.Bool false))),
         AstD.Cst (Ast.Bool false)
        )
      )
  | AstD.Binop (Ast.Add,e1,e2) ->
      let e1 = opt_expr "" e1 in
      let e2 = opt_expr "" e2 in
      mk_add e1 e2
  | AstD.Binop (Ast.Sub,e1,e2) ->
      let e1 = opt_expr "" e1 in
      let e2 = opt_expr "" e2 in
      mk_sub e1 e2
  | AstD.Binop (Ast.Mul,e1,e2) ->
      let e1 = opt_expr "" e1 in
      let e2 = opt_expr "" e2 in
      mk_mul e1 e2
  | AstD.Binop (Ast.Div,e1,e2) ->
      let e1 = opt_expr "" e1 in
      let e2 = opt_expr "" e2 in
      mk_div e1 e2
  | AstD.Binop (Ast.Mod,e1,e2) ->
      let e1 = opt_expr "" e1 in
      let e2 = opt_expr "" e2 in
      Binop (Mod,e1,e2)
  | AstD.Binop (op,e1,e2) ->
      let e1 = opt_expr "" e1 in
      let e2 = opt_expr "" e2 in
      Binop (convert_binop op,e1,e2)
  | AstD.Unop (op,e) ->
      let e = opt_expr "" e in
      Unop (convert_unop op,e)
  | AstD.IfThenElse (e,e1,e2) ->
      let e = opt_expr "" e in
      let e1 = opt_expr func e1 in
      let e2 = opt_expr func e2 in
      (match e with
        | Cst (Bool true) -> e1
        | Cst (Bool false) -> e2
        | _ -> IfThenElse (e,e1,e2)
      )
  | AstD.LetIn (p,e1,e2) ->
      delete_prime p;
      let e1 = opt_expr "" e1 in
      let e2 = opt_expr func e2 in
      LetIn (p,e1,e2)
  | AstD.MatchWith (e,e1,p1,p2,e2) ->
      delete_prime p1;
      delete_prime p2;
      let e = opt_expr "" e in
      let e1 = opt_expr func e1 in
      let e2 = opt_expr func e2 in
      MatchWith (e,e1,p1,p2,e2)

let rec opt_elt = function
  | AstD.LetIgnore e ->
      let e = opt_expr "" e in
      LetIgnore e
  | AstD.Let (p,[],e) ->
      delete_prime p;
      let e = opt_expr "" e in
      LetVar (p,e)
  | AstD.Let (p,fargs,e) ->
      delete_prime p;
      let var_unit =
        List.fold_left 
          (fun vu (b,p) ->
            if b then Sset.add p vu else vu)
          Sset.empty fargs
      in
      let e = delete_unit var_unit e in
      (* garde uniquement les éléments utiles *)
      let fargs = List.filter (fun (b,e)->not b) fargs in
      let fargs = List.map snd fargs in
      let e = opt_expr p e in
      LetFun (p,fargs,e)

let opt_prog prog =
  List.map opt_elt prog
