lsc.ml

type var = int

type term =
  (* x *)
  | T_var of var
  (* M N *)
  | T_apply of term * term
  (* x => M *)
  | T_lambda of var * term
  (* x = M; N *)
  | T_let of var * term * term

let rename : term -> term = failwith "TODO"
let ( let* ) = Option.bind

(* Pretty much any naive implementation that you do
      can be reasoned about it in terms of LSC + equivalences *)

(* This looks way better by having a notion of head_term
      but that would make it harder to get the intuition *)
module Small_step = struct
  let rec split_env env term =
    match term with
    | T_let (var, arg, body) ->
        let env = (var, arg) :: env in
        split_env env body
    | term -> (env, term)

  let rec merge_env env body =
    match env with
    | [] -> body
    | (var, arg) :: env ->
        let body = merge_env env body in
        T_let (var, arg, body)

  (* this is bad *)
  let rec try_subst from to_ term : term option =
    match term with
    | T_var var -> (
        match Int.equal var from with
        | true -> Some (rename to_)
        | false -> None)
    | T_apply (funct, arg) -> (
        match try_subst from to_ funct with
        | Some funct -> Some (T_apply (funct, arg))
        | None ->
            let* arg = try_subst from to_ arg in
            Some (T_apply (funct, arg)))
    | T_lambda (var, body) ->
        let* body = try_subst from to_ body in
        Some (T_lambda (var, body))
    | T_let (var, arg, body) -> (
        match try_subst from to_ arg with
        | Some arg -> Some (T_let (var, arg, body))
        | None ->
            let* body = try_subst from to_ body in
            Some (T_let (var, arg, body)))

  let try_subst from to_ term = try_subst from to_ @@ rename term

  let step_head term =
    match term with
    | T_var var -> (false, T_var var)
    | T_apply (funct, arg) -> (
        let env, funct = split_env [] funct in
        match funct with
        | T_lambda (var, body) ->
            (* dβ *)
            let env = (var, arg) :: env in
            (true, merge_env env body)
        | funct ->
            let funct = merge_env env funct in
            (false, T_apply (funct, arg)))
    | T_lambda (var, body) -> (false, T_lambda (var, body))
    | T_let (var, arg, body) -> (
        match try_subst var arg body with
        | Some body ->
            (* subst *)
            (true, T_let (var, arg, body))
        | None ->
            (* gc *)
            (true, body))

  let rec eval_head term =
    match step_head term with
    | true, term -> eval_head term
    | false, term -> term

  let rec normalize term =
    match eval_head term with
    | T_var var -> T_var var
    | T_apply (funct, arg) ->
        let funct = normalize funct in
        let arg = normalize arg in
        T_apply (funct, arg)
    | T_lambda (var, body) ->
        let body = normalize body in
        T_lambda (var, body)
    | T_let (var, arg, body) ->
        (* this can happen if we add meta variables *)
        let arg = normalize arg in
        let body = normalize body in
        T_let (var, arg, body)
end

module Big_step = struct
  (* here I work modulo L<M> N == M N *)
  (* eval_head ((x, M) :: e) N == eval_head e (x = M; N) *)
  let rec eval_head env term =
    match term with
    | T_var var ->
        (* there are more efficient ways of doing rename *)
        (* subst *)
        let term = List.assoc var env in
        let term = rename term in
        eval_head env term
    | T_apply (funct, arg) -> (
        let env, funct = eval_head env funct in
        match funct with
        | T_lambda (var, body) ->
            (* dβ *)
            let env = (var, arg) :: env in
            eval_head env body
        | T_var _ | T_apply (_, _) | T_let (_, _, _) ->
            (env, T_apply (funct, arg)))
    | T_lambda (var, body) -> (env, T_lambda (var, body))
    | T_let (var, arg, body) ->
        let env = (var, arg) :: env in
        eval_head env body

  let rec normalize env term =
    let env, term = eval_head env term in
    match term with
    | T_var var -> (env, T_var var)
    | T_apply (funct, arg) ->
        let env, funct = normalize env funct in
        let env, arg = normalize env arg in
        (env, T_apply (funct, arg))
    | T_lambda (var, body) ->
        let env, body = normalize env body in
        (env, T_lambda (var, body))
    | T_let (var, arg, body) ->
        (* TODO: I don't think this one ever happens *)
        let env, arg = normalize env arg in
        let env, body = normalize env body in
        (env, T_let (var, arg, body))

  let normalize term =
    let _env, term = normalize [] term in
    (* gc *)
    term
end