open Terms;;

type clos = Clos of env * lambda
 and env  = clos list
;;

let rec kn k = function
  (e, App(a,b), s)              -> kn k (e, a, (Clos(e,b)::s))
| (e, Lam f, (arg::s))          -> kn k ((arg::e), f, s)
| (((Clos(e',u))::_), Rel 0, s) -> kn k (e', u, s)
| ((_::e), Rel n, s)            -> kn k (e, Rel (n-1), s)
| (e, Lam f, [])                ->
  Lam (kn (k+1) (((Clos([], Rel (-k-1)))::e), f, []))
| ([], Rel n, s)                ->
  app_list (Rel (k+n)) (List.map (fun (Clos(e,t)) -> kn k (e,t,[])) s)
;;

let nkn t = kn 0 ([],t,[]);;


(* Very efficient optimization: expanding variables immediatly.
 * Makes the computation of pred(n) linear.
 *)
let rec clos_rel e i =
  match (e,i) with
    ((c::_), 0) -> c
  | ((_::e), n) -> clos_rel e (n-1)
  | ([], n)     -> Clos([], Rel n)
;;

let mk_clos e = function
    Rel i -> clos_rel e i
  | b -> Clos(e,b)
;;

let rec kn_opt k = function
| (e, App(a,b), s)            -> kn_opt k (e, a, (mk_clos e b::s))
| (e, Lam f, (arg::s))        -> kn_opt k ((arg::e), f, s)
| ((Clos(e',u)::_), Rel 0, s) -> kn_opt k (e', u, s)
| ((_::e), Rel n, s)          -> kn_opt k (e, Rel (n-1), s)
| (e, Lam f, [])              ->
    Lam (kn_opt (k+1) (((Clos([], Rel (-k-1)))::e), f, []))
| ([], Rel n, s)              -> fold_appl k (Rel (k+n)) s

and fold_appl k h = function
  [] -> h
| (Clos(e,t))::s -> fold_appl k (App(h, kn_opt k (e,t,[]))) s
;;

let norm_kam t = kn_opt 0 ([], t, [])
;;


(* Tail recursive versions *)
type stk =
    Ztop
  | Zapp of lambda * stk
  | Zlam of stk
  | Zarg of clos * stk
;;

let rec kn_tr k env t stk =
  match (t,stk) with
    (App(a,b),_) -> kn_tr k env a (Zarg (mk_clos env b, stk))
  | (Lam f, Zarg(arg,s)) -> kn_tr k (arg::env) f s
  | (Lam f, s) ->
      let nk = k+1 in
      let nv = Clos([], Rel (-nk)) in
        kn_tr nk (nv::env) f (Zlam s)
  | (Rel i, _) -> kn_tr_rel k env i stk

and kn_tr_rel k env i stk =
  match (i,env) with
    (0, (Clos(e,t)::_)) -> kn_tr k e t stk
  | (_, (_::e)) -> kn_tr_rel k e (i-1) stk
  | (_, []) -> zip_tr k (Rel (i+k)) stk

and zip_tr k t stk =
  match stk with
    Ztop -> t
  | Zapp(a,s) -> zip_tr k (App(a,t)) s
  | Zlam s -> zip_tr (k-1) (Lam t) s
  | Zarg(Clos(e,u), s) -> kn_tr k e u (Zapp(t,s))
;;


(* Imperative version: almost C code! *)
let k  = ref 0;;
let ge = ref [];;
let gs = ref Ztop;;

let rec kni = function
| App(a,b) ->
    gs := Zarg (mk_clos !ge b, !gs);
    kni a
| Lam f ->
    (match !gs with
       Zarg(arg,s) ->
         ge := arg :: !ge;
         gs := s;
         kni f
     | s ->
         k  := !k + 1;
         ge := (Clos ([], Rel(- !k))) :: !ge;
         gs := Zlam s;
         kni f)
| Rel i -> kni_rel (!ge, i)

and kni_rel = function
  (((Clos(e,t))::_), 0) ->
    ge := e;
    kni t
| ((_::e), i) -> kni_rel (e, i-1)
| ([], i) -> zip (Rel(i + !k))

and zip t =
  match !gs with
    Ztop -> t
  | Zapp(a,s) ->
      gs := s;
      zip (App(a,t))
  | Zlam s ->
      k  := !k - 1;
      gs := s;
      zip (Lam t)
  | Zarg(Clos(e,u), s) ->
      ge := e;
      gs := Zapp(t,s);
      kni u
;;

let norm_kni t =
  k  := 0; (* useless initialization *)
  ge := [];
  gs := Ztop;
  kni t
;;