(* Abstract Syntax Trees for types, values, expressions, and environments *)

type ty =
    Int
  | Fntype of ty * ty;;

type value =
    N of int
  | Lam of string * ty * exp
and exp =
    V of value
  | Var of string
  | App of exp * exp;;

type environment =
    EmptyEnv
  | Env of string * ty * environment;;

exception TypeErrorUnboundVar;;
exception TypeErrorApp;;

(* lookup: Get the type bound to x in environment env *)

let rec lookup env x =
  match env with
      EmptyEnv -> raise TypeErrorUnboundVar
    | Env(y, t, env2) ->
        if x = y then t else lookup env2 x;;
    
(* check: type check e in environment env *)

let rec check env e = 
  match e with
      V(N n)               -> Int
    | V(Lam(x, t, e2)) ->
        let nuenv = Env(x,t,env) in
        let te2 = check nuenv e2 in 
          Fntype(t,te2)
    | Var(x) -> lookup env x
    | App(e1,e2) ->
        match check env e1 with
            Fntype(te2,t) -> 
              if te2 = check env e2 
              then t 
              else raise TypeErrorApp
          | _ -> raise TypeErrorApp;;

(* some test programs *)

let prog1 = App(V(Lam("x",Int,Var("x"))),V(N(4)));;
let prog2 = App(V(Lam("x",Int,Var("y"))),V(N(4)));;
let prog3 = App(V(Lam("x",Int,Var("x"))),V(Lam("x",Int,Var("x"))));;
check EmptyEnv prog1;;
(* check EmptyEnv prog2;; *)

exception CoreDumpUnboundVar;;
exception CoreDumpAlreadyValue;;
exception CoreDumpBadFn;;

(* subst: Replace each free occurrence of x in e by ex *)

let rec subst ex x e =
  match e with
      V (N n) -> e
      | V(Lam(y, t, e2)) ->
        if x = y then e
        else V(Lam(y,t,subst ex x e2))
    | Var y -> if x = y then ex else e
    | App(e1,e2) -> App(subst ex x e1,subst ex x e2);;

(* step: perform a single step of e, which should not be a value *)

let rec step e =
  match e with
      V v -> raise CoreDumpAlreadyValue
    | Var(x) -> raise CoreDumpUnboundVar
    | App(V(Lam(x, t, e2)),V arg) ->
        subst (V arg) x e2 
    | App(V(_),V(_)) -> raise CoreDumpBadFn
    | App(V(v1),e2) -> App(V(v1),step e2)
    | App(e1,e2) -> App(step e1,e2);;

(* run: execute e until it yields a value *)

let rec run e =
  match e with
      V v -> e
    | _ -> run (step e);;