open Absyn

(* Sets and dictionaries in the library are implemented using balanced binary 
 * trees. Anything to be stored in such sets or used as keys in such dictionaries
 * needs to have a comparison function associated with it. *)

module VarType = struct
  type t = var
  let compare = Pervasives.compare 
end 

(* The type of sets of variables. 
 * See http://caml.inria.fr/ocaml/htmlman/libref/Set.S.html 
 * for the operations provided on sets *)

module VarSet = Set.Make (VarType)

type set = VarSet.t

(* This exception should be raised in your interpreter of 
 * part 5 when an expression cannot be evaluated. *)
exception RuntimeError of string

(* This exception should not appear in the code that you submit.  It
 *	is only used now so that you can run and test the program before
 *	completing the assignment. *)
exception TODO 

(* A function to print all of the vars contained in a set, with 
 * spaces between each variable. *)
let print_vars (vs:set) : unit = 
  let es = VarSet.elements vs in 
  Format.open_hbox();
  Absyn.print_list es Format.print_string Format.print_space;
  Format.close_box();
  Format.print_newline()


(*********************************************************************)
(* This is the part you need to write *)

(* 1 *)

let free_vars (exp:exp) : set = 
  let rec helper (bound:set) (exp:exp)  : set = 
	 match exp with 
		Var v -> (if VarSet.mem v bound
					 then VarSet.empty
					 else VarSet.singleton v)
	 | Lambda (v,e) -> 
		  helper (VarSet.add v bound) e
	 | App (e1,e2) -> VarSet.union (helper bound e1) (helper bound e2)
	 | Num i -> VarSet.empty
	 | Prim (p, exps) -> 
		  List.fold_left VarSet.union VarSet.empty 
		  (List.map (helper bound) exps)
	 | Let (v, e1, e2) -> 
		  let new_bound = VarSet.add v bound in
			 VarSet.union (helper bound e1) (helper new_bound e2)
  in 
	 helper VarSet.empty exp


(** Alternative 1. List variables are those that occur in the binding
  positions of terms such as lambda and let. For example, 
  bound_vars (lambda (x) (lambda (z) (z y))) returns "x z"
*)

let rec bound_vars (exp:exp) : set = 
  match exp with 
		Var v -> VarSet.empty
	 | Lambda (v,e) -> VarSet.add v (bound_vars e) 
	 | App (e1,e2) -> VarSet.union (bound_vars e1) (bound_vars e2)
	 | Num i -> VarSet.empty
	 | Prim (p, exps) -> 
		  List.fold_left 
		  VarSet.union  VarSet.empty (List.map bound_vars exps)
	 | Let (v,e1, e2) -> 
		  VarSet.union 
		  (* variables from the bindings in the let *)
		  (VarSet.add v (bound_vars e1))
		  (* Variables from the body of the let *)
		  (bound_vars e2)

(** Alternative 2. List variables that occur bound in an expression.
  For example, 
  bound_vars (lambda (x) (lambda (z) (z y))) returns "z"
 *)

let rec bound_vars (exp:exp) : set = 
  match exp with 
		Var v -> VarSet.empty
	 | Lambda (v,e) ->
		  let set = bound_vars e in 
		  if (VarSet.mem v (free_vars e)) then 
			 VarSet.add v set
		  else set
	 | App (e1,e2) -> VarSet.union (bound_vars e1) (bound_vars e2)
	 | Num i -> VarSet.empty
	 | Prim (p, exps) -> 
		  List.fold_left 
		  VarSet.union  VarSet.empty (List.map bound_vars exps)
	 | Let (v,e1, e2) -> 
		  let set = VarSet.union (bound_vars e1) (bound_vars e2) in  
			 if (VarSet.mem v (free_vars e2)) 
			 then VarSet.add v set 
			 else set
				
(* 2 *)

let is_closed (exp:exp) :bool = VarSet.is_empty (free_vars exp)

(* 3 *)

let rec subst  (subst_exp:exp) (subst_var:var) (exp:exp) : exp = 
  let go = subst subst_exp subst_var in 
	 match exp with 
		  Var v -> if v = subst_var then subst_exp else exp
		| Lambda (v,e) ->
			 if v = subst_var then 
				exp 
			 else 
				Lambda (v, go e)
		| App (e1,e2) -> 
			 App (go e1, go e2)
		| Num i -> exp
		| Prim (p, exps) -> 
			 Prim (p, List.map go exps)
		| Let (v, e1, e2) -> 
			 if v = subst_var 
			 then 
				Let (v, go e1,e2)
			 else 
				Let (v, go e1, go e2)	 
(* 4 *)
	
(* Alternative 1:  large step evaluation that we talked about in class. *)
				
