open Util.Vars

module Pure =
struct

type t =
  | V of int
  | A of t * t
  | L of t

let rec print ?(l=[]) =
 function
    V n -> print_name l n
  | A(t1,t2) -> "(" ^ print ~l t1 ^ " " ^ print ~l t2 ^ ")"
  | L t ->
     let name = string_of_var (List.length l) in
     "λ" ^ name ^ "." ^ print ~l:(name::l) t

let lift m =
 let rec aux l =
  function
   | V n -> V (if n < l then n else n+m)
   | A (t1,t2) -> A (aux l t1, aux l t2)
   | L t -> L (aux (l+1) t)
 in
  aux 0

(* Reference implementation.
   Reduction machine used below
let subst delift_by_one what with_what =
 let rec aux l =
  function
   | A(t1,t2) -> A(aux l t1, aux l t2)
   | V n ->
       if n = what + l then
        lift l with_what
       else
        V (if delift_by_one && n >= l then n-1 else n)
   | L t -> L (aux (l+1) t)
 in
  aux 0

let rec whd =
 function
  | A(t1, t2) ->
     let t2 = whd t2 in
     let t1 = whd t1 in
      (match t1 with
        | L f -> whd (subst true 0 t2 f)
        | V _
        | A _ -> A(t1,t2))
  | V _
  | L _ as t -> t
*)

let unwind ?(tbl = Hashtbl.create 317) m =
 let rec unwind (e,t,s) =
  let cache_unwind m =
   try
    Hashtbl.find tbl m
   with
    Not_found ->
     let t = unwind m in
      Hashtbl.add tbl m t ;
      t in
  let s = List.map cache_unwind s in
  let rec aux l =
   function
    | A(t1,t2) -> A(aux l t1, aux l t2)
    | V n as x when n < l -> x
    | V n ->
       (try
         lift l (cache_unwind (List.nth e (n - l)))
        with Failure _ -> V (n - l))
    | L t -> L (aux (l+1) t) in
  let t = aux 0 t in
  List.fold_left (fun f a -> A(f,a)) t s
in
 unwind m

let mwhd m =
 let rec aux =
  function
     (e,A(t1,t2),s) ->
       let t2' = aux (e,t2,[]) in
       aux (e,t1,t2'::s)
   | (e,L t,x::s) -> aux (x::e,t,s)
   | (e,V n,s) as m ->
      (try
        let e,t,s' = List.nth e n in
        aux (e,t,s'@s)
       with Failure _ -> m)
   | (_,L _,[]) as m -> m
 in
  unwind (aux m)

end

module Scott =
struct

open Pure

let rec mk_n n =
 if n = 0 then L (L (A (V 1, L (V 0)))) else L (L (A (V 0, mk_n (n-1))))

let dummy = V (max_int / 2)

let mk_match t bs =
 let bs = List.sort (fun (n1,_) (n2,_) -> compare n1 n2) bs in
 let rec aux m t =
  function
     [] -> dummy
   | (n,p)::tl as l ->
      if n = m then
       A (A (t, L (lift (m+1) p)), L (aux (m+1) (V 0) tl))
      else
       A (A (t, dummy), L (aux (m+1) (V 0) l))
 in
  aux 0 t bs

end