Untyped.ml

open Base.Export

type expr =
| Var of string
| Lam of string * expr
| App of expr * expr

let rec a_eq term =
  alpha 0 term
and alpha acc term =
    match term with
    | Var x -> Var x
    | Lam (x, t1) ->
        let y = "#" ^ string_of_int acc in
        Lam (y, alpha (acc + 1) (rename x y t1))
    | App (t1, t2) -> App (alpha acc t1, alpha acc t2)
and rename vx vy term =
  match term with
  | Var x when x == vx -> Var vy
  | Var x -> Var x
  | Lam (x, t1) when x == vx -> Lam (vy, rename vx vy t1)
  | Lam (x, t1) -> Lam (x, rename vx vy t1)
  | App (t1, t2) -> App (rename vx vy t1, rename vx vy t2)

let rec show term =
  match term with
  | Var(s) -> s
  | Lam(s, t) -> 
    "\\" ^ s ^ showFun t
    | App(t1, t2) -> 
      showL t1 ^ showR t2
and showFun t' = 
  match t' with
  | Lam(s1, t1) -> " " ^ s1 ^ showFun t1
  | _ -> "." ^ show t'
and showL t' =
  match t' with
  | Lam _ -> "(" ^ show t' ^ ")"
  | _ -> show t'
and showR t' =
  match t' with
  | Var s -> " " ^ s
  | _ -> "(" ^ show t' ^ ")"

let rec free_vars e =
  match e with
  | Var x -> [x]
  | Lam(x, e1) -> List.filter (fun y -> y != x) (free_vars e1)
  | App(e1, e2) -> (free_vars e1) @ (free_vars e2)

let%test_unit "free_vars" =
  [%test_eq: string list] (free_vars(Var "x")) (["x"]);
  [%test_eq: string list] (free_vars(App(Var "x", Var "y"))) (["x";"y"]);
  [%test_eq: string list] (free_vars(Lam("x", Var "y"))) (["y"]);
  [%test_eq: string list] (free_vars(Lam("x", Lam("y",Var "y")))) ([])

let symbol =
  let counter = ref 0 in
  (fun () -> 
    counter := (!counter) + 1;
    "#" ^ string_of_int !counter)

let rec subst target new_exp cur_exp =
  match cur_exp with
  | Var x -> 
    if target == x then
      new_exp 
    else 
      cur_exp
  | Lam (x, exp) ->
    let new_var = symbol () in
    let new_body = subst x (Var new_var) exp in
    Lam(new_var, subst target new_exp new_body)
  | App(exp1, exp2) -> 
    App(
      subst target new_exp exp1, 
      subst target new_exp exp2)

let rec eval exp =
  match exp with
  | Var _ -> exp
  | App(e1, e2) -> 
    (match eval e1 with
    | Lam(x, body_e) ->
      eval (subst x e2 body_e)
    | _ -> exp)
  | Lam (x, e) -> Lam(x, e)



let rec norm exp =
  match eval exp with
  | Var x -> Var x
  | Lam(x, body_e) -> Lam(x, norm body_e)
  | App(e1, e2) -> App(norm e1, norm e2)

let trueb = Lam("x", (Lam("y", (Var "x"))))
let falseb = Lam("x", (Lam("y", (Var "y"))))
let%test_unit "bool def" =
  [%test_eq : string] (show(a_eq(norm(trueb)))) ("\\#0 #1.#0");
  [%test_eq : string] (show(a_eq(norm(falseb)))) ("\\#0 #1.#1")

let nandb a b = App (App (App(App(a, b), falseb), falseb), trueb)
let notb a = nandb a a
let%test_unit "notb" =
  [%test_eq : string] (show(a_eq(norm(notb(trueb))))) (show(a_eq(falseb)));
  [%test_eq : string] (show(a_eq(norm(notb(falseb))))) (show(a_eq(trueb)))

let andb a b = notb(nandb a b)
let%test_unit "andb truth table" =
  [%test_eq : string] (show(a_eq(norm(andb falseb falseb)))) (show(a_eq(falseb)));
  [%test_eq : string] (show(a_eq(norm(andb trueb falseb)))) (show(a_eq(falseb)));
  [%test_eq : string] (show(a_eq(norm(andb falseb trueb)))) (show(a_eq(falseb)));
  [%test_eq : string] (show(a_eq(norm(andb trueb trueb)))) (show(a_eq(trueb)))

let orb a b = notb(andb (notb a) (notb b))
let%test_unit "orb truth table" =
  [%test_eq : string] (show(a_eq(norm(orb falseb falseb)))) (show(a_eq(falseb)));
  [%test_eq : string] (show(a_eq(norm(orb trueb falseb)))) (show(a_eq(trueb)));
  [%test_eq : string] (show(a_eq(norm(orb falseb trueb)))) (show(a_eq(trueb)));
  [%test_eq : string] (show(a_eq(norm(orb trueb trueb)))) (show(a_eq(trueb)))

let impliesb a b = orb (notb a) b
let%test_unit "impliesb truth table" =
  [%test_eq : string] (show(a_eq(norm(impliesb trueb trueb)))) (show(a_eq(trueb)));
  [%test_eq : string] (show(a_eq(norm(impliesb falseb trueb)))) (show(a_eq(trueb)));
  [%test_eq : string] (show(a_eq(norm(impliesb trueb falseb)))) (show(a_eq(falseb)));
  [%test_eq : string] (show(a_eq(norm(impliesb falseb falseb)))) (show(a_eq(trueb)))