let rec eval (exp:exp) : exp = 
	 match exp with 
		  Var v -> 
			 (* You can do whatever you want here, it's guaranteed not 
				 to happen *)
			 raise (RuntimeError ("Unbound variable " ^ v))
		| Lambda (v,e) -> exp 
		| App (e1,e2) -> 
			 let v1 = eval e1 in 
			 let v2 = eval e2 in 
				(match v1 with 
					 Lambda (v,e) -> 
						eval (subst v2 v e)
				  | _ -> raise (RuntimeError "Application not a function"))
		| Num i -> Num i
		| Prim (p, exps) -> 
			 let vs = List.map eval exps in 
				(match exps with 
					  [Num i;Num j] -> 
						(match p with 
							 Plus  -> Num(i+j) 
						  | Times -> Num(i*j)
						  | Minus -> Num(i-j))
					(* Add cases for new primitives here *)
					| _-> raise (RuntimeError "Invalid argument to primitive"))
		| Let (v, e1, e2) -> 
				eval (subst (eval e1) v e2)
				
(* Alternative 2: Single step version of evaluation *)

let is_val (e:exp) : bool = 
  match e with 
		Lambda _ -> true 
	 | Num _ -> true 
	 | _ -> false 

let rec ss_eval (exp:exp) : exp = 
	 match exp with 
		  Var v -> raise (RuntimeError ("Unbound variable " ^ v))
		| Lambda (v,e) -> exp 
		(* When clause puts extra constraints on pattern matching. *)
		| App (Lambda (x,e), e2) when is_val e2 -> 
			 ss_eval (subst e2 x e)
		| App (e1,e2) when is_val e1 && is_val e2 -> 
			 raise (RuntimeError "Runtime error! not a Lambda")
		| App (e1,e2) when is_val e1 -> 
			 App (e1, ss_eval e2) 
		| App (e1,e2) -> 
			 App (ss_eval e1, e2) 
		| Num i -> Num i
		| Prim (p, [e1;e2]) when is_val e1 && is_val e2 -> 
				(match e1,e2 with 
					  Num i,Num j -> 
						(match p with 
							 Plus  -> Num(i+j) 
						  | Times -> Num(i*j)
						  | Minus -> Num(i-j))
					| _-> raise (RuntimeError "Invalid argument to primitive"))
		| Prim (p, [e1;e2]) when is_val e1 -> 
			 Prim (p, [e1; ss_eval e2])
		| Prim (p, [e1;e2]) -> 
			 Prim (p, [ss_eval e1; e2])
		| Prim (p,_) -> 
			 raise (RuntimeError "Wrong number of args to primitive")
		| Let (x, e1, e2) when is_val e1 -> 
			 (subst e1 x e2)
		| Let (x, e1, e2) -> 
			 Let (x, (ss_eval e1), e2)


let rec eval (exp:exp) : exp = 
  if is_val exp then exp
  else
	 let t = ss_eval exp in 
		begin
(*		  print_exp t; *)
		  eval t
		end 

(* 5 *) 

(* Two lambda expression are alpha-equivalent if their bodies are aeq
 *	after substituting a new name for both variables in the each
 *	expression. We don't have to worry about variable capture because the
 *	new name will be chosen to be different than all of the bound and free
 *	variables in both expressions.*)

(* In order to choose that name, we need to list *all* of the
 * variables that occur anywhere in an expression. *)
let rec vars (e:exp) : set = 
  match e with 
		Var v -> VarSet.singleton v
	 | Lambda (v,e) -> VarSet.add v (vars e)
	 | App (e1,e2) -> VarSet.union (vars e1) (vars e2)
	 | Num i -> VarSet.empty
	 | Prim (p, exps) -> 
		  List.fold_left 
		  VarSet.union  VarSet.empty (List.map vars exps)
	 | Let (v,e1, e2) -> 
		  VarSet.add v (VarSet.union (vars e1) (vars e2))

(* Chooses a new variable that is not found in the set by picking the
 * smallest name (the one closest to the begining of the dictionary)
 * and prepending "A" to it. There are many other ways we could pick a 
 * new name. *)
let new_name (set:set) : var = 
  let smallest = VarSet.min_elt set in 
	 "A" ^ smallest

let rec is_aeq (exp1:exp) (exp2:exp) : bool =
  match exp1,exp2 with 
		Var v1, Var v2 -> v1 = v2
	 | Lambda (v1,e1), Lambda (v2,e2) -> 
		  if v1 = v2 then is_aeq e1 e2
		  else 
			 let new_var = new_name (VarSet.union (vars e1)(vars e2)) in 
				is_aeq (subst (Var new_var) v1 e1) (subst (Var new_var) v2 e2)
	 | App (e1,e1'), App (e2,e2') -> 
		  is_aeq e1 e2 && is_aeq e1' e2'
	 | Num i, Num j -> i = j
	 | Prim (p1,es1), Prim(p2,es2) -> 
		  p1 = p2 &&
				 (try 
					 List.fold_left2 (fun b e1 e2 -> 
											  b && is_aeq e1 e2)
						true es1 es2
				  (* Raised if the two lists are not the same length *)
				  with Invalid_argument _ -> false)
    | Let (v1,e1,e1'), Let(v2,e2,e2') -> 
		  is_aeq e1 e2 &&
		  if v1 = v2 then is_aeq e1' e2'
		  else 
			 let new_var = new_name (VarSet.union (vars e1')(vars e2')) in 
				is_aeq (subst (Var new_var) v1 e1') (subst (Var new_var) v2 e2')
	 | _,_ -> false