(*
 *
 * BU CAS CS 520, F2002
 * Assignment 3
 *
 *)

(* The code is written Hongwei Xi on Oct 10, 2002 *)

signature TYPE_CHECKER =
    sig
	val typeCheck: Parser.ttm -> bool
    end

structure TypeChecker :> TYPE_CHECKER =
    struct
	structure P = Parser

	exception Fatal

	fun fatal (msg) = (print msg; raise Fatal)

	type context = (string * P.stp) list

	val tpBool = P.TpBase "bool"
	val tpInt = P.TpBase "int"
	val tpStr = P.TpBase "string"
	val tpUnit = P.TpTup []

	val Sigma: context =
	    [
	     ("~", P.TpFun (tpInt, tpInt)),
             ("+", P.TpFun (P.TpTup [tpInt, tpInt], tpInt)),
	     ("-", P.TpFun (P.TpTup [tpInt, tpInt], tpInt)),
	     ("*", P.TpFun (P.TpTup [tpInt, tpInt], tpInt)),
	     ("/", P.TpFun (P.TpTup [tpInt, tpInt], tpInt)),
	     (">", P.TpFun (P.TpTup [tpInt, tpInt], tpBool)),
	     (">=", P.TpFun (P.TpTup [tpInt, tpInt], tpBool)),
	     ("<", P.TpFun (P.TpTup [tpInt, tpInt], tpBool)),
	     ("<=", P.TpFun (P.TpTup [tpInt, tpInt], tpBool)),
	     ("=", P.TpFun (P.TpTup [tpInt, tpInt], tpBool)),
	     ("<>", P.TpFun (P.TpTup [tpInt, tpInt], tpBool)),
	     ("print", P.TpFun (tpStr, tpUnit))
	    ]

	exception Undeclared

	fun lookup (G: context, x: string) =
	    case G of
		[] => raise Undeclared
	      | (x', T) :: G => if x = x' then T else lookup (G, x)

	exception IllTyped
	fun illTyped () = raise IllTyped

	fun typeOf (G: context, t: P.ttm) =
	    case t of
		P.TtmBool _ => tpBool
	      | P.TtmInt _ => tpInt
	      | P.TtmStr _ => tpStr
	      | P.TtmVar x =>
		    (lookup (G, x)
		     handle Undeclared => fatal ("Unbounded variable: " ^ x))
	      | P.TtmIf (t1, t2, t3) =>
		    let
			val T1 = typeOf (G, t1)
			val T2 = typeOf (G, t2)
			val T3 = typeOf (G, t3)
		    in
			if T1 = tpBool then
			    if T2 = T3 then T2 else illTyped ()
			else illTyped ()
		    end
	      | P.TtmOp (opname, ts) =>
		    let
			val T = lookup (Sigma, opname)
			val Ts = typesOf (G, ts)
		    in
			case T of
			    P.TpFun (P.TpTup (Ts1), T2) =>
				if Ts1 = Ts then T2 else illTyped ()
			  | P.TpFun (T1, T2) =>
				if [T1] = Ts then T2 else illTyped ()
			  | _ => illTyped ()				
		    end
	      | P.TtmLam (x, T, t) => P.TpFun (T, typeOf ((x, T) :: G, t))
	      | P.TtmApp (t1, t2) =>
		    let
			val T1 = typeOf (G, t1)
			val T2 = typeOf (G, t2)
		    in
			case T1 of
			    P.TpFun (T11, T12) =>
				if T11 = T2 then T12
				else illTyped ()
			  | _ => illTyped ()
		    end
	      | P.TtmLet (x, t1, t2) =>
		    let
			val T1 = typeOf (G, t1)
		    in
			typeOf ((x, T1) :: G, t2)
		    end
	      | P.TtmLetrec (f, T, t1, t2) =>
		    let
			val T1 = typeOf ((f, T) :: G, t1)
			val T2 = typeOf ((f, T) :: G, t2)
		    in
			if T = T1 then T2 else illTyped ()
		    end
	      | P.TtmTup (ts) => P.TpTup (typesOf (G, ts))
	      | P.TtmPro (t, i) =>
		    let
			val T = typeOf (G, t)
		    in 
			case T of
			    P.TpTup (ts) =>
				(List.nth (ts, i-1)
				 handle Subscript => illTyped ())
			  | _ => illTyped ()
		    end
	      | P.TtmFix (t) =>
		    let
			val T = typeOf (G, t)
		    in
			case T of
			    P.TpFun (T1, T2) =>
				if T1 = T2 then T1 else illTyped ()
			  | _ => illTyped ()
		    end    
	      | P.TtmAsc (t, T) =>
		    let
			val T' = typeOf (G, t)
		    in
			if T = T' then T else illTyped ()
		    end

	and typesOf (G, ts) = List.map (fn t => typeOf (G, t)) ts

	fun typeCheck (t: P.ttm): bool =
	    let
		val T = typeOf ([], t)
	    in
		true
	    end handle IllTyped => false

    end