type id = string

type ty_annot = id * Types.ty

type lam =
  Fun of ty_annot * lam
  | App of lam * lam
  | Var of id
  | Exf of lam * Types.ty

(** alpha renaming in a deterministic way
    if readable is set to true, original variable's names are used to rename *)
let alpha_convert ?(readable=false) (e : lam) : lam =
  let cpt_var = ref 0 in (* id cpt for variables *)
  let generator (v : id) (cpt : int ref) : id = (* id generator *)
    let start = if readable then v else "" in
    let v = start ^ "__"^(string_of_int !cpt) in
    incr cpt; v

  in let get_ren_var (var : id) (g : (id * id) list): id = (* returns var if not found *)
    match List.assoc_opt var g with
      | Some var' -> var'
      | None -> var

  in let rec alpha_aux (e : lam) (g : (id * id) list): lam = match e with (* actual renaming *)
    Fun ((v, t), e) ->
      let v' = generator v cpt_var in
      Fun ((v', t), alpha_aux e ((v, v')::g))
    | App (e1, e2) -> App (alpha_aux e1 g, alpha_aux e2 g)
    | Var v -> Var (get_ren_var v g)
    | Exf (e, t) -> Exf (alpha_aux e g, t)
  in alpha_aux e []


(** alpha equivalence *)
let (=~) (e1 : lam) (e2 : lam) : bool =
  alpha_convert e1 = alpha_convert e2

(** beta-reduction *)
(* subst m n x substitutes x for n in m *)
let rec subst (m : lam) (n : lam) (x : id) : lam =
  match m with
  Fun ((y, t), e) ->
    if x = y
    then m
    else Fun ((y, t), subst e n x)
  | App (e1, e2) ->
    let e1 = subst e1 n x in
    let e2 = subst e2 n x in
    App (e1, e2)
  | Var y ->
    if x = y then n else m
  | Exf (e, t) ->
    let e = subst e n x in
    Exf (e, t)

(* INVARIANT : e has already been alpha-converted *)
let rec betastep (e : lam) : lam option =
  match e with
  Fun ((x, t), e) ->
    begin match betastep e with
    None -> None
    | Some e -> Some (Fun ((x, t), e))
    end
  | App (e1, e2) ->
    (* reduce leftmost redex.
       e1 e2 -> e1' e2 then
       e1_nf e2 -> e1_nf e2 then
       (fun x -> e_nf) e2_nf -> e_nf[e2/x] *) 
    begin match betastep e1 with
    None ->
      begin match betastep e2 with
      None ->
        begin match e1 with
        Fun ((x, _), e) -> Some (subst e e2 x)
        | _ -> None
        end
      | Some e2 -> Some (App (e1, e2))
      end
    | Some e1 -> Some (App (e1, e2))
    end
  | Var _ -> None
  | Exf (e, t) ->
    begin match betastep e with
    None -> None
    | Some e -> Some (Exf (e, t))
    end