(* Data definitions *)

type fae = 
    Num of int
  | Add of fae * fae
  | Sub of fae * fae
  | Id of string
  | Fun of string * te * fae
  | App of fae * fae

and te =
    NumTE
  | BoolTE
  | ArrowTE of te * te

and faeValue =
    NumV of int
  | ClosureV of string * fae * subCache

and subCache =
    MTSub
  | ASub of string * faeValue * subCache

and ty =
    NumT
  | BoolT
  | ArrowT of ty * ty

and typeEnv =
    MTEnv
  | ABind of string * ty * typeEnv

(* Lookup *)

exception Failed of string

let rec lookup = function 
    (findName, MTSub) -> raise (Failed "free variable")
  | (findName, ASub(name, v, restSc)) ->
      if (name = findName)
      then v
      else lookup(findName, restSc)

(* Number operations *)

let mkNumOp = fun op
  -> function
      (NumV(a), NumV(b)) -> NumV(op a b)
    | _ -> raise (Failed "not numbers")
    
let numPlus = mkNumOp (fun a b -> a + b)
let numMinus = mkNumOp (fun a b -> a - b)

(* Interpreter *)

let rec interp : (fae * subCache -> faeValue ) = function
    (Num(n), sc) -> NumV(n)
  | (Add(l, r), sc) -> numPlus(interp(l, sc),
                               interp(r, sc))
  | (Sub(l, r), sc) -> numMinus(interp(l, sc),
                                interp(r, sc))
  | (Id(name), sc) -> lookup(name, sc)
  | (Fun(param, texpr, body), sc) -> ClosureV(param, body, sc)
  | (App(fn, arg), sc) ->
      let funV = interp(fn, sc)
      and argV = interp(arg, sc)
      in match funV with
          ClosureV(param, body, sc) -> 
            interp(body, ASub(param, argV, sc))
      | _ -> raise (Failed "not a function")

(* Type-environment lookup *)

let rec gettype = function 
    (findName, MTEnv) -> raise (Failed "free variable, so no type")
  | (findName, ABind(name, t, restEnv)) ->
      if (name = findName)
      then t
      else gettype(findName, restEnv)

(* Type expression to type *)

let rec parseType = function
    NumTE -> NumT
  | BoolTE -> BoolT
  | ArrowTE(a, b) -> ArrowT(parseType a, parseType b)

(* Type checking *)

exception NoType of fae * string

let rec typecheck : (fae * typeEnv -> ty) = function
    (Num(n), env) -> NumT
  | (Add(l, r), env) -> 
      (match (typecheck(l,env), typecheck(r, env)) with
        (NumT, NumT) -> NumT
      | _ -> raise (NoType (Add(l, r), "+: not both nums")))
  | (Sub(l, r), env) -> 
      (match (typecheck(l,env), typecheck(r, env)) with
        (NumT, NumT) -> NumT
      | _ -> raise (NoType (Add(l, r), "-: not both nums")))
  | (Id(name), env) -> gettype(name, env)
  | (Fun(param, texpr, body), env) -> 
      let argType = parseType(texpr)
      in ArrowT(argType, typecheck(body, ABind(param, argType, env)))
  | (App(fn, arg), env) ->
      (match typecheck(fn, env) with
        ArrowT(paramType, resultType) ->
          let argType = typecheck(arg, env)
          in if (argType = paramType)
          then resultType
          else raise (NoType (App(fn,arg),"argument/parameter type mismatch"))
      | _ -> raise (NoType (App(fn, arg), "application of non-function")))

(* Testing *)

let test = function 
    (a, b) ->
      if (a = b)
      then true
      else raise (Failed "test failure")
        
let testExn = function
     f -> test((try (let _ = f () in false) with NoType(e, s) -> true),
               true)

;;
  
(* Test cases *)

test(interp(Num(10), MTSub), NumV(10));;
test(interp(Add(Num(10), Num(17)), MTSub), NumV(27));;
test(interp(Sub(Num(10), Num(7)), MTSub), NumV(3));;
test(interp(App(Fun("x", NumTE, Add(Id("x"),Num(12))),
                Add(Num(1), Num(17))),
            MTSub),
     NumV(30));;
test(interp(App(Fun("x", NumTE,
                    App(Fun("f", NumTE,
                            Add(App(Id("f"), Num(1)),
                                App(Fun("x", NumTE, App(Id("f"), Num(2))),
                                    Num(3)))),
                        Fun("y", NumTE, Add(Id("x"), Id("y"))))),
                Num(0)),
            MTSub),
     NumV(3));;

test(typecheck(Num(10), MTEnv), NumT);;
test(typecheck(Add(Num(10), Num(17)), MTEnv), NumT);;
test(typecheck(Sub(Num(10), Num(7)), MTEnv), NumT);;
test(typecheck(Fun("x", NumTE, Add(Id("x"),Num(12))),
               MTEnv),
     ArrowT(NumT, NumT));;
test(typecheck(Fun("x", NumTE, Fun("y", BoolTE, Id("x"))),
               MTEnv),
     ArrowT(NumT, ArrowT(BoolT, NumT)));;
test(typecheck(App(Fun("x", NumTE, Add(Id("x"),Num(12))),
                   Add(Num(1), Num(17))),
               MTEnv),
     NumT);;
test(typecheck(App(Fun("x", NumTE,
                       App(Fun("f", ArrowTE(NumTE, NumTE),
                               Add(App(Id("f"), Num(1)),
                                   App(Fun("x", NumTE, App(Id("f"), Num(2))),
                                       Num(3)))),
                           Fun("y", NumTE, Add(Id("x"), Id("y"))))),
                   Num(0)),
               MTEnv),
     NumT);;

testExn(fun () -> typecheck(App(Num(1), Num(2)), MTEnv));;

testExn(fun () -> typecheck(Add(Fun("x",NumTE,Num(12)), Num(2)), MTEnv));;