(* 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
  | GuessTE

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
  | VarT of ty option ref

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)
  | GuessTE -> VarT(ref None)

(* Type checking *)

exception NoType of fae * string

let rec occurs = function
    (r, t) -> match t with
      NumT -> false
    | BoolT -> false
    | ArrowT(a, b) -> occurs(r, a) || occurs(r, b)
    | VarT(r2) -> 
	((r == r2) ||
	(match !r2 with
	  None -> false
	| Some(t) -> occurs(r, t)))

let rec resolve = function
    t -> match t with 
      VarT(r) ->
	(match !r with
	  None -> t
	| Some(t) -> resolve(t))
    | _ -> t

let samevar = function
    (VarT(r1), VarT(r2)) -> r1 == r2
  | _ -> false
	  
let rec unify = function
    (t1, t2, expr) -> match (t1, t2) with
      (VarT(r), _) ->
	(match !r with
	  Some(t1) -> unify(t1, t2, expr)
	| None ->
	    let t2 = resolve(t2)
	    in if samevar(t1, t2)
	    then ()
	    else if occurs(r, t2)
	    then raise (NoType(expr, "occurs check failed"))
	    else r := Some(t2))
    | (_, VarT(r)) -> unify(t2, t1, expr)
    | (ArrowT(a1, b1), ArrowT(a2, b2))->
	(unify(a1, a2, expr);
	 unify(b1, b2, expr))
    | (NumT, NumT) -> ()
    | (BoolT, BoolT) -> ()
    | _ ->  raise (NoType(expr, "unification failed"))

let rec typecheck : (fae * typeEnv -> ty) = function
    (expr, env) -> match expr with
      Num(n) -> NumT
    | Add(l, r) -> 
	(unify(typecheck(l,env), NumT, l);
	 unify(typecheck(r,env), NumT, r);
	 NumT)
    | Sub(l, r) -> 
	(unify(typecheck(l,env), NumT, l);
	 unify(typecheck(r,env), NumT, r);
	 NumT)
    | Id(name) -> gettype(name, env)
    | Fun(param, texpr, body) -> 
	let argType = parseType(texpr)
	in ArrowT(argType, typecheck(body, ABind(param, argType, env)))
    | App(fn, arg) ->
	let resultType = VarT(ref None)
	in (unify(typecheck(fn, env), 
		  ArrowT(typecheck(arg, env), resultType), 
		  expr);
	    resultType)
    | IfZ(tst, thn, els) ->
	(unify(typecheck(tst, env), NumT, tst);
	 let thnType = typecheck(thn, env)
	 and elsType = typecheck(els, env)
	 in (unify(thnType, elsType, expr);
	     thnType))
    | Rec(name, texpr, rhs, body) ->
	let bindType = parseType(texpr)
	in let env = ABind(name, bindType, env)
	in (unify(typecheck(rhs, env), bindType, expr);
	    typecheck(body, env))
	  
(* Testing *)

let test = function 
    (a, b) ->
      if (a = b)
      then true
      else raise (Failed "test failure")
        
let testnotype = function
    thunk ->
      test ((try (let v = thunk () in false) with NoType(a, b) -> 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(unify(typecheck(App(Fun("x", NumTE, Add(Id("x"),Num(12))),
			 Add(Num(1), Num(17))),
		     MTEnv),
	   NumT, 
	   Num(-1)),
     ());;

test(unify(typecheck(App(Fun("x", GuessTE, Add(Id("x"),Num(12))),
			 Add(Num(1), Num(17))),
		     MTEnv),
	   NumT, 
	   Num(-1)),
     ());;
test(unify(typecheck(Fun("x", GuessTE, Add(Id("x"),Num(12))),
		     MTEnv),
	   ArrowT(NumT,NumT),
	   Num(-1)),
     ());;
test(unify(typecheck(Fun("x", GuessTE, IfZ(Num(0),Id("x"),Id("x"))),
		     MTEnv),
	   ArrowT(NumT,NumT),
	   Num(-1)),
     ());;

test(unify(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,
	   Num(-1)),
     ());;

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(unify(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,
	   Num(-1)),
     ());;

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

testnotype(fun () -> unify(typecheck(Fun("x", GuessTE, Add(Id("x"),Num(12))),
				     MTEnv),
			   ArrowT(BoolT,NumT),
			   Num(-1)));;

testnotype(fun () -> typecheck(Fun("x", GuessTE, App(Id("x"),Id("x"))),
			       MTEnv));;