(* 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
  | IfZ of fae * fae * fae
  | Rec of string * te * 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
  | ARecSub of string * faeValue ref * 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)
  | (findName, ARecSub(name, rv, restSc)) ->
      if (name = findName)
      then !rv
      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"))
  | (IfZ(tst, thn, els), sc) ->
      (match interp(tst, sc) with
	NumV(0) -> interp(thn, sc)
      | _ -> interp(els, sc))
  | (Rec(name, texpr, rhs, body), sc) ->
      let r = ref(NumV(0))
      in let sc = ARecSub(name, r, sc)
      in let v = interp(rhs, sc)
      in (r := v; 
          interp(body, sc))

(* 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")))
  | (IfZ(tst, thn, els), env) ->
      (match typecheck(tst, env) with
	NumT -> 
	  let thnType = typecheck(thn, env)
	  and elsType = typecheck(els, env)
	  in if (thnType = elsType)
	  then thnType
	  else raise (NoType (IfZ(tst, thn, els), "results mismatch"))
      | _ -> raise (NoType (IfZ(tst, thn, els), "test of non-number")))
  | (Rec(name, texpr, rhs, body), env) ->
      let bindType = parseType(texpr)
      in let env = ABind(name, bindType, env)
      in if (typecheck(rhs, env) = bindType)
      then typecheck(body, env)
      else raise (NoType (Rec(name, texpr, rhs, body), "bind mismatch"))

(* 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(interp(IfZ(Num(0), Num(1), Num(2)), MTSub), NumV(1));;
test(interp(IfZ(Num(1), Num(1), Num(2)), MTSub), NumV(2));;
test(interp(Rec("a", NumTE, Num(10), Add(Num(1), Id("a"))),
	    MTSub),
     NumV(11));;
test(interp(Rec("fib", ArrowTE(NumTE, NumTE), 
		Fun("x", NumTE, IfZ(Id("x"), Num(1),
				    IfZ(Sub(Id("x"), Num(1)), 
					Num(1),
					Add(App(Id("fib"), Sub(Id("x"), Num(1))),
					    App(Id("fib"), Sub(Id("x"), Num(2))))))),
		App(Id("fib"), Num(4))),
	    MTSub),
     NumV(5));;

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);;

test(typecheck(IfZ(Num(0), Num(1), Num(2)), MTEnv), NumT);;
test(typecheck(IfZ(Num(0), 
		   Fun("x", NumTE, Id("x")), 
		   Fun("y", NumTE, Num(10))),
		 MTEnv), 
     ArrowT(NumT, NumT));;
test(typecheck(Rec("a", NumTE, Num(10), Add(Num(1), Id("a"))),
	       MTEnv),
     NumT);;
test(typecheck(Rec("fib", ArrowTE(NumTE, NumTE), 
		Fun("x", NumTE, IfZ(Id("x"), Num(1),
				    IfZ(Sub(Id("x"), Num(1)), 
					Num(1),
					Add(App(Id("fib"), Sub(Id("x"), Num(1))),
					    App(Id("fib"), Sub(Id("x"), Num(2))))))),
		App(Id("fib"), Num(4))),
	    MTEnv),
     NumT);;

testExn(fun () -> typecheck(App(Num(1), Num(2)), MTEnv));;
testExn(fun () -> typecheck(Add(Fun("x",NumTE,Num(12)), Num(2)), MTEnv));;
testExn(fun () -> typecheck(IfZ(Num(0), 
                                Num(9),
				Fun("y", NumTE, Num(10))),
			    MTEnv));;
testExn(fun () -> typecheck(IfZ(Fun("x", NumTE, Id("x")), 
			        Fun("x", NumTE, Id("x")), 
				Fun("y", NumTE, Num(10))),
			    MTEnv));;
testExn(fun () -> typecheck(Rec("a", ArrowTE(NumTE, NumTE), Num(10), Add(Num(1), Id("a"))),
			    MTEnv));;

test(typecheck(Rec("fib", ArrowTE(NumTE, NumTE), 
		   Id("fib"),
		   App(Id("fib"), Num(10))),
	       MTEnv),
     NumT);;

(*
interp(Rec("fib", ArrowTE(NumTE, NumTE), 
	   Id("fib"),
	   App(Id("fib"), Num(10))),
       MTSub);;
*)