(* HCoop Domtool (http://hcoop.sourceforge.net/)
 * Copyright (c) 2006, Adam Chlipala
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 *)

(* Domtool configuration language type checking *)

structure Tycheck :> TYCHECK = struct

open Ast Print Env

structure SM = StringMap

val externFlag = ref false
fun allowExterns () = externFlag := true
fun disallowExterns () = externFlag := false

local
    val unifCount = ref 0
in
fun resetUnif () = unifCount := 0

fun newUnif () =
    let
	val c = !unifCount
	val name =
	    if c < 26 then
		str (chr (ord #"A" + c))
	    else
		"UNIF" ^ Int.toString (c - 26)
    in
	unifCount := c + 1;
	TUnif (name, ref NONE)
    end
end
	    
exception UnequalDomains
	  
fun eqRecord f (r1, r2) =
    (SM.appi (fn (k, v1) =>
		 case SM.find (r2, k) of
		     NONE => raise UnequalDomains
		   | SOME v2 =>
		     if f (v1, v2) then
			 ()
		     else
			 raise UnequalDomains) r1;
     SM.appi (fn (k, v2) =>
		 case SM.find (r1, k) of
		     NONE => raise UnequalDomains
		   | SOME v1 =>
		     if f (v1, v2) then
			 ()
		     else
			 raise UnequalDomains) r2;
     true)
    handle UnequalDomains => false

fun eqPred ((p1, _), (p2, _)) =
    case (p1, p2) of
	(CRoot, CRoot) => true
      | (CConst s1, CConst s2) => s1 = s2
      | (CPrefix p1, CPrefix p2) => eqPred (p1, p2)
      | (CNot p1, CNot p2) => eqPred (p1, p2)
      | (CAnd (p1, q1), CAnd (p2, q2)) =>
	eqPred (p1, p2) andalso eqPred (q1, q2)

      | _ => false

fun eqTy (t1All as (t1, _), t2All as (t2, _)) =
    case (t1, t2) of
	(TBase s1, TBase s2) => s1 = s2
      | (TList t1, TList t2) => eqTy (t1, t2)
      | (TArrow (d1, r1), TArrow (d2, r2)) =>
	eqTy (d1, d2) andalso eqTy (r1, r2)

      | (TAction (p1, d1, r1), TAction (p2, d2, r2)) =>
	eqPred (p1, p2) andalso eqRecord eqTy (d1, d2)
	andalso eqRecord eqTy (r1, r2)

      | (TNested (p1, q1), TNested (p2, q2)) =>
	eqPred (p1, p2) andalso eqTy (q1, q2)

      | (TUnif (_, ref (SOME t1)), _) => eqTy (t1, t2All)
      | (_, TUnif (_, ref (SOME t2))) => eqTy (t1All, t2)

      | (TUnif (_, r1), TUnif (_, r2)) => r1 = r2

      | (TError, TError) => true

      | _ => false

datatype unification_error =
	 UnifyPred of pred * pred
       | UnifyTyp of typ * typ
       | UnifyOccurs of string * typ

exception Unify of unification_error

datatype type_error =
	 WrongType of string * exp * typ * typ * unification_error option
       | WrongForm of string * string * exp * typ * unification_error option
       | UnboundVariable of string
       | WrongPred of string * pred * pred

fun describe_unification_error t ue =
    case ue of
	UnifyPred (p1, p2) =>
	(print "Reason: Incompatible contexts.\n";
	 preface ("Have:", p_pred p1);
	 preface ("Need:", p_pred p2))
      | UnifyTyp (t1, t2) =>
	if eqTy (t, t1) then
	    ()
	else
	    (print "Reason: Incompatible types.\n";
	     preface ("Have:", p_typ t1);
	     preface ("Need:", p_typ t2))
      | UnifyOccurs (name, t') =>
	if eqTy (t, t') then
	    ()
	else
	    (print "Reason: Occurs check failed for ";
	     print name;
	     print " in:\n";
	     printd (p_typ t))

fun describe_type_error loc te =
    case te of
	WrongType (place, e, t1, t2, ueo) =>
	(ErrorMsg.error (SOME loc) (place ^ " has wrong type.");
	 preface (" Expression:", p_exp e);
	 preface ("Actual type:", p_typ t1);
	 preface ("Needed type:", p_typ t2);
	 Option.app (describe_unification_error t1) ueo)
      |	WrongForm (place, form, e, t, ueo) =>
	(ErrorMsg.error (SOME loc) (place ^ " has a non-" ^ form ^ " type.");
	 preface ("Expression:", p_exp e);
	 preface ("      Type:", p_typ t);
	 Option.app (describe_unification_error t) ueo)
      | UnboundVariable name =>
	ErrorMsg.error (SOME loc) ("Unbound variable " ^ name ^ ".\n")
      |	WrongPred (place, p1, p2) =>
	(ErrorMsg.error (SOME loc) ("Context incompatibility for " ^ place ^ ".");
	 preface ("Have:", p_pred p1);
	 preface ("Need:", p_pred p2))

fun predImplies (p1All as (p1, _), p2All as (p2, _)) =
    case (p1, p2) of
	(_, CAnd (p1, p2)) => predImplies (p1All, p1) andalso predImplies (p1All, p2)
      | (CAnd (p1, p2), _) => predImplies (p1, p2All) orelse predImplies (p2, p2All)

      | (_, CPrefix (CRoot, _)) => true
      | (CNot (CPrefix (CRoot, _), _), _) => true

      | (CRoot, CRoot) => true

      | (CConst s1, CConst s2) => s1 = s2

      | (CPrefix p1, CPrefix p2) => predImplies (p1, p2)
      | (_, CPrefix p2) => predImplies (p1All, p2)

      | (CNot p1, CNot p2) => predImplies (p2, p1)
      | (CRoot, CNot (CConst _, _)) => true
      | (CConst s1, CNot (CConst s2, _)) => s1 <> s2

      | _ => false

fun predSimpl (pAll as (p, loc)) =
    case p of
	CRoot => pAll
      | CConst _ => pAll
      | CPrefix p => (CPrefix (predSimpl p), loc)
      | CNot p => (CNot (predSimpl p), loc)
      | CAnd (p1, p2) =>
	let
	    val p1' = predSimpl p1
	    val p2' = predSimpl p2
	in
	    case p1' of
		(CAnd (c1, c2), _) => predSimpl (CAnd (c1, (CAnd (c2, p2'), loc)), loc)
	      | _ => if predImplies (p2', p1') then
			 p2'
		     else if predImplies (p1', p2') then
			 p1'
		     else
			 (CAnd (p1', p2'), loc)
	end
   
fun subPred (p1, p2) =
    if predImplies (p1, p2) then
	()
    else
	raise (Unify (UnifyPred (p1, p2)))

fun subRecord f (r1, r2) =
    SM.appi (fn (k, v2) =>
		case SM.find (r1, k) of
		    NONE => raise UnequalDomains
		  | SOME v1 => f (v1, v2)) r2

fun occurs u (t, _) =
    case t of
	TBase _ => false
      | TList t => occurs u t
      | TArrow (d, r) => occurs u d orelse occurs u r
      | TAction (_, d, r) =>
	List.exists (occurs u) (SM.listItems d)
	orelse List.exists (occurs u) (SM.listItems r)
      | TNested (_, t) => occurs u t
      | TError => false
      | TUnif (_, ref (SOME t)) => occurs u t
      | TUnif (_, u') => u = u'

fun subTyp (t1All as (t1, _), t2All as (t2, _)) =
    case (t1, t2) of
	(TBase s1, TBase s2) =>
	if s1 = s2 then
	    ()
	else
	    raise Unify (UnifyTyp (t1All, t2All))
      | (TList t1, TList t2) => subTyp (t1, t2)
      | (TArrow (d1, r1), TArrow (d2, r2)) =>
	(subTyp (d2, d1);
	 subTyp (r1, r2))

      | (TAction (p1, d1, r1), TAction (p2, d2, r2)) =>
	((subPred (p2, p1);
	  subRecord subTyp (d2, d1);
	  subRecord subTyp (r1, r2);
	  subRecord subTyp (r2, r1))
	 handle UnequalDomains => raise Unify (UnifyTyp (t1All, t2All)))

      | (TNested (d1, r1), TNested (d2, r2)) =>
	(subPred (d2, d1);
	 subTyp (r1, r2))

      | (TUnif (_, ref (SOME t1)), _) => subTyp (t1, t2All)
      | (_, TUnif (_, ref (SOME t2))) => subTyp (t1All, t2)

      | (TUnif (_, r1), TUnif (_, r2)) =>
	if r1 = r2 then
	    ()
	else
	    r1 := SOME t2All

      | (TUnif (name, r), _) =>
	if occurs r t2All then
	    raise (Unify (UnifyOccurs (name, t2All)))
	else
	    r := SOME t2All

      | (_, TUnif (name, r)) =>
	if occurs r t1All then
	    raise (Unify (UnifyOccurs (name, t1All)))
	else
	    r := SOME t1All

      | (TError, _) => ()
      | (_, TError) => ()

      | _ => raise Unify (UnifyTyp (t1All, t2All))

fun isError t =
    case t of
	(TError, _) => true
      | _ => false

fun whnorm (tAll as (t, loc)) =
    case t of
	TUnif (_, ref (SOME tAll)) => whnorm tAll
      | _ => tAll

fun baseCondition t =
    case whnorm t of
	(TBase name, _) => typeRule name
      | (TList t, _) =>
	(case baseCondition t of
	     NONE => NONE
	   | SOME f => SOME (fn (EList ls, _) => List.all f ls
			      | _ => false))
      | _ => NONE

fun hasTyp (e, t1, t2) =
    if (case baseCondition t2 of
	    NONE => false
	  | SOME rule => rule e) then
	()
    else
	subTyp (t1, t2)

fun checkPred G (p, loc) =
    let
	val err = ErrorMsg.error (SOME loc)
    in
	case p of
	    CRoot => ()
	  | CConst s =>
	    if lookupContext G s then
		()
	    else
		err ("Unbound context " ^ s)
	  | CPrefix p => checkPred G p
	  | CNot p => checkPred G p
	  | CAnd (p1, p2) => (checkPred G p1; checkPred G p2)
    end

fun checkTyp G (tAll as (t, loc)) =
    let
	val err = ErrorMsg.error (SOME loc)
    in
	case t of
	    TBase name =>
	    if lookupType G name then
		tAll
	    else
		(err ("Unbound type name " ^ name);
		 (TError, loc))
	  | TList t => (TList (checkTyp G t), loc)
	  | TArrow (d, r) => (TArrow (checkTyp G d, checkTyp G r), loc)
	  | TAction (p, d, r) => (checkPred G p;
				  (TAction (p, SM.map (checkTyp G) d,
					    SM.map (checkTyp G) r), loc))
	  | TNested (p, t) => (checkPred G p;
			       (TNested (p, checkTyp G t), loc))
	  | TError => raise Fail "TError in parser-generated type"
	  | TUnif _ => raise Fail "TUnif in parser-generated type"
    end

fun envVarSetFrom v (e, _) =
    case e of
	ESet (v', e) =>
	if v = v' then
	    SOME e
	else
	    NONE
      | EGet (_, _, e) => envVarSetFrom v e
      | ESeq es => foldr (fn (e, found) =>
			     case found of
				 SOME _ => found
			       | NONE => envVarSetFrom v e)
			 NONE es
      | ELocal (_, e) => envVarSetFrom v e

      | _ => NONE

fun checkExp G (eAll as (e, loc)) =
    let
	val dte = describe_type_error loc
    in
	case e of
	    EInt _ => (TBase "int", loc)
	  | EString _ => (TBase "string", loc)
	  | EList es =>
	    let
		val t = (newUnif (), loc)
	    in
		foldl (fn (e', ret) =>
			  let
			      val t' = checkExp G e'
			  in
			      (hasTyp (eAll, t', t);
			       if isError t' then
				   (TList (TError, loc), loc)
			       else
				   ret)
			      handle Unify ue =>
				     (dte (WrongType ("List element",
						      e',
						      t',
						      t,
						      SOME ue));
				      (TError, loc))
			  end) (TList t, loc) es
	    end

	  | ELam (x, to, e) =>
	    let
		val t =
		    case to of
			NONE => (newUnif (), loc)
		      | SOME t => checkTyp G t

		val G' = bindVal G (x, t, NONE)
		val t' = checkExp G' e
	    in
		(TArrow (t, t'), loc)
	    end
	  | EVar x =>
	    (case lookupVal G x of
		 NONE => (dte (UnboundVariable x);
			  (TError, loc))
	       | SOME t => t)
	  | EApp (func, arg) =>
	    let
		val dom = (newUnif (), loc)
		val ran = (newUnif (), loc)

		val tf = checkExp G func
		val ta = checkExp G arg
	    in
		(hasTyp (func, tf, (TArrow (dom, ran), loc));
		 hasTyp (arg, ta, dom)
		 handle Unify ue =>
			dte (WrongType ("Function argument",
					arg,
					ta,
					dom,
					SOME ue));
		 ran)
		handle Unify ue =>
		       (dte (WrongForm ("Function to be applied",
					"function",
					func,
					tf,
					SOME ue));
			(TError, loc))
	    end

	  | EALam (x, p, e) =>
	    let
		val p' = checkPred G p

		val G' = bindVal G (x, (TAction (p, SM.empty, SM.empty), loc), NONE)
		val t' = whnorm (checkExp G' e)
	    in
		case t' of
		    (TAction _, _) => (TNested (p, t'), loc)
		  | _ => (dte (WrongForm ("Body of nested configuration 'fn'",
					  "action",
					  e,
					  t',
					  NONE));
			  (TError, loc))
	    end

	  | ESet (evar, e) =>
	    let
		val t = checkExp G e
	    in
		(TAction ((CPrefix (CRoot, loc), loc),
			  SM.empty,
			  SM.insert (SM.empty, evar, t)),
		 loc)
	    end
	  | EGet (x, evar, rest) =>
	    let
		val xt = (newUnif (), loc)
		val G' = bindVal G (x, xt, NONE)

		val rt = whnorm (checkExp G' rest)
	    in
		case rt of
		    (TAction (p, d, r), _) =>
		    (case SM.find (d, evar) of
			 NONE => (TAction (p, SM.insert (d, evar, xt), r), loc)
		       | SOME xt' =>
			 (subTyp (xt', xt)
				handle Unify ue =>
				       dte (WrongType ("Retrieved environment variable",
						       (EVar x, loc),
						       xt',
						       xt,
						       SOME ue));
			  rt))
		  | (TError, _) => rt
		  | _ => (dte (WrongForm ("Body of environment variable read",
					  "action",
					  rest,
					  rt,
					  NONE));
			  (TError, loc))
	    end

	  | ESeq [] => raise Fail "Empty ESeq"
	  | ESeq [e1] => checkExp G e1
	  | ESeq (e1 :: rest) =>
	    let
		val e2 = (ESeq rest, loc)

		val t1 = whnorm (checkExp G e1)
		val t2 = whnorm (checkExp G e2)
	    in
		case t1 of
		    (TAction (p1, d1, r1), _) =>
		    (case t2 of
			 (TAction (p2, d2, r2), _) =>
			 let
			     val p' = predSimpl (CAnd (p1, p2), loc)

			     val d' = SM.foldli (fn (name, t, d') =>
						    case SM.find (r1, name) of
							NONE =>
							(case SM.find (d', name) of
							     NONE => SM.insert (d', name, t)
							   | SOME t' =>
							     ((case envVarSetFrom name e1 of
								   NONE => subTyp (t, t')
								 | SOME e => hasTyp (e, t, t'))
							      handle Unify ue =>
								     dte (WrongType ("Shared environment variable",
										     (EVar name, loc),
										     t',
										     t,
										     SOME ue));
							      d'))
						      | SOME t' =>
							((case envVarSetFrom name e1 of
								   NONE => subTyp (t, t')
								 | SOME e => hasTyp (e, t, t'))
							 handle Unify ue =>
								dte (WrongType ("Shared environment variable",
										(EVar name, loc),
										t',
										t,
										SOME ue));
							 d'))
						d1 d2

			     val r' = SM.foldli (fn (name, t, r') => SM.insert (r', name, t))
						r1 r2
			 in
			     (TAction (p', d', r'), loc)
			 end
		       | (TError, _) => t2
		       | _ => (dte (WrongForm ("Action to be sequenced",
					       "action",
					       e2,
					       t2,
					       NONE));
			       (TError, loc)))
		  | (TError, _) => t1
		  | _ => (dte (WrongForm ("Action to be sequenced",
					  "action",
					  e1,
					  t1,
					  NONE));
			  (TError, loc))
	    end

	  | ELocal (e1, e2) =>
	    let
		val t1 = whnorm (checkExp G e1)
		val t2 = whnorm (checkExp G e2)
	    in
		case t1 of
		    (TAction (p1, d1, r1), _) =>
		    (case t2 of
			 (TAction (p2, d2, r2), _) =>
			 let
			     val p' = predSimpl (CAnd (p1, p2), loc)

			     val d' = SM.foldli (fn (name, t, d') =>
						    case SM.find (r1, name) of
							NONE =>
							(case SM.find (d', name) of
							     NONE => SM.insert (d', name, t)
							   | SOME t' =>
							     ((case envVarSetFrom name e1 of
								   NONE => subTyp (t', t)
								 | SOME e => hasTyp (e, t', t))
							      handle Unify ue =>
								     dte (WrongType ("Shared environment variable",
										     (EVar name, loc),
										     t',
										     t,
										     SOME ue));
							      d'))
						      | SOME t' =>
							((case envVarSetFrom name e1 of
								   NONE => subTyp (t', t)
								 | SOME e => hasTyp (e, t', t))
							 handle Unify ue =>
								dte (WrongType ("Shared environment variable",
										(EVar name, loc),
										t',
										t,
										SOME ue));
							 d'))
						d1 d2
			 in
			     (TAction (p', d', r2), loc)
			 end
		       | (TError, _) => t2
		       | _ => (dte (WrongForm ("Action to be sequenced",
					       "action",
					       e2,
					       t2,
					       NONE));
			       (TError, loc)))
		  | (TError, _) => t1
		  | _ => (dte (WrongForm ("Action to be sequenced",
					  "action",
					  e1,
					  t1,
					  NONE));
			  (TError, loc))
	    end


	  | EWith (e1, e2) =>
	    let
		val t1 = whnorm (checkExp G e1)
		val t2 = whnorm (checkExp G e2)
	    in
		case t1 of
		    (TNested (pd, (TAction (pr, d1, r1), _)), _) =>
		    (case t2 of
			 (TAction (p, d, r), _) =>
			 if predImplies (pd, p) then
			     let
				 val combineRecs =
				     SM.unionWithi (fn (name, t1, t2) =>
						       (subTyp (t1, t2)
							handle Unify ue =>
							       dte (WrongType ("Environment variable",
									       (EVar name, loc),
									       t1,
									       t2,
									       SOME ue));
							t2))
			     in
				 (TAction (pr, combineRecs (d, d1),
					   combineRecs (r, r1)), loc)
			     end
			 else
			     (dte (WrongPred ("nested action",
					      pd,
					      p));
			      (TError, loc))
		       | (TError, _) => t2
		       | _ =>
			 (dte (WrongForm ("Body of nested action",
					  "action",
					  e2,
					  t2,
					  NONE));
			  (TError, loc)))
		  | (TError, _) => t1
		  | _ =>
		    (dte (WrongForm ("Container of nested action",
				     "action",
				     e1,
				     t1,
				     NONE));
		     (TError, loc))
	    end

	  | ESkip => (TAction ((CPrefix (CRoot, loc), loc),
			       SM.empty, SM.empty), loc)
    end

exception Ununif

fun ununif (tAll as (t, loc)) =
    case t of
	TBase _ => tAll
      | TList t => (TList (ununif t), loc)
      | TArrow (d, r) => (TArrow (ununif d, ununif r), loc)
      | TAction (p, d, r) => (TAction (p, SM.map ununif d, SM.map ununif r), loc)
      | TUnif (_, ref (SOME t)) => ununif t
      | TNested _ => tAll
      | TError => tAll

      | TUnif (_, ref NONE) => raise Ununif

fun hasError (t, _) =
    case t of
	TBase _ => false
      | TList t => hasError t
      | TArrow (d, r) => hasError d orelse hasError r
      | TAction (p, d, r) => List.exists hasError (SM.listItems d)
			     orelse List.exists hasError (SM.listItems r)
      | TNested _ => false
      | TError => false
      | TUnif (_, ref (SOME t)) => hasError t
      | TUnif (_, ref NONE) => false


fun checkUnit G (eAll as (_, loc)) =
    let
	val _ = resetUnif ()
	val t = checkExp G eAll
    in
	if hasError t then
	    t
	else
	    ununif t
	    handle Ununif =>
		   (ErrorMsg.error (SOME loc) "Unification variables remain in type:";
		    printd (p_typ t);
		    t)
    end

fun checkDecl G (d, _, loc) =
    case d of
	DExternType name =>
	if !externFlag then
	    bindType G name
	else
	    (ErrorMsg.error (SOME loc) "'extern type' not allowed in untrusted code";
	     G)
      | DExternVal (name, t) =>
	if !externFlag then
	    bindVal G (name, checkTyp G t, NONE)
	else
	    (ErrorMsg.error (SOME loc) "'extern val' not allowed in untrusted code";
	     G)
      | DVal (name, to, e) =>
	let
	    val to =
		case to of
		    NONE => (newUnif (), loc)
		  | SOME to => checkTyp G to

	    val t = checkExp G e
	in
	    hasTyp (e, t, to)
	    handle Unify ue =>
		   describe_type_error loc
				       (WrongType ("Bound value",
						   e,
						   t,
						   to,
						   SOME ue));
	    bindVal G (name, to, SOME e)
	end
      | DContext name => bindContext G name

fun checkFile G tInit (_, ds, eo) =
    let
	val G' = foldl (fn (d, G) => checkDecl G d) G ds
    in
	case eo of
	    NONE => ()
	  | SOME (e as (_, loc)) =>
	    let
		val t = checkExp G' e
	    in
		hasTyp (e, t, tInit)
		handle Unify ue =>
		       (ErrorMsg.error (SOME loc) "Bad type for final expression of source file.";
			preface ("Actual:", p_typ t);
			preface ("Needed:", p_typ tInit))
	    end;
	G'
    end

end