small_cbv.ml

(* index when term, level when value *)
type var = int
type index = int
type level = int

type term =
  (* x *)
  | T_var of { var : var }
  (* x = M; N *)
  | T_let of { arg : term; body : term }
  (* x => M *)
  | T_lambda of { body : term }
  (* M(N) *)
  | T_apply of { funct : term; arg : term }
  (* [M, N] *)
  | T_pair of { fst : term; snd : term }
  (* M.0 *)
  | T_fst of { pair : term }
  (* M.1 *)
  | T_snd of { pair : term }

(* value is a term where var is interpreted as a level *)
type value = term

(* a term with a hole *)
type context =
  (* • *)
  | C_hole
  (* C<x = •; K> *)
  | C_let_arg of { ctx : context; body : term }
  (* C<x = M; •> *)
  | C_let_body of { ctx : context; arg : value }
  (* C<M(•)> *)
  | C_apply_arg of { ctx : context; funct : term }
  (* C<•(N)> *)
  | C_apply_funct of { ctx : context; arg : value }
  (* C<[•, N]> *)
  | C_pair_fst of { ctx : context; snd : term }
  (* C<[M, •]> *)
  | C_pair_snd of { ctx : context; fst : value }
  (* C<•.0> *)
  | C_fst of { ctx : context }
  (* C<•.1> *)
  | C_snd of { ctx : context }

type head =
  | H_var of { var : var }
  | H_lambda of { body : term }
  | H_pair of { fst : value; snd : value }

let rec lookup ctx var =
  match ctx with
  | C_hole -> failwith "lookup: unbound variable"
  | C_let_body { ctx; arg } -> (
      match var with 0 -> arg | _ -> lookup ctx (var - 1))
  | C_let_arg { ctx; body = _ }
  | C_apply_arg { ctx; funct = _ }
  | C_apply_funct { ctx; arg = _ }
  | C_pair_fst { ctx; snd = _ }
  | C_pair_snd { ctx; fst = _ }
  | C_fst { ctx }
  | C_snd { ctx } ->
      lookup ctx var

(* reduce *)
let step ctx head =
  match (ctx, head) with
  | ctx, H_var { var } ->
      let value = lookup ctx var in
      (ctx, `subst value)
  | C_apply_funct { ctx; arg }, H_lambda { body } ->
      let ctx = C_let_body { ctx; arg } in
      (ctx, `beta body)
  | C_fst { ctx }, H_pair { fst; snd = _ } -> (ctx, `fst fst)
  | C_snd { ctx }, H_pair { fst = _; snd } -> (ctx, `snd snd)
  | ctx, H_lambda { body } -> (ctx, `head (T_lambda { body }))
  | ctx, H_pair { fst; snd } -> (ctx, `head (T_pair { fst; snd }))

(* search *)
let rec next_down ctx term =
  match term with
  | T_var { var } -> (ctx, H_var { var })
  | T_let { arg; body } ->
      let ctx = C_let_arg { ctx; body } in
      next_down ctx arg
  | T_lambda { body } -> (ctx, H_lambda { body })
  | T_apply { funct; arg } ->
      (* arg first *)
      let ctx = C_apply_arg { ctx; funct } in
      next_down ctx arg
  | T_pair { fst; snd } ->
      (* fst first *)
      let ctx = C_pair_fst { ctx; snd } in
      next_down ctx fst
  | T_fst { pair } ->
      let ctx = C_fst { ctx } in
      next_down ctx pair
  | T_snd { pair } ->
      let ctx = C_snd { ctx } in
      next_down ctx pair

let rec next_up ctx value =
  match ctx with
  | C_hole -> `halt value
  | C_let_arg { ctx; body } ->
      let ctx = C_let_body { ctx; arg = value } in
      `next (ctx, body)
  | C_let_body { ctx; arg } ->
      (* TODO: gc? *)
      let value = T_let { arg; body = value } in
      next_up ctx value
  | C_apply_arg { ctx; funct } ->
      let ctx = C_apply_funct { ctx; arg = value } in
      `next (ctx, funct)
  | C_apply_funct { ctx; arg } ->
      let value = T_apply { funct = value; arg } in
      next_up ctx value
  | C_pair_fst { ctx; snd } ->
      let ctx = C_pair_snd { ctx; fst = value } in
      `next (ctx, snd)
  | C_pair_snd { ctx; fst } ->
      let value = T_pair { fst; snd = value } in
      next_up ctx value
  | C_fst { ctx } ->
      let value = T_fst { pair = value } in
      next_up ctx value
  | C_snd { ctx } ->
      let value = T_snd { pair = value } in
      next_up ctx value

let rec eval_down ctx term =
  let ctx, head = next_down ctx term in
  let ctx, step = step ctx head in
  match step with
  | `beta body -> eval_down ctx body
  | `subst value | `fst value | `snd value | `head value -> eval_up ctx value

and eval_up ctx value =
  match next_up ctx value with
  | `halt value -> value
  | `next (ctx, term) -> eval_down ctx term