(* vim: set fileencoding=utf-8 : *)

type unary  = Neg | Not
type binary =   Add | Sub | Mul | Div | Mod    (* arithmetic operators *)
              | Eq | Ne | Lt | Le | Gt | Ge    (* relational operators *)
              | And | Or                       (* logic operators *)
type expr =   Num of int                       (* integer constant *)
            | Bool of bool                     (* boolean constant *)
            | Var of string                    (* variable *)
            | UnaryAp of unary * expr          (* unary operator application *)
            | BinaryAp of binary * expr * expr (* binary operator application *)
            | Cond of expr * expr * expr       (* conditional expression *)
            | Fun of string * expr             (* function abstraction *)
            | FunAp of expr * expr             (* function application *)
            | Let of string * expr * expr      (* variable binding *)
            | LetRec of string * expr * expr   (* recursive variable binding *)

(* Unfortunately, this little gem did not work as expected:
 *
 *   let s f l = List.fold_left Printf.sprintf f l
 *
 * otherwise, I could eliminate all parentheses after the format string. Since
 * pattern matching requires an application after "->". *)

(* Convert an expression to a string. *)
let rec expr2str =
    (* Translate an unary operator to a string. *)
    let unary2str = function
    | Neg -> "-"
    | Not -> "not "
    (* Translate a binary operator to a string. *)
    in let binary2str = function
    | Add -> "+"
    | Sub -> "-"
    | Mul -> "*"
    | Div -> "/"
    | Mod -> "mod"
    | Eq -> "="
    | Ne -> "!="
    | Lt -> "<"
    | Le -> "<="
    | Gt -> ">"
    | Ge -> ">="
    | And -> "&&"
    | Or -> "||"
    (* Aliases: *)
    in let e = expr2str
    in let s = Printf.sprintf
    (* Perform the heavy lifting: *)
    in function
    | Num i                -> s "%i" i
    | Bool b               -> s "%b" b
    | Var v                -> v
    | UnaryAp(e1, e2)      -> s "%s%s" (unary2str e1) (e e2)
    | BinaryAp(e1, e2, e3) -> s "%s %s %s" (e e2) (binary2str e1) (e e2)
    | Cond(e1, e2, e3)     -> s "if %s then %s else %s" (e e1) (e e2) (e e3)
    | Fun(e1, e2)          -> s "fun %s -> %s" e1 (e e2)
    | FunAp(e1, e2)        -> (
        (* Add parentheses for ambiguous cases. *)
        match e2 with 
        | BinaryAp(_,_,_)
        | Cond(_,_,_)
        | Fun(_, _)
        | FunAp(_,_)       -> s "%s (%s)" (e e1) (e e2)
        | _                -> s "%s %s" (e e1) (e e2)
    )
    | Let(e1, e2, e3)      -> s "let %s = %s in %s" e1 (e e2) (e e3)
    | LetRec(e1, e2, e3)   -> s "let rec %s = %s in %s" e1 (e e2) (e e3)
;;

(* Find all free variables of an expression. *)
let freevars e =
    let rec free bound = function
    | Num _
    | Bool _ -> []
    | Var v when List.mem v bound -> []
    (* Bound: FV(x) = {x}. *)
    | Var v -> [v]
    | UnaryAp(_, e) -> free bound e
    | BinaryAp(_, e1, e2)
    (* Application: FV((e1 e2)) = FV(e1) U FV(e2). *)
    | FunAp(e1, e2) -> free bound e1 @ free bound e2
    | Cond(cond, e1, e2) -> free bound cond @ free bound e1 @ free bound e2
    (* Abstraction: FV(λ x.e) = FV(e) \ {x}. *)
    | Fun (arg, e) -> free (arg::bound) e
    | Let (v, e1, e2) -> free (v::bound) e2 @ free (bound) e1
    | LetRec (v, e1, e2) -> free (v::bound) e2 @ free (v::bound) e1
    in free [] e
;;

(* This function implements substitution, i.e. it yields the expression that
 * arises from substituting all free occurrences of the variable specified by
 * the second argument in the expression given as the first argument by the
 * expression given as third argument. *)
