nbe-optimal.ml

(* this whole thing is a sketch *)
type var = int [@@deriving show { with_path = false }]

type term =
  | T_var of var
  | T_let of var * term * term
  | T_lambda of var * term
  | T_apply of term * term
  | T_reify of thunk

and value = V_var of var | V_apply of value * thunk | V_lambda of var * term
and env = E_hole | E_let of env * var * thunk

and thunk = {
  mutable env : env;
  mutable term : term;
  mutable done_ : value;
  mutable reify : term;
}

(* TODO: physical identity *)
let v_null : value = V_var (-1)
let is_v_null value = value == v_null
let t_null : term = T_var (-1)
let is_t_null term = term == t_null
let thunk env term = { env; term; done_ = v_null; reify = t_null }
let beta_counter = ref 0

let rec eval_head env term =
  match term with
  | T_var var -> eval_var env var
  | T_let (var, arg, body) ->
      let arg = thunk env arg in
      let env = E_let (env, var, arg) in
      eval_head env body
  | T_lambda (var, body) ->
      let body = eval_head env body in
      let body = reify_head body in
      V_lambda (var, body)
  | T_apply (funct, arg) -> (
      let funct = eval_head env funct in
      let arg = thunk env arg in
      match funct with
      | V_lambda (var, body) ->
          (* TODO: beta counter *)
          incr beta_counter;
          let env = E_let (E_hole, var, arg) in
          eval_head env body
      | V_var _ | V_apply _ -> V_apply (funct, arg))
  | T_reify thunk ->
      let term = eval_reify thunk in
      eval_head env term

and eval_var env var =
  match env with
  | E_hole ->
      (* TODO: is this okay? *)
      V_var var
  | E_let (env, env_var, arg) -> (
      match Int.equal var env_var with
      | true -> eval_force arg
      | false -> eval_var env var)

and eval_force thunk =
  let { env; term; done_; reify = _ } = thunk in
  match is_v_null done_ with
  | true ->
      let done_ = eval_head env term in
      thunk.done_ <- done_;
      (* free env and term *)
      thunk.env <- E_hole;
      thunk.term <- t_null;
      done_
  | false -> done_

and eval_reify thunk =
  let { env = _; term = _; done_ = _; reify } = thunk in
  match is_t_null reify with
  | true ->
      let done_ = eval_force thunk in
      let reify = reify_head done_ in
      thunk.reify <- reify;
      reify
  | false -> reify

and reify_head value =
  match value with
  | V_var var -> T_var var
  | V_apply (funct, arg) ->
      let funct = reify_head funct in
      let arg = T_reify arg in
      T_apply (funct, arg)
  | V_lambda (var, body) -> T_lambda (var, body)

(* just to run the counter *)
let rec eval_all value =
  match value with
  | V_var _ -> ()
  | V_apply (funct, arg) ->
      eval_all funct;
      eval_all_thunk arg
  | V_lambda (_var, body) -> eval_all_term body

and eval_all_thunk thunk = eval_all @@ eval_force thunk

and eval_all_term term =
  match term with
  | T_var _ -> ()
  | T_let (_, arg, body) ->
      eval_all_term arg;
      eval_all_term body
  | T_lambda (_, body) -> eval_all_term body
  | T_apply (funct, arg) ->
      eval_all_term funct;
      eval_all_term arg
  | T_reify thunk -> eval_all_thunk thunk

(* test *)

let ( @@ ) funct arg = T_apply (funct, arg)

let ( @=> ) var body =
  let var = match var with T_var var -> var | _ -> assert false in
  T_lambda (var, body)

(* let id = 10 @=> fun x -> x *)

(* let term =
  let f = T_var 0 in
  T_let
    ( 0,
      ( 1 @=> fun k ->
        2 @=> fun x -> id @@ k @@ id @@ x ),
      f @@ (3 @=> fun _x -> (f @@ id) @@ id) @@ id ) *)

(* let term =
  let f = T_var 0 in
  let x = T_var 1 in
  let y = T_var 2 in
  f @=> (x @=> f @@ x @@ x) @@ y @=> f @@ y @@ x *)

let term =
  let x = T_var 0 in
  let y = T_var 1 in
  let z = T_var 2 in
  let w = T_var 3 in
  let u = T_var 4 in
  (x @=> (x @@ y) @@ x @@ z) @@ w @=> (u @=> u) @@ w

let () = eval_all (eval_head E_hole term)
let () = Format.printf "beta: %d@." !beta_counter