open Abstract_syntax_tree
open Util

open Value_domain
open Environment_domain


module NonRelational (V : VALUE_DOMAIN) : ENVIRONMENT_DOMAIN = struct
	type env = V.t VarMap.t

	type t = Env of env | Bot

	let init = Env VarMap.empty
	let bottom = Bot

	let get_var env id =
		try VarMap.find id env
		with Not_found -> V.top
	
	(* utilities *)

	let rec eval env expr =
		begin match fst expr with
		| AST_identifier (id, _) -> get_var env id
		| AST_int_const (s, _) -> V.const (Z.of_string s)
		| AST_int_rand ((s, _), (t, _)) -> V.rand (Z.of_string s) (Z.of_string t)
		| AST_unary (AST_UNARY_PLUS, e) -> eval env e
		| AST_unary (AST_UNARY_MINUS, e) -> V.neg (eval env e)
		| AST_unary (_, e) -> V.bottom
		| AST_binary (op, e1, e2) ->
			let f = match op with
			| AST_PLUS -> V.add
			| AST_MINUS -> V.sub
			| AST_MULTIPLY -> V.mul
			| AST_DIVIDE -> V.div
			| AST_MODULO -> V.rem
			| _ -> fun _ _ -> V.bottom
			in f (eval env e1) (eval env e2)
		| _ -> assert false (* unimplemented extension *)
		end

	let strict env =					(*		env -> t		*)
		if VarMap.exists (fun _ x -> x = V.bottom) env
			then Bot
			else Env env
	let strict_apply f = function		(* (env -> t) -> t -> t *)
		| Bot -> Bot
		| Env e -> match f e with
			| Bot -> Bot
			| Env e -> strict e

	(* implementation *)

	let is_bot env = match env with
		| Bot -> true
		| Env env -> strict env = Bot

	let addvar x id = strict_apply (fun y -> Env (VarMap.add id V.top y)) x
	let rmvar x id = strict_apply (fun y -> Env (VarMap.remove id y)) x

	let assign x id expr = strict_apply
		(fun env -> Env (VarMap.add id (eval env expr) env))
		x
	let compare x e1 e2 =
		strict_apply (fun env -> match fst e1, fst e2 with
				| AST_identifier(a, _), AST_identifier(b, _) ->
					let v1, v2 = get_var env a, get_var env b in
					let v1', v2' = V.leq v1 v2 in
					Env (VarMap.add a v1' (VarMap.add b v2' env))
				| AST_identifier(a, _), _ ->
					let v1, v2 = get_var env a, eval env e2 in
					let v1', v2' = V.leq v1 v2 in
					if V.bottom = v2'
						then Bot
						else Env (VarMap.add a v1' env)
				| _, AST_identifier(b, _) ->
					let v1, v2 = eval env e1, get_var env b in
					let v1', v2' = V.leq v1 v2 in
					if V.bottom = v1'
						then Bot
						else Env (VarMap.add b v2' env)
				| _ ->
					let v1, v2 = eval env e1, eval env e2 in
					let v1', v2' = V.leq v1 v2 in
					if V.bottom = v1' ||V.bottom = v2'
						then Bot
						else Env env)
			x

	let join a b = match a, b with
		| Bot, x | x, Bot -> x
		| Env m, Env n ->
			strict (VarMap.map2z (fun _ a b -> V.join a b) m n)
	
	let meet a b = match a, b with
		| Bot, _ | _, Bot -> Bot
		| Env m, Env n ->
			strict (VarMap.map2z (fun _ a b -> V.meet a b) m n)

	let widen a b = match a, b with
		| Bot, x | x, Bot -> x
		| Env m, Env n ->
			strict (VarMap.map2z (fun _ a b -> V.widen a b) m n)
	
	let subset a b = match a, b with
		| Bot, x -> true
		| Env _, Bot -> false
		| Env m, Env n ->
			VarMap.for_all2o
				(fun _ _ -> true)
				(fun _ _ -> true)
				(fun _ a b -> V.subset a b)
				m n
		
	(* pretty-printing *)
	let var_str env vars =
		match env with
		| Bot -> "bottom"
		| Env env ->
			let v = List.fold_left
				(fun s id ->
					(if s = "" then s else s ^ ", ") ^
						id ^ " in " ^ (V.to_string (get_var env id))
					)
				""
				vars
			in
				"[ " ^ v ^ " ]"
	
	let vars env = match env with
		| Bot -> []
		| Env env -> List.map fst (VarMap.bindings env)
end