(* codegen-fn.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Generic support for translating LowIL code to the target representation.  We
 * assume that the LowIL has first been run through the splitting pass to match
 * the target's vector widths.
 *)

functor CodeGenFn (T : TARGET) : sig

    val generate : LowIL.program -> T.program

  end = struct

    structure IL = LowIL
    structure V = LowIL.Var
    structure Op = LowOps
    structure Nd = LowIL.Node
    structure CFG = LowIL.Node

  (* a mapping from LowIL variables to target expressions.  Variables get
   * removed when the expressions are used as arguments.
   *)
    structure VDefTbl : sig
	type tbl

	val newTbl : unit -> tbl
	val getDefOf : tbl -> V.var -> T.exp
	val useDefOf : tbl -> V.var -> T.exp
	val setDefOf : tbl -> (V.var * T.exp) -> unit

	val clear : tbl -> unit

	val bind : tbl -> (V.var * T.exp) -> T.stm list

      (* force all pending expressions into variables *)
	val flush : tbl -> T.stm list

      end = struct

	type info = {
	    cnt : int ref,	(* count of oustanding uses (usually 1) *)
	    bind : T.exp
	  }

	type tbl = info V.Tbl.hash_table

	fun newTbl () = V.Tbl.mkTable (512, Fail "vtbl")

	fun getDefOf tbl x = (case V.Tbl.find tbl x
	       of NONE => ??
		| SOME{bind, cnt} => bind
	      (* end case *))

	fun useDefOf tbl x = (case V.Tbl.find tbl x
	       of NONE => ??
		| SOME{cnt=ref 1, bind} => (
		    ignore (V.Tbl.remove tbl x);
		    bind)
		| SOME{cnt, bind} =>  => (
		    cnt := !cnt - 1;
		    bind)
	      (* end case *))

	fun setDefOf tbl (x, exp) =
	      V.Tbl.insert tbl (x, {cnt = ref(V.useCount x), bind = exp})

	fun assign tbl (x, exp) = let
		    val lhs : T.local_var = ??
		    in
		      V.Tbl.insert tbl
			(x, {cnt = V.useCount x, bind = T.Expr.var lhs});
		      [T.Stmt.assign(lhs, exp)]
		    end

	fun bind tbl (x, exp) = (case V.useCount lhs
	       of 1 => (V.Tbl.insert tbl (x, {cnt = 1, bind = exp}); [])
		| n => let (* bind exp to a new target variable *)
		    val lhs : T.local_var = ??
		    in
		      V.Tbl.insert tbl (x, {cnt = n, bind = T.Expr.var lhs});
		      [T.Stmt.assign(lhs, exp)]
		    end
	      (* end case *))

      end (*  VDefTbl *)

