(****************************************
  typing.ml - typage par l'algorithme W
 ****************************************)

open Ast

module Int =
    struct
        type t = int
        let compare x y = x - y
    end

module Imap = Map.Make(Int)

(* Retourne soit un type de base, soit un type non défini *)
let rec head = function
    | Tvar tvar ->
        begin
        match tvar.def with
            | None -> Tvar tvar
            | Some t -> head t
        end
    | t -> t

(* Applique récursivement l'opération "head" *)
let rec canon t =
    match head t with
        | Tarrow (t1, t2) -> Tarrow (canon t1, canon t2)
        | Tproduct l -> Tproduct (List.map canon l)
        | Tlist t -> Tlist (canon t)
        | t -> t

(* Pour le debug et les erreurs : produit la représentation canonique d'un
   type sous forme de chaîne de caractères *)
let canon_string t tvar_names =
    let next_tvar_char = ref 97
    and tvar_names = ref tvar_names
    in
        let rec aux t =
            match t with
                | Tint -> "int"
                | Tstring -> "string"
                | Tunit -> "unit"
                | Tbool -> "bool"
                | Tarrow (t1, t2) -> aux_arrow t1 t2
                | Tproduct (t::l) ->
                    Printf.sprintf "%s"
                        (List.fold_left (fun base t ->
                                        Printf.sprintf "%s * %s"
                                                       base (aux t))
                                        (aux t)
                                        l)
                | Tproduct _ -> failwith "canon_string : Unreachable matching"
                | Tlist t ->
                    Printf.sprintf "%s list" (aux_list t)
                | Tvar tvar ->
                    try
                        Imap.find tvar.id !tvar_names
                    with Not_found ->
                        let tvar_name =
                            (Printf.sprintf "'%c" (char_of_int !next_tvar_char))
                        in
                            incr next_tvar_char;
                            tvar_names := Imap.add tvar.id
                                                   tvar_name
                                                   !tvar_names;
                            tvar_name
        and aux_arrow t1 t2 =
            let s1 =
                match t1 with
                    | Tarrow _ -> Printf.sprintf "(%s)" (aux t1)
                    | _ -> (aux t1)
            in
                let s2 = aux t2
                in
                    Printf.sprintf "%s -> %s" s1 s2
        and aux_list t =
            match t with
                | Tarrow (_, _)
                | Tproduct _ -> Printf.sprintf "(%s)" (aux t)
                | _ -> aux t
        in
            let s = aux (canon t)
            in
                s, !tvar_names

(* Pour le debug *)
let print_canon_string t =
    let s, _ = (canon_string t Imap.empty)
    in
        Printf.printf "%s\n" s

(* Erreurs *)
exception TypingFailure of loc option * string

let typing_error l err = raise (TypingFailure (l, err))

let unification_error l t1 t2 =
    let t1s, tvar_names = canon_string t1 Imap.empty
    in
        let t2s, _ = canon_string t2 tvar_names
        in
            let err = Printf.sprintf "Cette expression a le type %s mais est \
                                      ici utilisée avec le type %s" t1s t2s
            in
                typing_error l err

(* Détermine si la variable var apparaît dans le type t_raw *)
let rec occur var t_raw =
    match head t_raw with
        | Tarrow (t1, t2) -> (occur var t1) || (occur var t2)
        | Tproduct l -> List.exists (occur var) l
        | Tlist t -> occur var t
        | Tvar tvar when tvar.id = var.id -> true
        | t -> false

(* Réalise l'unification des types t1 et t2 *)
let rec unify loc t1 t2 =
    match (head t1, head t2) with
        | (a, b) when a = b -> ()
        | (Tarrow (t1a, t1b), Tarrow (t2a, t2b)) ->
            unify loc t1a t2a;
            unify loc t1b t2b
        | (Tproduct p1, Tproduct p2) ->
            List.iter2 (fun t1 t2 -> unify loc t1 t2) p1 p2
        | (Tlist t1, Tlist t2) -> unify loc t1 t2
        | (Tvar tv1, t2) when not (occur tv1 t2) -> (* tv1.def = None puisque
                                                       c'est le head de t1 *)
            tv1.def <- Some t2
        | (Tvar tv1, t2) -> unification_error loc (Tvar tv1) t2
        | (t1, Tvar tv2) ->
            unify loc (Tvar tv2) t1
        | (t1, t2) -> unification_error loc t1 t2

module V = struct
    type t = tvar
    let compare v1 v2 = Pervasives.compare v1.id v2.id
    let create = let r = ref 0 in fun () -> incr r; { id = !r; def = None }
end

module Vset = Set.Make(V)

(* Détermination des variables libres d'un type *)
let rec fvars = function
    | Tint
    | Tstring
    | Tbool
    | Tunit -> Vset.empty
    | Tarrow (t1, t2) -> Vset.union (fvars t1) (fvars t2)
    | Tproduct l ->
        List.fold_left
            (fun base t -> Vset.union base (fvars t)) Vset.empty l
    | Tlist t -> fvars t
    | Tvar { def = Some t } -> fvars t
    | Tvar tv -> Vset.singleton tv

type schema = { vars : Vset.t; typ : typ }
module Smap = Map.Make(String)
type env = { bindings : schema Smap.t; fvars : Vset.t }

(* Pour le debug : affiche de manière brute une liste de vars *)
let print_vars vars =
    Vset.iter (function tv -> Printf.printf "%d " tv.id) vars;
    print_newline ()

(* Pour le debug : affiche un type (représentation non canonique) *)
let rec format_type = function
    | Tint -> "int"
    | Tstring -> "string"
    | Tbool -> "bool"
    | Tunit -> "unit"
    | Tarrow (t1, t2) -> Printf.sprintf "%s -> %s" (format_type t1) (format_type t2)
    | Tproduct (t::l) ->
        List.fold_left (fun base t ->
                            Printf.sprintf "%s * %s" base (format_type t))
                       (format_type t)
                       l
    | Tproduct _ -> failwith "format_type : Unreachable matching"
    | Tlist t ->
        Printf.sprintf "%s list" (format_type t)
    | Tvar { id = id ; def = Some t } ->
        Printf.sprintf "var { %d : %s }" id (format_type t)
    | Tvar tvar ->
        Printf.sprintf "var %d" tvar.id

let empty_env = { bindings = Smap.empty ; fvars = Vset.empty }

(* Ajoute un identifiant typé à l'environnement sans généralisation *)
let add ident t env =
    { bindings = Smap.add ident { vars = Vset.empty ; typ = t } env.bindings;
      fvars = Vset.union env.fvars (fvars t) }

(* Folder pour la mise à jour des variables libres *)
let update_fvars_folder tv current =
    Vset.union (fvars (Tvar tv)) current

(* Ajoute un identifiant typé à l'environnement avec généralisation *)
let add_gen ident t env =
    (* Met à jour les variables libres *)
    let envfvars =
        Vset.fold update_fvars_folder env.fvars Vset.empty
    in
        (* Calcules les variables libres du type *)
        let binding_vars = Vset.diff (fvars t) envfvars
        in
            { bindings = Smap.add ident
                                  { vars = binding_vars ; typ = t }
                                  env.bindings ;
              fvars = envfvars }

(* Remplace la variable d'id orig par la fresh var new_tv  *)
let rec replace_var orig new_tv = function
    | Tarrow (t1, t2) ->
        Tarrow (replace_var orig new_tv t1, replace_var orig new_tv t2)
    | Tproduct l ->
        Tproduct (List.map (replace_var orig new_tv) l)
    | Tlist t ->
        Tlist (replace_var orig new_tv t)
    | Tvar tv when tv.id = orig ->
        new_tv
    | Tvar { id = id ; def = Some t } ->
        Tvar { id = id ; def = Some (replace_var orig new_tv t) }
    | t -> t

(* Remplace si besoin les occurences de la variable tv par des fresh vars *)
let use_freshvar tv t =
    match head (Tvar tv) with
        | Tvar tv -> replace_var tv.id (Tvar (V.create())) t
        | _ -> t

(* Retourne le type associé à l'identifiant ident, muni de fresh vars *)
let find ident env =
    let binding = Smap.find ident env.bindings
    in
        Vset.fold use_freshvar binding.vars binding.typ

module Sset = Set.Make(String)

(* Détermine l'ensemble des variables apparaîssant dans le motif *)
let rec motif_vars vars m =
    match m.m with
        | Munderscore -> vars
        | Mident ident when Sset.mem ident vars ->
                typing_error m.motif_loc
                             (Printf.sprintf "Variable %s non unique \
                                              dans ce motif" ident)
        | Mident ident -> Sset.add ident vars
        | Mtuple l -> List.fold_left motif_vars vars l

(* Ajoute les identifiants du motif m typé comme t à l'environnement env
   avec la fonction d'ajout add_func *)
let rec add_motif_base add_func env m t =
    match m.m with
        | Munderscore -> env
        | Mident ident -> add_func ident t env
        | Mtuple l ->
            begin
                match head t with
                    | Tproduct lp ->
                        List.fold_left2 (fun env m t -> add_motif_base add_func
                                                                       env
                                                                       m
                                                                       t)
                                        env
                                        l
                                        lp
                    | _ -> failwith "add_motif_base : Unreachable matching"
            end

(* Wrappers de la fonction précédente avec et sans généralisation *)
let add_motif_gen env m t = add_motif_base add_gen env m t
and add_motif env m t = add_motif_base add env m t

(* L'algorithme w ! *)
let rec w env e =
    let w_aux e =
        match e.e with
            (* Typage des identificateurs : on va chercher le type de
               l'identificateur correspondant dans l'environnement
               Au passage les free vars du type sont remplacées par des fresh
               vars si besoin *)
            | Eident ident ->
                begin
                try
                    find ident env
                with Not_found ->
                    typing_error e.loc
                                 (Printf.sprintf "Identifiant invalide : %s"
                                                 ident)
                end
            (* Typage des constantes : on retourne le type adéquat *)
            | Econst (Cint _) -> Tint
            | Econst (Cstring _) -> Tstring
            | Econst (Cbool _) -> Tbool
            | Econst Cunit -> Tunit
            | Econst (Cemptylist) ->
                Tlist (Tvar (V.create ()))
            (* Typage des tuples : on retourne juste le type produit
               correspondant *)
            | Etuple l ->
                Tproduct (List.map (w env) l)
            (* Typage des opérateurs : on unifie les types avec ceux attendus
               et on retourne le type correspondant *)
            | Ebinop (e1, OAnd, e2)
            | Ebinop (e1, OOr, e2) ->
                let t1 = w env e1
                and t2 = w env e2
                in
                    unify e1.loc t1 Tbool;
                    unify e2.loc t2 Tbool;
                    Tbool
            | Ebinop (e1, OEq, e2)
            | Ebinop (e1, ONeq, e2)
            | Ebinop (e1, OLt, e2)
            | Ebinop (e1, OLe, e2)
            | Ebinop (e1, OGt, e2)
            | Ebinop (e1, OGe, e2) ->
                let t1 = w env e1
                and t2 = w env e2
                in
                    unify e1.loc t1 Tint;
                    unify e2.loc t2 Tint;
                    Tbool
            | Ebinop (e1, OAdd, e2)
            | Ebinop (e1, OSub, e2)
            | Ebinop (e1, OMul, e2)
            | Ebinop (e1, ODiv, e2) ->
                let t1 = w env e1
                and t2 = w env e2
                in
                    unify e1.loc t1 Tint;
                    unify e2.loc t2 Tint;
                    Tint
            | Eunop (ONot, e) ->
                let t = w env e
                in
                    unify e.loc t Tbool;
                    Tbool
            | Eunop (ONeg, e) ->
                let t = w env e
                in
                    unify e.loc t Tint;
                    Tint
            | Eletin (m, val_expr, body_expr) ->
                (* Typage du let in :
                   * On type le motif et l'expression de la valeur assignée
                   * On unifie ces deux types
                   * On ajoute les identificateurs du motif à l'environnement
                     avec généralisation
                   * On calcule et retourne le type du corps du let in *)
                let m_t = w_motif m
                and val_t = w env val_expr
                in
                    unify val_expr.loc m_t val_t;
                    let env2 = add_motif_gen env m val_t
                    in
                        w env2 body_expr
            | Efunc f ->
                (* Typage des fonctions non récursives :
                   * On ajoute éventuellement le nom de la fonction à
                     l'environnement avec généralisation
                   * On type l'argument
                   * On l'ajoute sans généralisation à l'environnement
                   * On type le corps de la fonction
                   * On retourne arg_t -> return_t *)
                let arg_t = w_motif f.arg
                and return_t = Tvar (V.create())
                in
                    let env2 = add_motif env f.name (Tarrow (arg_t, return_t))
                    in
                        let env2 = add_motif env2 f.arg arg_t
                        in
                            unify f.body.loc (w env2 f.body) return_t;
                            Tarrow (arg_t, return_t)
            | Eif (cond, then_expr, else_expr) ->
                (* Typage des if :
                   * On type la condition et les deux expressions de retour
                   * On unifie le type de la condition avec bool
                   * On unifie les types des deux expressions de retour
                   * On retourne le type d'une des expressions de retour *)
                let cond_t = w env cond
                and then_t = w env then_expr
                and else_t = w env else_expr
                in
                    unify cond.loc cond_t Tbool;
                    unify else_expr.loc else_t then_t;
                    then_t
            | Elistcons (head, tail) ->
                (* Typage des listcons :
                   * On type la tête et la queue
                   * On unifie (head_t list) et tail_t
                   * On retourne le type de la queue (ou head_t list) *)
                let head_t = w env head
                and tail_t = w env tail
                in
                    unify head.loc (Tlist head_t) tail_t;
                    Tlist head_t
            | Ecall (func_expr, arg) ->
                (* Typage des calls :
                   * On type la fonction et son argument
                   * On unifie le type de la fonction avec argument -> 'a
                   * On retourne le type donné à 'a par unification *)
                let func_t = w env func_expr
                and arg_t = w env arg
                and alpha_t = Tvar (V.create ())
                and beta_t = Tvar (V.create ())
                in
                    unify func_expr.loc func_t (Tarrow (alpha_t, beta_t));
                    unify arg.loc arg_t alpha_t;
                    beta_t
            | Ematch (match_expr, empty_expr, head_m, tail_m, expr) ->
                (* Typage du match :
                   * On type l'expression matchée, l'expression retournée
                     dans le cas où la liste est vide, on détermine un type
                     brut pour les motifs dans le cas non vide ; 
                   * On réalise l'unification du type de l'expression matchée
                     et des motifs représentant la liste dans le cas non vide
                   * On met à jour l'environnement avec les identifieurs
                     éventuels des motifs du cas non vide
                     Question : dans la sémantique de ocaml, on ne généralise
                     pas ici, ce qui peut provoquer d'ennuyeux problèmes à cet
                     endroit, par exemple avec :
                     let _ =
                        match [(function x -> x)] with
                            [] -> ()
                          | id::l ->
                            let _ = print_int (id 42)
                            in
                                print_string (id "42")
                     Au contraire, généraliser empêcherait de correctement
                     typer certaines autres expressions, donc le comportement
                     à adopter est difficiler à déterminer
                   * On détermine le type de l'expression retournée dans le cas
                     non vide, et on l'unifie avec celui de l'expression dans
                     le cas vide
                   * On retourne le type de l'expression retournée dans le cas
                     non vide *)
                let match_t = w env match_expr
                and empty_t = w env empty_expr
                and head_t = w_motif head_m
                and tail_t = w_motif tail_m
                in
                    unify match_expr.loc match_t (Tlist head_t);
                    begin
                        match tail_m.m with
                            | Mtuple _ ->
                                typing_error tail_m.motif_loc
                                             "Motif interdit dans le \
                                              filtrage de la queue de \
                                              liste"
                            | Mident ident -> ()
                            | Munderscore -> ()
                    end;
                    unify match_expr.loc tail_t (Tlist head_t);
                    let env2 = add_motif (add_motif env head_m head_t)
                                         tail_m
                                         tail_t
                    and _ = motif_vars Sset.empty
                                       ({ m = (Mtuple [head_m ; tail_m]) ;
                                          motif_loc = None ;
                                          motif_t = None })
                    in
                        let result_t = w env2 expr
                        in
                            unify expr.loc result_t empty_t;
                            result_t
            | Eclos _ -> failwith "w : Unreachable matching"
    in
        let t = w_aux e
        in
            e.t <- Some t;
            t
and w_motif m =
    (* Type un motif, tout en vérifiant l'unicité des variables apparaîssant
       dans le motif *)
    let vars = ref Sset.empty
    in
        let rec aux m =
            match m.m with
                | Munderscore -> Tvar (V.create ())
                | Mident ident when Sset.mem ident !vars ->
                    typing_error m.motif_loc
                                 (Printf.sprintf "Variable %s non unique \
                                                  dans ce motif" ident)
                | Mident ident ->
                    let t = Tvar (V.create ())
                    in
                        vars := Sset.add ident !vars;
                        t
                | Mtuple l ->
                    Tproduct (List.map aux l)
        in
            let t = aux m
            in
                m.motif_t <- Some t;
                t

(* Environnement de base, contenant les primitives *)
let base_env =
    add "print_int" (Tarrow (Tint, Tunit))
        (add "print_string" (Tarrow (Tstring, Tunit))
            (add "print_newline" (Tarrow (Tunit, Tunit))
                (add "read_int" (Tarrow (Tunit, Tint))
                    (add "read_line" (Tarrow (Tunit, Tstring)) empty_env))))

(* L'ultime helper : type un ast dans l'environnement de base *)
let type_ast a =
    w base_env a