(* This is a simple compiler from Core ML to Rube VM.
   It is based upon the Core ML to OCaml compiler we
   previosly wrote. *)
open Rubevm
open Instr

type var = string
type expr =
  | EInt of int
  | EVar of var
  | ELam of var * expr
  | EApp of expr * expr

let compile (e : expr) : prog = 
  let ht = Hashtbl.create 123456 in

  let gensym : string -> string =
    let i = ref 0 in
    fun s ->
    i := !i + 1;
    s ^ (string_of_int !i)
  in

  let emit : string -> instr array -> string =
    fun s is ->
    let s' = gensym s in
    Hashtbl.add ht s' is;
    s'
  in

  let rec comp (e : expr) : string =
    match e with
    | EInt i -> 
       emit "int" 
	    [| I_const (`L_Reg 0, `L_Int i); 
	       I_ret (`L_Reg 0) |]
    | EVar x ->
       emit "var"
	    [| I_const (`L_Reg 1, `L_Str x);
	       I_const (`L_Reg 2, `L_Id "lookup");
	       I_call (`L_Reg 2, 0, 2);
	       I_ret (`L_Reg 0); |]
    | ELam (x, e) ->
       let c = comp e in
       emit "lam"
	    [| I_mk_tab (`L_Reg 1);
	       I_const (`L_Reg 2, `L_Str "var");
	       I_const (`L_Reg 3, `L_Str x);
	       I_const (`L_Reg 4, `L_Str "code");
	       I_const (`L_Reg 5, `L_Id c);
	       I_const (`L_Reg 6, `L_Str "env");
	       
	       I_wr_tab (`L_Reg 1, `L_Reg 2, `L_Reg 3);
	       I_wr_tab (`L_Reg 1, `L_Reg 4, `L_Reg 5);
	       I_wr_tab (`L_Reg 1, `L_Reg 6, `L_Reg 0);
	       I_ret (`L_Reg 1) |]

    | EApp (e1, e2) ->
       let c1 = comp e1 in
       let c2 = comp e2 in
       emit "app"
	    [| I_mov (`L_Reg 1, `L_Reg 0);
	       I_mov (`L_Reg 2, `L_Reg 0);

	       I_const (`L_Reg 3, `L_Id c1);
	       I_const (`L_Reg 4, `L_Id c2);
	       I_call (`L_Reg 3, 1, 2);
	       I_call (`L_Reg 4, 2, 3);

	       I_const (`L_Reg 5, `L_Str "var");
	       I_const (`L_Reg 6, `L_Str "code");
	       I_const (`L_Reg 7, `L_Str "env");

	       I_rd_tab (`L_Reg 8, `L_Reg 1, `L_Reg 5);  (* var *)
	       I_mov (`L_Reg 9, `L_Reg 2);               (* val *)
	       I_rd_tab (`L_Reg 10, `L_Reg 1, `L_Reg 7); (* env *)

	       I_const (`L_Reg 11, `L_Id "extend");
	       I_call (`L_Reg 11, 8, 11);

	       I_rd_tab (`L_Reg 12, `L_Reg 1, `L_Reg 6); (* code *)
	       I_call (`L_Reg 12, 8, 9);
	       I_ret (`L_Reg 8); |]
  in
  let main = comp e in

  Hashtbl.add ht "extend"
	      (* r0: var *)
	      (* r1: val *)
	      (* r2: env *)
	      [| I_const (`L_Reg 3, `L_Str "var");
		 I_const (`L_Reg 4, `L_Str "val");
		 I_const (`L_Reg 5, `L_Str "env");
		 
		 I_mk_tab (`L_Reg 6);
		 
		 I_wr_tab (`L_Reg 6, `L_Reg 3, `L_Reg 0);
		 I_wr_tab (`L_Reg 6, `L_Reg 4, `L_Reg 1);
		 I_wr_tab (`L_Reg 6, `L_Reg 5, `L_Reg 2);

		 I_ret (`L_Reg 6) |];	      

  Hashtbl.add ht "lookup" 
	      (* r0: env *)
	      (* r1: var *)
	      [| I_const (`L_Reg 2, `L_Str "var");

		 I_rd_tab (`L_Reg 3, `L_Reg 0, `L_Reg 2);
		 I_eq (`L_Reg 5, `L_Reg 3, `L_Reg 1);
		 I_if_zero (`L_Reg 5, 3);

		 I_const (`L_Reg 2, `L_Str "val");
		 I_rd_tab (`L_Reg 6, `L_Reg 0, `L_Reg 2);
		 I_ret (`L_Reg 6);

		 I_const (`L_Reg 7, `L_Str "env");
		 I_rd_tab (`L_Reg 8, `L_Reg 0, `L_Reg 7);
		 I_mov (`L_Reg 0, `L_Reg 8);
		 
		 I_jmp (-10);
		|];
		 
  
  Hashtbl.add ht "main" 
	      [| I_const (`L_Reg 1, `L_Id main);
		 I_mk_tab (`L_Reg 0);		 
		 I_call (`L_Reg 1, 0, 1);
		 I_ret (`L_Reg 0) |];
  ht

let go () =
  let p = compile (EApp ((EApp (ELam ("z", (ELam ("q", EVar "z"))), EInt 7)), EInt 3))  in
  let v = Rubevm.run_prog p in
  Printf.printf "> %a\n" Disassembler.value v
;;

go ()