let rec subs haystack needle replace =
    (* Alias for normal substitution. *)
    let s e = subs e needle replace
    (* Replace x in e with a new variable w, if x will cause a name clash. *)
    in let unique_var x e =
        let free_replace = freevars replace in
        let free_e = freevars e in
        if List.mem needle free_e && List.mem x free_replace then
            let all_free_vars = free_replace @ free_e in
            (* Add primes to the new variable until it is free *)
            let rec add_prime x =
                let new_x = x ^ "'" in
                if List.mem new_x all_free_vars then add_prime x else new_x
            in
            let w = add_prime x in
            Fun(w, s (subs e x (Var w)))
        else
            Fun(x, s e)
    (* Apply substitution rules: *)
    in match haystack with
    | Var x when x = needle -> replace
    (* Substitute operator arguments. *)
    | UnaryAp(op, e) -> UnaryAp (op, s e)
    | BinaryAp(op, e1, e2) -> BinaryAp (op, s e1, s e2)
    (* 1. x = needle: (λ y.e)[a/x] → λ y.e *)
    | Fun(x, _) as e when x = needle -> e
    (* 2. x ≠ needle: (λ y.e)[a/x] →
     *
     *  if x ∊ FV(e) and y ∊ FV(a):
     *    create new variable w
     *    λ w.(e[w/y][a/x])
     *  else:
     *    λ y.(e[a/x])
     *)
    | Fun(x, e) -> unique_var x e
    (* Substitute in all sub-expressions. *)
    | FunAp(e1, e2) -> FunAp(s e1, s e2)
    | Cond(e1, e2, e3) ->  Cond(s e1, s e2, s e3)
    (* Only substitue in a 'let' or 'let rec' expression if there is no name
     * clash with the assigned variable. *)
    | Let(x, e1, e2) when x != needle -> Let(x, s e1, s e2)
    | LetRec(x, e1, e2) when x != needle -> LetRec(x, s e1, s e2)
    (* No substitition rule. *)
    | _ -> haystack
;;

(* eval yields the normal form of the given expression using applicative order
 * reduction (where applicative order reduction yields the normal form). *)
let rec eval = function
    (* Unary application. *)
    | UnaryAp (op, e) -> (
        match op, eval e with
        | Neg, Num i -> Num(-i)
        | Not, Bool b -> Bool(not b)
        | _ -> raise (Failure "Unknown unary operator")
    )
    (* Binary application. *)
    | BinaryAp (op, e1, e2) -> (
        match op, e1, e2 with
        (* An arithmethic or a relational operation is possible on numbers. *)
        | op, Num a, Num b -> (
            match op with
            | Add -> Num(a + b)
            | Sub -> Num(a - b)
            | Mul -> Num(a * b)
            | Div -> Num(a / b)
            | Mod -> Num(a mod b)
            | Eq  -> Bool(a = b)
            | Ne  -> Bool(a != b)
            | Lt  -> Bool(a < b)
            | Le  -> Bool(a <= b)
            | Gt  -> Bool(a > b)
            | Ge  -> Bool(a >= b)
            | _   -> raise (Failure "Unknown binary operator")
        )
        (* Locical operation is possible on booleans. *)
        | op, Bool a, Bool b -> (
            match op with
            | And -> Bool(a && b)
            | Or  -> Bool(a || b)
            | Eq  -> Bool(a = b)
            | Ne  -> Bool(a != b)
            | _   -> raise (Failure "Unknown boolean operator")
        )
        | op, e1, e2 -> BinaryAp (op, eval e1, eval e2)
    )
    | Fun (x, e) -> Fun (x, eval e)
    (* Function application and expression application. *)
    | FunAp (Fun (x, Var y), a) when y = x -> eval a
    | FunAp (Fun (y, e), a)                -> (subs e y a)
    | FunAp (e1, e2)                       -> FunAp(eval e1, eval e2)
    (* Condition evaluates to a boolean, else evaluate its arguments. *)
    | Cond (cond, e1, e2) -> (
        match eval cond with
        | Bool b -> if b then eval e1 else eval e2
        | c      -> Cond(c, eval e1, eval e2)
    )
    (* Eval'ing ``let'' and ``let rec'' is basic substitution. *)
    | Let (x, e1, e2) -> eval (subs e2 x e1)
    | LetRec (x, e1, e2) -> LetRec(x, eval e1, eval e2)
    | exp -> exp
;;
