lcdbjn.ml

open Format

type term =
  | Var of int
  | Lamb of term
  | App of term * term

let rec shift i depth = function
  | Var x when x < depth -> Var x
  | Var x -> Var (x + i)
  | App (e1, e2) -> App (shift i depth e1, shift i depth e2)
  | Lamb body -> Lamb (shift i (depth + 1) body)

let rec subst x replacement depth = function
  | Var y when x = y -> replacement
  | Var y -> Var y
  | App (e1, e2) ->
    App (subst x replacement depth e1, subst x replacement depth e2)
  | Lamb body ->
    let replacement = shift 1 (depth + 1) replacement in
    let body = subst (x + 1) replacement (depth + 1) body in
    Lamb body

let rec pp_term f = function
  | Var i -> fprintf f "%d" i
  | App (e1, e2) -> fprintf f "(%a %a)" pp_term e1 pp_term e2
  | Lamb body -> fprintf f "λ%a" pp_term body

let rec hnf depth = function
  | Var x -> Var x
  | Lamb body -> Lamb body
  | App (App (_, _) as f, arg) ->
    hnf depth (App (hnf depth f, arg))
  | App (Lamb body, arg) ->
    let body = shift (-1) (depth + 1) body in
    let body = subst 0 arg (depth + 1) body in
    hnf depth body
  | App (_, _) as term ->
    failwith (asprintf "bad application %a" pp_term term)

let rec is_redex = function
  | Var x -> false
  | Lamb body -> is_redex body
  | App (Lamb _, _) -> true
  | App (e1, e2) -> is_redex e1 || is_redex e2 

let rec norm depth = function
  | Var x -> Var x
  | Lamb body -> Lamb (norm (depth + 1) body)
  | App (Lamb body, arg) ->
    let arg = norm depth arg in
    let body = shift (-1) (depth + 1) body in
    let body = subst 0 arg (depth + 1) body in
    if is_redex body
    then norm depth body
    else body
  | App (e1, e2) ->
    let e1 = norm depth e1 in
    let e2 = norm depth e2 in
    let app = App (e1, e2) in
    if is_redex app
    then norm depth app
    else app

let p msg term =
  printf "%s: %a\n" msg pp_term term

let () =
  let f_true = Lamb (Lamb (Var 1)) in
  let f_false = Lamb (Lamb (Var 0)) in
  let f_and = Lamb (Lamb (App (App (Var 1, Var 0), Var 1))) in
  let f_or = Lamb (Lamb (App (App (Var 1, Var 1), Var 0))) in
  let zero = Lamb (Lamb (Var 0)) in
  let succ = Lamb (Lamb (Lamb (App (Var 1, App (App (Var 2, Var 1), Var 0))))) in
  let plus = Lamb (Lamb (Lamb (Lamb (App (App (Var 3, Var 1), App (App (Var 2, Var 1), Var 0)))))) in
  p "true" f_true;
  p "false" f_false;
  p "and true true" (hnf 0 (App (App (f_and, f_true), f_true)));
  p "and true false" (hnf 0 (App (App (f_and, f_true), f_false)));
  p "and false true" (hnf 0 (App (App (f_and, f_false), f_true)));
  p "and false false" (hnf 0 (App (App (f_and, f_false), f_false)));
  p "or true true" (hnf 0 (App (App (f_or, f_true), f_true)));
  p "or true false" (hnf 0 (App (App (f_or, f_true), f_false)));
  p "or false true" (hnf 0 (App (App (f_or, f_false), f_true)));
  p "or false false" (hnf 0 (App (App (f_or, f_false), f_false)));
  p "zero" (hnf 0 zero);
  p "1" (norm 0 (App (succ, zero)));
  let two = norm 0 (App (succ, App (succ, zero))) in
  p "2" two;
  p "4" (norm 0 (App (App (plus, two), two)));
  ()