open Abstract_syntax_tree
open Environment_domain
open Util

module Make (E : ENVIRONMENT_DOMAIN) = struct

	let neg (e, l) =
		(AST_unary(AST_NOT, (e, l))), l

	let binop op (e, l) e2 =
		(AST_binary (op, (e,l), e2)), l
	let m1 (e, l) =
		binop AST_MINUS (e, l) (AST_int_const("1", l), l)
	let p1 (e, l) =
		binop AST_PLUS (e, l) (AST_int_const("1", l), l)

	let rec condition cond env =
		begin match fst cond with
		| AST_binary (AST_LESS_EQUAL, e1, e2) ->
			E.compare env e1 e2
		| AST_binary (AST_AND, e1, e2) ->
			E.meet (condition e1 env) (condition e2 env)
		| AST_binary (AST_OR, e1, e2) ->
			E.join (condition e1 env) (condition e2 env)

		(* transformations : remove not *)
		| AST_unary (AST_NOT, (AST_unary(AST_NOT, cond), x)) ->
			condition cond env
		| AST_unary (AST_NOT, (AST_binary(AST_AND, e1, e2), x)) ->
			condition
				(AST_binary(AST_OR, neg e1, neg e2), x) env
		| AST_unary (AST_NOT, (AST_binary(AST_OR, e1, e2), x)) ->
			condition
				(AST_binary(AST_AND, neg e1, neg e2), x) env

		| AST_unary (AST_NOT, (AST_binary(op, e1, e2), x)) ->
			let op2 = match op with
			| AST_LESS_EQUAL -> AST_GREATER
			| AST_LESS -> AST_GREATER_EQUAL
			| AST_GREATER_EQUAL -> AST_LESS
			| AST_GREATER -> AST_LESS_EQUAL
			| AST_EQUAL -> AST_NOT_EQUAL
			| AST_NOT_EQUAL -> AST_EQUAL
			| _ -> assert false
			in
			condition (binop op2 e1 e2) env
		
		(* transformations : encode everything with leq *)
		| AST_binary(AST_LESS, e1, e2) ->
			condition
				(binop AST_AND (binop AST_LESS_EQUAL e1 (m1 e2))
							   (binop AST_LESS_EQUAL (p1 e1) e2))
				env
		| AST_binary (AST_GREATER_EQUAL, e1, e2) ->
			condition
				(binop AST_LESS_EQUAL e2 e1)
				env
		| AST_binary (AST_GREATER, e1, e2) ->
			condition
				(binop AST_LESS e2 e1)
				env
		| AST_binary (AST_EQUAL, e1, e2) ->
			condition
				(binop AST_AND (binop AST_LESS_EQUAL e1 e2) (binop AST_LESS_EQUAL e2 e1))
				env
		| AST_binary (AST_NOT_EQUAL, e1, e2) ->
			condition
				(binop AST_OR (binop AST_LESS e1 e2) (binop AST_LESS e2 e1))
				env

		| _ -> env
		end

	let rec interp_stmt env stat =
		begin match fst stat with
		| AST_block b ->
			(* remember to remove vars that have gone out of scope at the end *)
			let prevars = E.vars env in
			let env2 = List.fold_left interp_stmt env b in
			let postvars = E.vars env2 in
			let rmvars = List.filter (fun x -> not (List.mem x prevars)) postvars in
			List.fold_left E.rmvar env2 rmvars
		| AST_assign ((id, _), exp) ->
			E.assign env id exp
		| AST_if (cond, tb, None) ->
			E.join
				(interp_stmt (condition cond env) tb)
				(condition (neg cond) env)
		| AST_if (cond, tb, Some eb) ->
			let e1 = (interp_stmt (condition cond env) tb) in
			let e2 = (interp_stmt (condition (neg cond) env) eb) in
			E.join e1 e2
		| AST_while (cond, body) ->
			let unroll_count = 3 in
			let rec iter n i x =
				(* Format.printf "(iter %d) i:%s x:%s @." n (E.var_str i (E.vars i))
								(E.var_str x (E.vars x)); *)
				let out_state = condition (neg cond) i in
				let next_step = interp_stmt (condition cond i) body in
				(* Format.printf ". next step: %s@." (E.var_str next_step (E.vars next_step)); *)
				if n < unroll_count then
					E.join
						out_state
						(iter (n+1) next_step x)
				else
					let env2 =
						E.widen
							i
							next_step
					in
						if env2 = i
							then env2
							else E.join
								out_state
								(iter (n+1) env2 x)
			in
			condition (neg cond) (iter 0 env env)
		| AST_HALT -> E.bottom
		| AST_assert cond ->
			if not
				(E.is_bot (condition (neg cond) env))
			then begin
				Format.printf "%s: ERROR: assertion failure@."
					(Abstract_syntax_printer.string_of_extent (snd stat));
			end;
			condition cond env
		| AST_print items ->
			Format.printf "%s: %s@."
				(Abstract_syntax_printer.string_of_extent (snd stat))
				(E.var_str env (List.map fst items));
			env
		| AST_local ((ty, _), vars) ->
			List.fold_left
				(fun env ((id, _), init) ->
					let env2 = E.addvar env id in
					match init with
					| Some e -> E.assign env2 id e
					| None -> env2)
				env
				vars
		| _ -> assert false (* not implemented *)
		end
	
	let interpret prog =
		let result = List.fold_left
			(fun env x -> match x with
				| AST_stat st -> interp_stmt env st
				| _ -> env)
			E.init
			(fst prog)
		in
			Format.printf "Output: %s@."
				(E.var_str result (E.vars result))
end