better-optimal.ml

(*
  Start with an environment based interpreter that does
    all the reductions, including under lambdas.

  Then make everything not on the head path lazy.
  When you go to evaluate everything on the side,
    you force the head producing a potentially stuck value
    and then apply the new environment, progressing further.

  Lazy act like graph sharing for every reduction on the side.
  And everything on the head is shared by the environment.
  
  To some extent, the env captured by lazy, a form
    of environment sharing, which seems to be fundamental.
  
  This one makes all the laziness explicit, this is a huge
    win for performance and future optimizations.
*)

type var = int

type term =
  | T_var of var
  | T_let of var * term * term
  | T_lambda of var * term
  | T_apply of term * term
  | T_thunk of term lazy_t
  | T_force of term

and env = E_hole | E_let of env * var * term

let beta_counter = ref 0

let rec lookup env var =
  match env with
  | E_hole -> None
  | E_let (env, env_var, arg) -> (
      match Int.equal var env_var with
      | true -> Some arg
      | false -> lookup env var)

let rec eval_head env term =
  match term with
  | T_var var -> (
      match lookup env var with Some arg -> arg | None -> T_var var)
  | T_let (var, arg, body) ->
      let arg = eval_head 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
      T_lambda (var, body)
  | T_apply (funct, arg) -> (
      let funct = eval_head env funct in
      let arg = eval_head env arg in
      match funct with
      | T_lambda (var, body) ->
          incr beta_counter;
          let env = E_let (E_hole, var, arg) in
          eval_head env body
      | T_var _ | T_let _ | T_apply _ | T_thunk _ | T_force _ ->
          T_apply (funct, arg))
  | T_thunk body ->
      let body =
        lazy
          (let body = Lazy.force body in
           eval_head env body)
      in
      T_thunk body
  | T_force thunk -> (
      let thunk = eval_head env thunk in
      match thunk with
      | T_thunk body -> Lazy.force body
      | T_var _ | T_let _ | T_lambda _ | T_apply _ | T_force _ -> T_force thunk)

let rec lazify term =
  match term with
  | T_var var -> T_force (T_var var)
  | T_let (var, arg, body) ->
      let arg = T_thunk (lazy (lazify arg)) in
      let body = lazify body in
      T_let (var, arg, body)
  | T_lambda (var, body) ->
      let body = T_thunk (lazy (lazify body)) in
      T_lambda (var, body)
  | T_apply (funct, arg) ->
      let funct = lazify funct in
      let arg = T_thunk (lazy (lazify arg)) in
      T_apply (funct, arg)
  | T_thunk _ -> failwith "compile: T_thunk reached"
  | T_force _ -> failwith "compile: T_force reached"

(* test *)
(* just to run the counter *)
let rec eval_side term =
  match term with
  | T_var _ -> ()
  | T_let (_var, _arg, _body) -> failwith "eval_all: let should not happen"
  | T_apply (funct, arg) ->
      eval_side funct;
      eval_side arg
  | T_lambda (_var, body) -> eval_side body
  | T_thunk thunk -> eval_side @@ Lazy.force thunk
  | T_force thunk -> eval_side thunk

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)

(* I = (x) => x; f = (k) => I(k(I(k))); f((y) => f((z) => I)) *)
(* Strong CBV would do 6 instead of 5 *)
let anti_strong_cbv =
  let i_var = 0 in
  let i = T_var i_var in
  let x = T_var 1 in
  let f_var = 2 in
  let f = T_var f_var in
  let k = T_var 3 in
  let y = T_var 4 in
  let z = T_var 5 in
  T_let
    ( i_var,
      x @=> x,
      T_let (f_var, k @=> i @@ k @@ i @@ k, f @@ y @=> f @@ z @=> i) )

(* I = (x) => x; f = (k) => I(k(I(k))); f(f(I)) *)
(* Strong CBN would do 7 instead of 6 *)
let anti_strong_cbn =
  let i_var = 0 in
  let i = T_var i_var in
  let x = T_var 1 in
  let f_var = 2 in
  let f = T_var f_var in
  let k = T_var 3 in
  T_let (i_var, x @=> x, T_let (f_var, k @=> i @@ k @@ i @@ k, f @@ f @@ i))

(* 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_eager name term =
  beta_counter := 0;
  (* Strong CBV-like *)
  eval_side (eval_head E_hole term);
  Format.printf "%s: %d@." name !beta_counter

let eval_optimal name term =
  let term = lazify term in
  eval_eager name term

let () = eval_eager "anti_strong_cbv.eager" anti_strong_cbv
let () = eval_optimal "anti_strong_cbv.optimal" anti_strong_cbv
let () = eval_eager "anti_strong_cbn.eager" anti_strong_cbn
let () = eval_optimal "anti_strong_cbn.optimal" anti_strong_cbn