(******************************************************************************
 * Définitions du type lambda_term et utilitaires de bases
 *
 * Pour utiliser au toplevel caml :
 * #load "term_lexer.cmo";;
 * #load "lambda.cmo";;
 *
 *****************************************************************************)

open List;;
open Term_lexer;;

(*****************************************************************************
 * Type des lambda-termes
 *****************************************************************************)

type lambda_term = Var of int * string
                 | App of lambda_term * lambda_term
                 | Abs of string * lambda_term
                 | Star
;;


(*****************************************************************************
 * Un normalisateur de lambda-termes basé sur la KAM
 *****************************************************************************)

(* types utilisées par la KAM *)
type closure    = Clos of lambda_term * environment
                | Free of int
and environment = closure list
and       stack = closure list
and       state = (lambda_term * environment * stack)
;;

type lambdas_list = (string * string) list
;;

(* Affichage *)
let rec string_of_lambda : lambda_term -> string =
  function
      Abs(x, t) -> "[" ^ x ^ "]" ^ (string_of_lambda t)
    | App(f, a) -> "(" ^ (string_of_lambda f) ^")" ^ (string_of_lambda a)
    | Var(_, x) -> x
    | Star      -> "*"
;;

let max_depth_limit = 100000;;

(* Rend la forme normale du terme passé en argument et le nombre de
   beta-réduction pour y arriver *)
let normalize_count : lambda_term -> (lambda_term * int) =
  let beta_count = ref 0 in
  let rec kam : int -> lambdas_list -> state -> lambda_term =
    fun depth lambdas ->
      if depth > max_depth_limit then failwith "Max recursion depth exceeded"
      else function
          (App(u, v), e, s) -> kam (depth + 1) lambdas (u, e, (Clos(v, e))::s)
        | (Abs(_, u), e, c::s) -> incr beta_count;
                                  kam (depth + 1) lambdas (u, c::e, s)
        | (Abs(x, u), e, []) ->
           let var_name = alpha lambdas x
           and lambdas_count = length lambdas
           in Abs(var_name,
                  kam (depth + 1) ((x, var_name)::lambdas)
                      (u, Free(lambdas_count + 1)::e, []))
        | (Var(i, _), e, s) -> (
          try
            match nth e i with
              Clos(vi, ei) -> kam (depth + 1) lambdas (vi, ei, s)
            | Free(lambda_index) ->
               let var_index = (length lambdas) - lambda_index
               in let (_, var_name) = nth lambdas var_index
                  in fold_left
                       (kam_arg depth lambdas)
                       (Var(var_index, var_name)) s
          with
            Failure(_) -> failwith "Free variable"
        )
        | (Star, _, []) -> Star
        | (Star, _, _) -> failwith "Arguments to constant *"

  and kam_arg: int -> lambdas_list -> lambda_term -> closure -> lambda_term =
    fun depth lambdas f ->
    function
      (Clos(a, m)) -> App(f, kam (depth + 1) lambdas (a, m, []))
    | _ -> failwith "Kam failure : invalid closure in stack"

  and alpha : lambdas_list -> string -> string =
    fun lambdas var_name ->
    let var_count = length (find_all (fun (x, _) -> x = var_name) lambdas)
    in if var_count = 0
       then var_name
       else var_name ^ (string_of_int var_count)

  in fun t -> let t0 = kam 0 [] (t, [], []) in (t0, !beta_count)
;;

let normalize : lambda_term -> lambda_term =
  fun t -> fst (normalize_count t)
;;


(*****************************************************************************
 * Parser : lambda-termes
 *****************************************************************************
 *
 * terme  = '$' ident | '#' ident
 *        | var |'[' var ']' terme | '(' terme ')' terme
 *
 * define_lambda "delta" "[x](x)x";;
 * define_lambda "deltadelta" "($delta)$delta"
 * define_macro  "xx" "(x)x";;
 * define_lambda "deltadelta" "([x]#xx)[x]#xx";;
 *
 *****************************************************************************)

(* Tables des termes et terme-macros prédéfinis *)
let lambda_table : (string, lambda_term) Hashtbl.t = Hashtbl.create(17);;
let lambda_macros_table : (string, string) Hashtbl.t = Hashtbl.create(17);;

let rec lambda_of_stream : token_stream -> lambda_term =
  let rec parse_term (binders : string list) : token_stream -> lambda_term =
    parser
        [< 'Kwd '$'; 'Ident name >]       -> Hashtbl.find lambda_table name
      | [< 'Kwd '#'; 'Ident name >]       ->
          parse_term binders (term_lexer
                                (Hashtbl.find lambda_macros_table name))
      | [< 'Kwd '[';
           'Ident x ?? "no variable name in lambda";
           'Kwd ']' ?? "lambda not terminated";
           t = parse_term (x::binders) >] -> Abs(x, t)
      | [< 'Kwd '(';
           f = parse_term binders;
           'Kwd ')' ?? "missing closing paren in application";
           a = parse_term binders >]      -> App(f, a)
      | [< 'Ident x >]                    -> Var(index_of x binders, x)
      | [< 'Kwd '*' >]                    -> Star
      | [< 'Kwd _ >]                      -> raise (Stream.Error
                                                      "invalid char")
      | [< 'Int _ >]                      -> raise (Stream.Error
                                                      "invalid char")
  in parse_term []

and lambda_of_string s = lambda_of_stream (term_lexer s)
;;

let define_lambda : ?dont_normalize : bool -> string -> string -> unit =
  fun ?(dont_normalize = false) name term ->
    Hashtbl.replace lambda_table name (
      if dont_normalize
      then lambda_of_string term
      else (normalize (lambda_of_string term))
    )
and define_macro : string -> string -> unit = fun name term ->
  Hashtbl.replace lambda_macros_table name term
;;

(* Normalisation avec conversion de et vers les strings *)
let exec : string -> string =
  fun s -> string_of_lambda (normalize (lambda_of_string s))
;;