(* FIXME: what about splitting code where the target width doesn't match the
 * source width?
 *)
    fun doRator (vtbl, lhs, rator, args) = let
	  val args' = List.map (VDefTbl.useDefOf vtbl) args
	  val rhs' = (case (rator, args)
		 of (Op.Add ty, [a, b]) =>
		  | (Op.Sub ty, [a, b]) =>
		  | (Op.Mul ty, [a, b]) =>
		  | (Op.Div ty, [a, b]) =>
		  | (Op.Neg ty, [a]) =>
		  | (Op.LT ty, [a, b]) =>
		  | (Op.LTE ty, [a, b]) =>
		  | (Op.EQ ty, [a, b]) =>
		  | (Op.NEQ ty, [a, b]) =>
		  | (Op.GT ty, [a, b]) =>
		  | (Op.GTE ty, [a, b]) =>
		  | (Op.Not, [a]) =>
		  | (Op.Max, [a, b]) =>
		  | (Op.Min, [a, b]) =>
		  | (Op.Sin, [a]) =>
		  | (Op.Cos, [a]) =>
		  | (Op.Pow, [a, b]) =>
		  | (Op.Dot d, [a, b]) =>
		  | (Op.Cross, [a, b]) =>
		  | (Op.Select(ty, i), [a]) =>
		  | (Op.Norm d, [a]) =>
		  | (Op.Scale d, [a, b]) =>
		  | (Op.InvScale d, [a, b]) =>
		  | (Op.CL
		  | (Op.PrincipleEvec of ty
		  | (Op.Subscript of ty
		  | (Op.Floor of int
		  | (Op.IntToReal
		  | (Op.TruncToInt of int
		  | (Op.RoundToInt of int
		  | (Op.CeilToInt of int
		  | (Op.FloorToInt of int
		  | (Op.ImageAddress of ImageInfo.info
		  | (Op.LoadVoxels of RawTypes.ty * int
		  | (Op.PosToImgSpace of ImageInfo.info
		  | (Op.GradToWorldSpace of ImageInfo.info
		  | (Op.LoadImage of ImageInfo.info
		  | (Op.Inside of ImageInfo.info
		  | (Op.Input of ty * string
		  | (Op.InputWithDefault of ty * string
		(* end case *))
	  in
	    VDefTbl.bind vtbl (lhs, rhs')
	  end

  (* translate a LowIL assignment to a list of zero or more target statements *)
    fun doAssign vtbl (lhs, rhs) = let
	  fun setDef rhs = (VTbl.setDefOf vtbl (lhs, rhs); [])
	  in
	    case rhs
	     of IL.VAR x => setDef (T.Expr.var(VDefTbl.useDefOf vtbl x))
	      | IL.LIT(Literal.Int n) => setDef (T.Expr.intLit n)
	      | IL.LIT(Literal.Bool b) => setDef (T.Expr.boolLit b)
	      | IL.LIT(Literal.Float f) => setDef (T.Expr.floatLit f)
	      | IL.LIT(Literal.String s) => setDef (T.Expr.stringLit s)
	      | IL.OP(rator, args) => doRator(vtbl, lhs, rator, args)
	      | IL.CONS args =>
		  VTbl.assign ctbl (lhs, T.Expr.vector (List.map (VDefTbl.useDefOf vtbl) args))
	    (* end case *)
	  end

    datatype open_if = T.stm list * IL.node -> stm

    fun gen (vtbl, cfg) = let
	  val doAssign = doAssign vtbl
	  fun doNode (vtbl, ifCont, stms, nd) = (case Nd.kind nd
		 of IL.NULL => raise Fail "unexpected NULL"
		  | IL.ENTRY{succ} => doNode (vtbl, ifStk, stms, !succ)
		  | IL.JOIN{phis, succ, ...} => ifCont (stms, nd)
		  | IL.COND{cond, trueBranch, falseBranch, ...} => let
		      fun kThen (stms', _) = let
			    val thenBlk = T.Stmt.block (List.rev stms')
			    fun kElse (stms', IL.JOIN{phis, succ, ...}) = let
				  val stm = T.Stmt.ifthenelse (
					VDefTbl.useDefOf vtbl cond,
					thenBlk,
					T.Stmt.block (List.rev stms'))
				  in
(* FIXME: what do we do about phis? *)
				    doNode (vtbl, ifStk, stm::stms, !succ)
				  end
			    in
			      doNode (vtbl, kElse::ifStk, [], !falseBranch)
			    end
		      in
			doNode (vtbl, kThen::ifStk, [], !trueBranch)
		      end
		  | IL.COM {text, succ, ...} =>
		      doNode (vtbl, ifStk, T.Stmt.comment text :: stms, !succ)
		  | IL.ASSIGN{stm, succ, ...} =>
		      doNode (vtbl, ifStk, doAssign stm :: stms, !succ)
		  | IL.NEW{strand, args, succ, ...} => raise Fail "NEW unimplemented"
		  | IL.DIE _ =>
		      doNode (vtbl, ifStk, T.Stmt.die() :: stms, !succ)
		  | IL.STABILIZE _ =>
		      doNode (vtbl, ifStk, T.Stmt.stabilize() :: stms, !succ)
		  | IL.EXIT _ => T.Stmt.mkBlock (List.rev stms)
		(* end case *))
	  in
	    doNode (vtbl, fn _ => raise Fail "bogus ifCont at JOIN node", [], CFG.entry cfg)
	  end

  end