(* HCoop Domtool (http://hcoop.sourceforge.net/)
 * Copyright (c) 2006, Adam Chlipala
 * Copyright (c) 2014 Clinton Ebadi <clinton@unknownlamer.org>
 *
 * 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 type-checking and reduction environments *)

structure Env :> ENV = struct

open Ast

structure SS = StringSet
structure SM = StringMap

type typeRule = exp -> bool
val typeRules : typeRule SM.map ref = ref SM.empty
fun registerType (name, rule) = typeRules := SM.insert (!typeRules, name, rule)
fun typeRule name = SM.find (!typeRules, name)

type env_vars = exp SM.map
type action = env_vars * Ast.exp list -> env_vars
val actions : action SM.map ref = ref SM.empty
fun registerAction (name, action) = actions := SM.insert (!actions, name, action)
fun action name = SM.find (!actions, name)

val containers : (action * (unit -> unit)) SM.map ref = ref SM.empty
fun registerContainer (name, befor, after) =
    containers := SM.insert (!containers, name, (befor, after))
fun container name = SM.find (!containers, name)

val functions : (exp list -> exp option) SM.map ref = ref SM.empty
fun registerFunction (name, f) =
    functions := SM.insert (!functions, name, f)
fun function name = SM.find (!functions, name)

local
    val pr = ref (fn () => ())
in

fun registerPre f =
    let
	val old = !pr
    in
	pr := (fn () => (old (); f ()))
    end
fun pre () = !pr ()

end

local
    val pst = ref (fn () => ())
in

fun registerPost f =
    let
	val old = !pst
    in
	pst := (fn () => (old (); f ()))
    end
fun post () = !pst ()

end

local
    val pr = ref (fn () => ())
in

fun registerPreTycheck f =
    let
	val old = !pr
    in
	pr := (fn () => (old (); f ()))
    end
fun preTycheck () = !pr ()

end

fun badArgs (name, args) =
    (print ("Invalid arguments to " ^ name ^ "\n");
     app (fn arg => Print.preface ("Argument: ", Print.p_exp arg)) args;
     raise Domain)
fun badArg (func, arg, v) =
    (print ("Invalid " ^ arg ^ " argument to " ^ func ^ "\n");
     Print.preface ("Argument: ", Print.p_exp v);
     raise Domain)

type 'a arg = exp -> 'a option

fun int (EInt n, _) = SOME n
  | int _ = NONE

fun string (EString s, _) = SOME s
  | string _ = NONE

fun bool (EVar "false", _) = SOME false
  | bool (EVar "true", _) = SOME true
  | bool _ = NONE

fun mapFail f [] = SOME []
  | mapFail f (h :: t) =
    case f h of
	NONE => NONE
      | SOME h' =>
	case mapFail f t of
	    NONE => NONE
	  | SOME t' => SOME (h' :: t')

fun list f (EList ls, _) = mapFail f ls
  | list _ _ = NONE

fun none func f (_, []) = (f ();
			   SM.empty)
  | none func _ (_, es) = badArgs (func, es)

fun one func (name, arg) f (_, [e]) =
    (case arg e of
	 NONE => badArg (func, name, e)
       | SOME v => (f v;
		    SM.empty))
  | one func _ _ (_, es) = badArgs (func, es)

fun two func (name1, arg1, name2, arg2) f (_, [e1, e2]) =
    (case (arg1 e1, arg2 e2) of
	 (NONE, _) => badArg (func, name1, e1)
       | (_, NONE) => badArg (func, name2, e2)
       | (SOME v1, SOME v2) => (f (v1, v2);
				SM.empty))
  | two func _ _ (_, es) = badArgs (func, es)

fun three func (name1, arg1, name2, arg2, name3, arg3) f (_, [e1, e2, e3]) =
    (case (arg1 e1, arg2 e2, arg3 e3) of
	 (NONE, _, _) => badArg (func, name1, e1)
       | (_, NONE, _) => badArg (func, name2, e2)
       | (_, _, NONE) => badArg (func, name3, e3)
       | (SOME v1, SOME v2, SOME v3) => (f (v1, v2, v3);
					 SM.empty))
  | three func _ _ (_, es) = badArgs (func, es)

fun four func (name1, arg1, name2, arg2, name3, arg3, name4, arg4) f (_, [e1, e2, e3, e4]) =
    (case (arg1 e1, arg2 e2, arg3 e3, arg4 e4) of
	 (NONE, _, _, _) => badArg (func, name1, e1)
       | (_, NONE, _, _) => badArg (func, name2, e2)
       | (_, _, NONE, _) => badArg (func, name3, e3)
       | (_, _, _, NONE) => badArg (func, name4, e4)
       | (SOME v1, SOME v2, SOME v3, SOME v4) => (f (v1, v2, v3, v4);
						  SM.empty))
  | four func _ _ (_, es) = badArgs (func, es)

fun noneV func f (evs, []) = (f evs;
			      SM.empty)
  | noneV func _ (_, es) = badArgs (func, es)

fun oneV func (name, arg) f (evs, [e]) =
    (case arg e of
	 NONE => badArg (func, name, e)
       | SOME v => (f (evs, v);
		    SM.empty))
  | oneV func _ _ (_, es) = badArgs (func, es)

fun twoV func (name1, arg1, name2, arg2) f (evs, [e1, e2]) =
    (case (arg1 e1, arg2 e2) of
	 (NONE, _) => badArg (func, name1, e1)
       | (_, NONE) => badArg (func, name2, e2)
       | (SOME v1, SOME v2) => (f (evs, v1, v2);
				SM.empty))
  | twoV func _ _ (_, es) = badArgs (func, es)


fun env arg (evs, name) =
    case SM.find (evs, name) of
	NONE => raise Fail ("Unavailable environment variable " ^ name)
      | SOME e =>
	case arg e of
	    NONE => (Print.preface ("Unexpected value for " ^ name ^ ":",
				    Print.p_exp e);
		     raise Fail ("Bad format for environment variable " ^ name))
	  | SOME v => v

fun type_one func arg f =
    registerType (func, fn e =>
			   case arg e of
			       NONE => false
			     | SOME v => f v)

fun action_none name f = registerAction (name, none name f)
fun action_one name args f = registerAction (name, one name args f)
fun action_two name args f = registerAction (name, two name args f)
fun action_three name args f = registerAction (name, three name args f)
fun action_four name args f = registerAction (name, four name args f)

fun actionV_none name f = registerAction (name, fn (env, _) => (f env; env))
fun actionV_one name args f = registerAction (name, oneV name args f)
fun actionV_two name args f = registerAction (name, twoV name args f)

fun container_none name (f, g) = registerContainer (name, none name f, g)
fun container_one name args (f, g) = registerContainer (name, one name args f, g)

fun containerV_none name (f, g) = registerContainer (name, noneV name f, g)
fun containerV_one name args (f, g) = registerContainer (name, oneV name args f, g)

datatype env = Env of SS.set * (typ * exp option * env) SM.map * SS.set * (typ * exp * env) SM.map
val empty : env = Env (SS.empty, SM.empty, SS.empty, SM.empty)


fun initialDynEnvTypes (Env (_, _, _, ds)) =
    SM.map (fn (t, _, _) => t) ds

(* hack ahead: These are not reduced when declared and must be before
   starting evaluation. Pass in reduceExp, and force an awkward
   calling convention so no one thinks this is the Right Way (tm) *)
fun initialDynEnvVals f (Env (_, _, _, ds)) =
    SM.map (fn (t, v, env) => f env v) ds

fun lookupType (Env (ts, _, _, _)) name = SS.member (ts, name)
fun lookupVal (Env (_, vs, _, _)) name =
    case SM.find (vs, name) of
	NONE => NONE
      | SOME (t, _, _) => SOME t
fun lookupEquation (Env (_, vs, _, _)) name =
    case SM.find (vs, name) of
	NONE => NONE
      | SOME (_, NONE, _) => NONE
      | SOME (_, SOME eq, env) => SOME (eq, env)
fun lookupContext (Env (_, _, cs, _)) name = SS.member (cs, name)
fun lookupInitialDynEnvVal (Env (_, _, _, ds)) name =
    case SM.find (ds, name) of
	NONE => NONE
      | SOME (t, _, _) => SOME t

fun bindType (Env (ts, vs, cs, ds)) name = Env (SS.add (ts, name), vs, cs, ds)
fun bindVal (env as (Env (ts, vs, cs, ds))) (name, t, eqo) = Env (ts, SM.insert (vs, name, (t, eqo, env)), cs, ds)
fun bindContext (Env (ts, vs, cs, ds)) name = Env (ts, vs, SS.add (cs, name), ds)
fun bindInitialDynEnvVal (env as (Env (ts, vs, cs, ds))) (name, t, eqn) = Env (ts, vs, cs, SM.insert (ds, name, (t, eqn, env)))

fun types (Env (ts, _, _, _)) = ts
fun vals (Env (_, vs, _, _)) = SM.foldli (fn (name, _, vs) => SS.add (vs, name)) SS.empty vs
fun contexts (Env (_, _, cs, _)) = cs
fun dynamics (Env (_, _, _, ds)) = SM.foldli (fn (name, _, ds) => SS.add (ds, name)) SS.empty ds

end