Lambda calculus on OCaml

Ast.ml

module Varset = Set.Make(String)

type ast  =
    | Var of        name
    | Val of        int
    | Fun of info * varname *  ast
    | App of info * ast     *  ast
    | BIF of        name    * (ast -> ast)
    | Now of info * ast

and varname = string
and name    = string
and info    = { context : varset }
and varset  = Varset.t

let rec contextOf ast = let open Varset in
    match ast with
    | Var (name)    -> { context = singleton name }
    | Fun (c, _, _) -> c
    | App (c, _, _) -> c
    | Now (c, _)    -> c
    | _             -> { context = empty }

let unite c1 c2 = { context = Varset.union c1.context c2.context }

let v x = Var(x)
let c i = Val(i)
let f a b = Fun(contextOf b, a, b)
let ( @ ) f x = App(unite (contextOf f) (contextOf x), f, x)
let b n f = BIF(n, f)
let n x   = Now(contextOf x, x)

let ( <| ) f x = f x
let ( <@> ) f g = fun x -> f (g x)

let rec sprint_ast ast = let open Printf in
    match ast with
    | Var (name)         -> name
    | Val (i)            -> sprintf "%d" i
    | Fun (_, arg, body) -> sprintf "%s. %s" arg (sprint_ast body)
    
    | App (_, Fun(d, e, f), App(_, g, x)) -> 
        sprintf "(%s) (%s %s)" 
            (sprint_ast (Fun(d, e, f)))
            (sprint_ast g)
            (sprint_ast x)
    
    | App (_, f, App(_, g, x)) -> 
        sprintf "%s (%s %s)" 
            (sprint_ast f)
            (sprint_ast g)
            (sprint_ast x)
    
    | App (_, Fun(d, e, f), Fun(c, a, b)) ->
        sprintf "(%s) (%s)"
            (sprint_ast (Fun(d, e, f)))
            (sprint_ast (Fun(c, a, b)))

    | App (_, Fun(c, a, b), x) ->
        sprintf "(%s) %s"
            (sprint_ast (Fun(c, a, b)))
            (sprint_ast x)

    | App (_, f, Fun(c, a, b)) ->
        sprintf "%s (%s)"
            (sprint_ast f)
            (sprint_ast (Fun(c, a, b)))

    | App (_, f, x) ->
        sprintf "%s %s"
            (sprint_ast f)
            (sprint_ast x)

    | Now (_, x) ->
        sprintf "[%s]" (sprint_ast x)

    | BIF (name, _) -> name

let without elem set = { context = Varset.remove elem set.context }

exception VariableEscaped of varname
exception NotApplicable   of string

let rec unless pred transform x =
    match pred x with
    | false -> 
        begin x 
            :: try (unless pred transform (transform x)) 
                with NotApplicable(msg) ->
                    Printf.printf "! Error - not applicable: %s\n\n" msg ;
                    []
        end

    | true  -> [x]

let rec is_terminal tree = match tree with
    | Val   (_)   -> true
    | Fun (_,_,_) -> true
    | App (_,_,_) -> false
    | BIF   (_,_) -> true
    | Now (_,x)   -> is_terminal x
    | Var   (n)   -> raise <| VariableEscaped (n)

let rec step tree =
    match tree with 
    | App (c1, Now (c2, g),      x)  -> App (c1, step (Now (c2, g)),    x)
    | App (c1, f,       Now (c2, x)) -> App (c1, f,      step (Now (c2, x)))
    | App (c1, App (c2, g,   y), x)  -> App (c1, step (App (c2, g, y)), x)
    | App (c1, BIF (name,    f), x)  -> f x
    | App (c1, Fun (c2, var, f), x)  -> (x ==> var) f
    | App (c1, f,                x)  -> raise 
                                     <| NotApplicable 
                                        ( "`"
                                        ^ sprint_ast f 
                                        ^ "` to `" 
                                        ^ sprint_ast x 
                                        ^ "`"
                                        ) 
    
    | Now(c1, x) when is_terminal x  -> x
    | Now(c1, x)                     -> Now(c1, step x)

    | Var(name)                      -> raise <| VariableEscaped (name)

    | _                              -> tree

and ( ==> ) fact formal = 
    let rec substitute tree =
        if not (Varset.mem formal (contextOf tree).context)
        then tree
        else match tree with
        | Var(name) when name = formal -> 
            fact

        | App(c1, f, x) -> 
            App
                ( without formal c1
                , substitute f
                , substitute x
                )

        | Fun(c1, arg, body) when arg <> formal ->
            Fun
                ( without formal c1
                , arg
                , substitute body
                )
    
        | Now(c1, x) ->
            Now(without formal c1, substitute x)

        | _ -> tree
    in
        substitute

let head = f "x" <| (f "y" <| v "x")
let tail = f "x" <| (f "y" <| v "y")

let cons = f "x" <| (f "y" <| (f "s" <| ((v "s" @ v "x") @ v "y")))

let lit a b c = f a c @ b

let test 
    =  lit "cons" (f "x" <| (f "y" <| (f "s" <| ((v "s" @ v "x") @ v "y")))) 
    <| lit "tail" (f "x" <| (f "y" <| v "y"))
        (((v "cons" @ c 1) @ c 2) @ v "tail")

let rec prettify pairs =
    let rec inner code = 
        match List.filter (fun (ast, name) -> ast = code) pairs with
        | [] ->
            begin
                match code with
                | Fun (c,a,b) -> Fun (c, a, inner b)
                | App (c,f,x) -> App (c, inner f, inner x)
                | Now (c,x)   -> Now (c, inner x)
                | _           -> code
            end

        | ((_, name) :: _) -> v name
    in inner

let uncode = prettify 
    [ f "x" <| (f "y" <| v "y")
    , "tail"

    ; f "x" <| (f "y" <| (f "s" <| ((v "s" @ v "x") @ v "y")))
    , "cons"
    ]

let () = List.iter (Printf.printf "%s\n" <@> sprint_ast <@> uncode) <| unless is_terminal step test