Capture-avoiding substitution

module String_map = Map.Make (String)

type term =
    | Global of string
    | Local of string
    | Function of string * term
    | Application of term * term

let rec pretty : term -> string = function
    | Global name ->
        "@" ^ name

    | Local name ->
        "$" ^ name

    | Function (parameter, body) ->
        let body' = pretty body in
        "fun $" ^ parameter ^ " is " ^ body' ^ " end"

    | Application (function_, argument) ->
        let function_' = pretty function_ in
        let argument'  = pretty argument  in
        function_' ^ "(" ^ argument' ^ ")"

let next_fresh_variable = ref 0
let fresh_variable () =
    let variable = !next_fresh_variable in
    let _ = next_fresh_variable := variable + 1 in
    "%" ^ string_of_int variable

let rec substitute (work : term String_map.t) : term -> term = function
    | Global name ->
        Global name

    | Local name ->
        begin
            try String_map.find name work with
            | Not_found -> Local name
        end

    | Function (parameter, body) ->
        let fresh = fresh_variable () in
        let work' = String_map.add parameter (Local fresh) work in
        let body' = substitute work' body in
        Function (fresh, body')

    | Application (function_, argument) ->
        let function_' = substitute work function_ in
        let argument'  = substitute work argument  in
        Application (function_', argument')

let rec beta_reduce : term -> term = function
    | Global name ->
        Global name

    | Local name ->
        Local name

    | Function (parameter, body) ->
        let body' = beta_reduce body in
        Function (parameter, body')

    | Application (function_, argument) ->
        begin
            let function_' = beta_reduce function_ in
            let argument'  = beta_reduce argument  in
            match function_' with
            | Function (parameter, body) ->
                let work = String_map.singleton parameter argument in
                substitute work body
            | _ ->
                Application (function_', argument')
        end

let () =
    let term = Application (Function ("x", Application (Local "x", Local "y")), Local "y") in
    let _ = print_endline (pretty term) in
    let _ = print_endline (pretty (substitute (String_map.singleton "y" (Local "x")) term)) in
    let _ = print_endline (pretty (beta_reduce term)) in
    ()