open Tfbsrxast;;
open Tfbsrxpp;;

exception TypeError

(*
 * If you would like typechecking to be enabled by your interpreter by default,
 * then change the following value to true.  Whether or not typechecking is
 * enabled by default, you can explicitly enable it or disable it using
 * command-line arguments.
 *)
let typecheck_default_enabled = true;;

(*
 * Replace this with your typechecker code.  Your code should not throw the
 * following exception; if you need to raise an exception, create your own
 * exception type here.
 *)

let rec typecheck_aux gamma e = match e with
    | Var(ident) -> lookup gamma ident
    | Function(ident, fbtype, expr) ->
         let t = typecheck_aux ((ident, fbtype)::gamma) expr in TArrow(fbtype, t)
    | Letrec(ident, ident2, fbtype, expr1, fbtype2, expr2) -> TInt
    | Appl(expr1, expr2) -> let TArrow(t1, t2) = (typecheck_aux gamma expr1) in
              if typecheck_aux gamma expr2 = t1 then t2 else raise TypeError
    | Plus(expr1, expr2) ->
        if (typecheck_aux gamma expr1 = TInt && typecheck_aux gamma expr2 = TInt)
        then TInt else raise TypeError
    | Minus(expr1, expr2) ->
        if (typecheck_aux gamma expr1 = TInt && typecheck_aux gamma expr2 = TInt)
        then TInt else raise TypeError
    | Equal(expr1, expr2) ->
        if (typecheck_aux gamma expr1 = TInt && typecheck_aux gamma expr2 = TInt)
        then TBool else raise TypeError
    | And(expr1, expr2) ->
        if (typecheck_aux gamma expr1 = TBool && typecheck_aux gamma expr2 = TBool)
        then TBool else raise TypeError
    | Or(expr1, expr2) ->
    if (typecheck_aux gamma expr1 = TBool && typecheck_aux gamma expr2 = TBool)
        then TBool else raise TypeError
    | Not(expr1) -> if (typecheck_aux gamma expr1 = TBool)
        then TBool else raise TypeError
    | If(expr1, expr2, expr3) ->
         if (typecheck_aux gamma expr1) = TBool
         then (let (t1, t2) = ((typecheck_aux gamma expr2), (typecheck_aux gamma expr3))
               in (if t1 = t2 then t1 else raise TypeError))
     else raise TypeError
    | Int(int) -> TInt
    | Bool(bool) -> TBool
    | Ref(expr) -> TRef(typecheck_aux gamma expr)
    | Set(expr1, expr2) -> let TRef(t) = (typecheck_aux gamma expr1) in
          if (typecheck_aux gamma expr2) = t then t else raise TypeError
    | Get(expr) -> let TRef(t) = (typecheck_aux gamma expr) in t
    | Record(recordlist) -> TRec(rec_helper recordlist gamma)
    | Select(label, expr) -> let TRec(t) = (typecheck_aux gamma expr) in (select_helper t label)
    | Raise(exnid, fbtype, expr) -> TInt
    | Try(expr1, exnid, ident, fbtype, expr2) -> TInt
    | Cell(int) -> TInt

 and lookup gamma ident = match gamma with
      [] -> raise TypeError
     | (id, t)::xs -> if id = ident then t
     	            else lookup xs ident

 and rec_helper recordlist gamma = match recordlist with
     [] -> []
     | (l, e)::rs -> (l, (typecheck_aux gamma e)) :: (rec_helper rs gamma)

 and select_helper reclist label =
     match reclist with
      [] -> raise TypeError
     | (l, t)::rs -> if l = label then t else (select_helper rs label)

;;

let typecheck e = typecheck_aux [] e;;