let ifb a b c = andb (impliesb a b) (impliesb (notb a) c)
  let%test_unit "ifb truth table" =
  let left = trueb in
  let right = falseb in
  [%test_eq : string] (show(a_eq(norm(ifb falseb left right)))) (show(a_eq(right)));
  [%test_eq : string] (show(a_eq(norm(ifb trueb left right)))) (show(a_eq(left)))

let eqb a b = orb (notb(orb a b)) (andb a b)
let%test_unit "eqb truth table" =
  [%test_eq : string] (show(a_eq(norm(eqb falseb falseb)))) (show(a_eq(trueb)));
  [%test_eq : string] (show(a_eq(norm(eqb falseb trueb)))) (show(a_eq(falseb)));
  [%test_eq : string] (show(a_eq(norm(eqb trueb falseb)))) (show(a_eq(falseb)));
  [%test_eq : string] (show(a_eq(norm(eqb trueb trueb)))) (show(a_eq(trueb)))

let succ n = 
  App(
    Lam("n", Lam("s", Lam("z",
      App(
        Var "s", 
        App(
          App(Var "n", Var "s"), 
          Var "z"))))), 
    n)

let n0 = Lam("s", Lam("z", Var "z"))
let n1 = succ n0
let n2 = succ n1
let n3 = succ n2
let n4 = succ n3
let n5 = succ n4
let n6 = succ n5
let n7 = succ n6
let n8 = succ n7
let n9 = succ n8
let n10 = succ n9

let plus m n = 
  App(App(
      Lam("m", Lam("n", Lam("s", Lam("z",
        App(
          App(Var "m", Var "s"),
          App(
            App(Var "n", Var "s"), 
            Var "z")))))), 
    m), n)
let%test_unit "plus" =
  [%test_eq : string] (show(a_eq(norm(plus n0 n0)))) (show(a_eq(norm(n0))));
  [%test_eq : string] (show(a_eq(norm(plus n0 n1)))) (show(a_eq(norm(n1))));
  [%test_eq : string] (show(a_eq(norm(plus n1 n0)))) (show(a_eq(norm(n1))));
  [%test_eq : string] (show(a_eq(norm(plus n1 n1)))) (show(a_eq(norm(n2))));
  [%test_eq : string] (show(a_eq(norm(plus n1 n2)))) (show(a_eq(norm(n3))))

let mult m n = 
  App(App(
    Lam("m", Lam("n", Lam("s",
      App(
        Var "m", 
        App(Var "n", Var "s"))))),
    m), n)

let%test_unit "mult" =
  [%test_eq : string] (show(a_eq(norm(mult n0 n0)))) (show(a_eq(norm(n0))));
  [%test_eq : string] (show(a_eq(norm(mult n1 n0)))) (show(a_eq(norm(n0))));
  [%test_eq : string] (show(a_eq(norm(mult n0 n1)))) (show(a_eq(norm(n0))));
  [%test_eq : string] (show(a_eq(norm(mult n1 n2)))) (show(a_eq(norm(n2))));
  [%test_eq : string] (show(a_eq(norm(mult n2 n1)))) (show(a_eq(norm(n2))));
  [%test_eq : string] (show(a_eq(norm(mult n3 n2)))) (show(a_eq(norm(n6))))

let pred n = 
  App(
    Lam("n", Lam("f", Lam("x", 
        App(
          App(
            App(
              Var "n",
              Lam("g", Lam("h", App(Var "h", App(Var"g", Var "f"))))),
            Lam("u", Var "x")),
          Lam("u", Var "u"))))), 
    n)

let%test_unit "pred" =
  [%test_eq : string] (show(a_eq(norm (pred n0)))) (show(a_eq(norm n0)));
  [%test_eq : string] (show(a_eq(norm (pred n1)))) (show(a_eq(norm n0)));
  [%test_eq : string] (show(a_eq(norm (pred n2)))) (show(a_eq(norm n1)))

let is0 n = 
  App(Lam("n", 
    App(
      App(Var "n", Lam("x", falseb)), 
      trueb)), 
    n)

let%test_unit "is0" =
  [%test_eq : string] (show(a_eq(norm(is0 n0)))) (show(norm(a_eq(trueb))));
  [%test_eq : string] (show(a_eq(norm(is0 n1)))) (show(norm(a_eq(falseb))))

let yc g =
  let f = Var "f" in
  let x = Var "x" in
  App(
    Lam ("f", 
      App(
        (Lam ("x", App (f, App (x, x)))),
        (Lam ("x", App (f, App(f, App (x, x))))))), 
    g);;

let y f =
  let rec g x = f (g x) in
  g

let fact g = 
  App(
    yc(
      Lam("f", Lam("n", 
        App(
          App((is0 (Var "n")), n1), 
          (mult (Var "n") (App(Var "f", pred (Var "n")))))))),
    g)
let%test_unit "y combinator" =
  [%test_eq : string] (show(a_eq(norm(fact n0)))) (show(a_eq(norm(n1))));
  [%test_eq : string] (show(a_eq(norm(fact n1)))) (show(a_eq(norm(n1))));
  [%test_eq : string] (show(a_eq(norm(fact n2)))) (show(a_eq(norm(n2))));
  [%test_eq : string] (show(a_eq(norm(fact n3)))) (show(a_eq(norm(n6))));