(* A trivial input language of sums and ints *)
module Ast = struct
  type t = Int of int32 | Plus of t * t

  (* ((1 + 2) + (3 + (4 + 5))) *)
  let ast1 = Plus(Plus(Int 1l, Int 2l),
		  Plus(Int 3l, Plus(Int 4l, Int 5l)))

end

(* X86micro -- only a few instructionss -- no labels, etc. *)
module Asm = struct

  type reg = EAX | EBX | ECX | EDX | ESI | EDI | EBP | ESP
  type op = Reg of reg | Imm of int32 | Ind of reg * int
  type insn = Add of op * op | Mov of op * op | Push of op | Sub of op * op

end

(* Example1 shows a simple, recursive translation to assembly code. *)
module Example1 = struct
  open Ast
  open Asm

  (* Invariant: final value of the translated expression is stored in EAX *)
  (* Temporary values are stored on the stack. *)
  (* Note that the order of evaluation is determined here too.*)
  let compile (ast:Ast.t) : Asm.insn list = 
    let rec translate (ast:Ast.t) (k:Asm.insn list) =
      match ast with
	| Int n -> Mov(Reg EAX, Imm n) :: k
	| Plus(l, r) ->
	    let kl = translate l k in
	    let kr = translate r (Push(Reg EAX):: kl) in
	      Sub(Reg ESP, Imm 4l)::Add(Reg EAX, Ind(ESP, 0))::kr
    in
      List.rev(translate ast [])
end


(* Example2 shows a simple intermediate language that uses named temporaries *)
(* This IR is a form of SSA -- single static assignment *)
module Example2 = struct
  open Ast

  type var = string
  type exp = Var of var | Imm of int32 
  type op  = Add of exp * exp   
  type cmd = Let of var * op * cmd 
           | Return of exp

  (* generate a fresh variable *)
  let mk_var =
    let ctr = ref 0 in
      fun hint -> let x = !ctr in ctr := x + 1; (hint ^ string_of_int x)

  (* Convert Ast's to SSA using an auxiliary list in reverse order: *)
  (* Note the similarity of rev_trans with Example1.compile *)
  type aux = ALet of var * op 
  let ast_to_cmd1 (ast:Ast.t):cmd =
    (* Result is a pair of the "returned expression" and 
     * auxiliary list of defined temoraries (in reverse order) *)
    let rec rev_trans (ast:Ast.t) (k:aux list) =
      match ast with
	| Int n -> (Imm n, k)
	| Plus(l, r) ->
	    let (e1, kl) = rev_trans l k in
	    let (e2, kr) = rev_trans r kl in
	    let tmp = mk_var "TMP" in
	      (Var tmp, ALet(tmp, Add(e1,e2))::kr) 
    in
    let (ret, rev_aux) = rev_trans ast [] in
      List.fold_left (fun c -> fun (ALet(v,o)) -> Let(v,o,c)) (Return ret) rev_aux

  (* Here, I'm using a technique called "continuation-passing" to generate the 
   * intermediate format.*)
  let ast_to_cmd2 (ast:Ast.t):cmd =
    let rec translate (ast:Ast.t) (k:exp -> cmd) =
      match ast with
	| Int n -> k (Imm n)
	| Plus(l, r) ->
	    translate l (fun e1 -> 
            translate r (fun e2 -> 
              let tmp = mk_var "TMP" in Let(tmp, Add(e1, e2), k (Var tmp))))
    in
      translate ast (fun e -> Return e)

  (* For those who are interested, this is a "defunctionalized" 
   * version of the above code *)
  type cont = KRet | KLeft of Ast.t * cont | KRight of exp * cont

  let ast_to_cmd3 (ast:Ast.t):cmd =
    let rec translate (ast:Ast.t) (k:cont) =
      match ast with
	| Int n -> apply k (Imm n)
	| Plus(l, r) ->
	    translate l (KLeft(r, k))
    and apply (k:cont) (e1:exp) =
      match k with
	| KRet -> Return e1
	| KLeft(r, k) -> translate r (KRight(e1, k))
	| KRight(e2, k) -> let tmp = mk_var "TMP" in 
	    Let(tmp, Add(e1, e2), apply k (Var tmp))
    in
      translate ast KRet


end