(* Basic lambda calculus definition, printing and evaluation. *)

(* ocamlc lambda.ml -o lambda *)

type lambda =
  | Var of string
  | Lam of string * lambda
  | App of lambda * lambda;;

module Free = Set.Make(String);;

(* Turn a lambda expression into a human-readable string *)
let rec print_lambda = function
  | Var x -> x
  | Lam (var, lambda) -> String.concat "" ["(\\"; var; "."; (print_lambda lambda); ")"]
  | App (l1, l2) -> String.concat "" [(print_lambda l1); " "; (print_lambda l2)]

(* Compute the set of unbound variables in a given lambda expression *)
let rec free_variables lambda =
  let rec inner lambda bound_vars =
    match lambda with
	  | Var x -> if Free.mem x bound_vars then Free.empty else Free.singleton x
	  | Lam (var, lambda) -> inner lambda (Free.add var bound_vars)
	  | App (l1, l2) -> Free.union (inner l1 bound_vars) (inner l2 bound_vars)
  in
    inner lambda Free.empty

(* Turn a set of unbound variables into a human-readable string *)
let print_free vars = String.concat ", " (Free.elements vars)

(* Substitutes a lambda expression for a variable in a given expression. *)
let rec substitute lambda var expr =
  match lambda with
    | Var x -> if x = var then expr else Var x
	| Lam (x, inner) -> if x = var then Lam (x, inner) else Lam (x, substitute inner var expr)
	| App (l1, l2) -> App (substitute l1 var expr, substitute l2 var expr)

(* Beta-reduce a lambda expression in a given environment of variable bindings *)
let rec reduce = function
  | Var x -> Var x
  | App (l1, l2) ->
    (match l1 with
	  | Lam (var, lambda) ->
	    substitute lambda var l2
	  | something_else ->
	    App (reduce something_else, l2))
  | Lam (var, lambda) -> Lam (var, lambda)

(* Beta-reduce a lambda expression a given number of times with no prior variable
bindings *)
let rec reduce_n lambda n =
  if n = 0 then lambda else
    reduce_n (reduce lambda) (n - 1)

(* Display the results of each stage of beta-reducing a lambda expression. Used
for debugging. *)
let display_n_stages lambda n = begin
  print_string "Evaluating expression: ";
  print_string (print_lambda lambda);
  print_string "\n";

  for iteration = 1 to n do
    print_string "Iteration ";
	print_int iteration;
	print_string ": ";
    print_string (print_lambda (reduce_n lambda iteration));
    print_string "\n";
  done
end

(* Beta-reduce a lambda expression until it can't be reduced any further *)
let rec eval lambda =
  let lambda' = reduce lambda in
    if lambda = lambda' then lambda
	else eval lambda'

(* Testing things. *)
let x_id = Lam ("x", Var "x")
let y_id = Lam ("y", Var "y")
let swap_f_g = Lam ("f", (Lam ("g", (App ((Var "g"), (Var "f"))))))
(* (\f.(\g.g f)) (\x.x) (\y.y) *)
let swap_ids = App (App(swap_f_g, x_id), y_id)

(* (\f.(\g.g f)) y x
Free vars: x, y *)
let mess_of_frees = App ((App (swap_f_g, (Var "y"))), (Var "x"))

(* Construct a swap lambda which doesn't work without alpha conversion.
(\x.(\y.y x)) (\y.y) (\x.x) *)
let swap_x_y = Lam ("x", (Lam ("y", (App ((Var "y"), (Var "x"))))))
let swap_ids_pre_alpha = App (App(swap_x_y, y_id), x_id)
;;

print_string "Free var computation for ";
print_string (print_lambda mess_of_frees);
print_string "\n";
print_string (print_free (free_variables mess_of_frees));
print_string "\n\n";
display_n_stages swap_ids 3;
print_string "\n";
display_n_stages swap_ids_pre_alpha 3;
print_string "\n";
print_string "Evaluating lambda: ";
print_string (print_lambda swap_ids_pre_alpha);
print_string "\n";
print_string "Result: ";
print_string (print_lambda (eval swap_ids_pre_alpha));
;;