(****************************************
    closure.ml - gestion des clotures
 ****************************************)

open Ast
open Utils

module Sset = Set.Make(String)

let next_fun_id = ref 1

(* Gestion des conflits de labels sur les fonctions *)
let rec new_function_label idents =
    let label = Printf.sprintf "fun___%d" !next_fun_id
    in
        incr next_fun_id;
        if Sset.mem label idents then
            new_function_label idents
        else
            label

(* Petit helper, générant un identifiant à la volée si nécessaire *)
let get_function_name idents f =
    match f.name.m with
        | Munderscore -> new_function_label idents
        | Mident s -> s
        | _ -> failwith "get_function_name : Unreachable matching"

(* Calcul des identifiants d'un motif *)
let rec motif_vars m =
    match m.m with
        | Munderscore -> Sset.empty
        | Mident id -> Sset.singleton id
        | Mtuple l -> List.fold_left (fun base motif ->
                                        Sset.union base (motif_vars motif))
                                     Sset.empty l

(* Détermination des variables libres d'une expression,
   utilisée lors de la construction des clotures et letfuns *)
let rec free_vars expr =
    match expr.e with
        | Econst _ -> Sset.empty
        | Eident s -> Sset.singleton s
        | Etuple l -> List.fold_left (fun base expr ->
                                        Sset.union base (free_vars expr))
                                     Sset.empty l
        | Ebinop (e1, _, e2)
        | Elistcons (e1, e2)
        | Ecall (e1, e2) -> Sset.union (free_vars e1) (free_vars e2)
        | Eunop (_, e) -> free_vars e
        | Eif (e1, e2, e3) ->
            Sset.union (Sset.union (free_vars e1) (free_vars e2))
                       (free_vars e3)
        | Eletin ({ m = Mident id }, { e = (Efunc f) }, e2) when f.recursive ->
            Sset.diff (Sset.union (free_vars (raw_expr (Efunc f)))
                                  (free_vars e2))
                      (Sset.singleton id)
        | Eletin (m, e1, e2) ->
            Sset.union (free_vars e1)
                       (Sset.diff (free_vars e2) (motif_vars m))
        | Efunc f -> Sset.diff (free_vars f.body) (motif_vars f.arg)
        | Ematch (match_expr, empty_expr, head_m, tail_m, result_expr) ->
            let motifs_idents = Sset.union (motif_vars head_m)
                                           (motif_vars tail_m)
            in
                Sset.union (Sset.union (free_vars match_expr)
                                       (free_vars empty_expr))
                           (Sset.diff (free_vars result_expr) motifs_idents)
        | Eclos _ -> failwith "free_vars : Unreachable matching"

(* Parcours récursif de l'arbre pour remplacer les fonctions par des
   clotures et des letfuns, retourne un (expr * (letfun list)) *)
let rec replace_funcs idents expr =
    match expr.e with
        | Econst _ -> expr, []
        | Eident _ -> expr, []
        | Etuple l ->
            let l, letfuns =
                List.fold_right (fun e (l, letfuns) ->
                                    let e2, letfuns2 = replace_funcs idents e
                                    in
                                        (e2::l, letfuns2 @ letfuns))
                                l  ([], [])
            in
                (raw_expr (Etuple l)), letfuns
        | Ebinop (e1, o, e2) ->
            let e1, letfuns1 = replace_funcs idents e1
            and e2, letfuns2 = replace_funcs idents e2
            in
                (raw_expr (Ebinop (e1, o, e2))), letfuns1 @ letfuns2
        | Eunop (o, e) ->
            let e, letfuns = replace_funcs idents e
            in
                (raw_expr (Eunop (o, e))), letfuns
        | Eletin (m, val_expr, body_expr) ->
            let val_expr, letfuns1 = replace_funcs idents val_expr
            and body_expr, letfuns2 = replace_funcs idents body_expr
            in
                (raw_expr (Eletin (m, val_expr, body_expr))),
                letfuns1 @ letfuns2
        | Efunc f ->
            let fname = get_function_name idents f
            in
                let fvars = Sset.elements (free_vars expr)
                in
                    (raw_expr (Eclos (f.recursive, fname, fvars))),
                    (build_letfun idents fname f fvars)
        | Eif (cond_expr, then_expr, else_expr) ->
            let cond_expr, letfuns1 = replace_funcs idents cond_expr
            and then_expr, letfuns2 = replace_funcs idents then_expr
            and else_expr, letfuns3 = replace_funcs idents else_expr
            in
                (raw_expr (Eif (cond_expr, then_expr, else_expr))),
                letfuns1 @ letfuns2 @ letfuns3
        | Elistcons (head_expr, tail_expr) ->
            let head_expr, letfuns1 = replace_funcs idents head_expr
            and tail_expr, letfuns2 = replace_funcs idents tail_expr
            in
                (raw_expr (Elistcons (head_expr, tail_expr))),
                letfuns1 @ letfuns2
        | Ecall (func_expr, arg_expr) ->
            let func_expr, letfuns1 = replace_funcs idents func_expr
            and arg_expr, letfuns2 = replace_funcs idents arg_expr
            in
                (raw_expr (Ecall (func_expr, arg_expr))),
                letfuns1 @ letfuns2
        | Ematch (match_expr, empty_expr, head_m, tail_m, result_expr) ->
            let match_expr, letfuns1 = replace_funcs idents match_expr
            and empty_expr, letfuns2 = replace_funcs idents empty_expr
            and result_expr, letfuns3 = replace_funcs idents result_expr
            in
                (raw_expr (Ematch (match_expr, empty_expr,
                                   head_m, tail_m, result_expr))),
                letfuns1 @ letfuns2 @ letfuns3
        | Eclos _ -> failwith "replace_funcs : Unreachable matching"
and build_letfun idents fname f fvars =
    let fbody, letfuns = replace_funcs idents f.body
    in
        letfuns @ [(fname, fvars, f.arg, fbody)]