(* 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 : string * LowIL.program -> unit

  end = struct

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

  (* convert LowIL types to T types *)
    fun cvtTy ty = (case ty
	   of Ty.BoolTy => T.boolTy
	    | Ty.StringTy => raise Fail "FIXME: StringTy"
	    | Ty.IVecTy 1 => T.intTy
	    | Ty.IVecTy n => T.ivecTy n		(* FIXME: what about vector splits? *)
	    | Ty.VecTy 1 => T.realTy
	    | Ty.VecTy n => T.vecTy n		(* FIXME: what about vector splits? *)
	    | Ty.AddrTy => raise Fail "FIXME: AddrTy"
	    | Ty.ImageTy =>  raise Fail "FIXME: ImageTy"
	  (* end case *))

  (* 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 -> IL.var -> T.exp
	val useDefOf : tbl -> IL.var -> T.exp
	val setDefOf : tbl -> (IL.var * T.exp) -> unit

	val clear : tbl -> unit

	val assign : tbl -> (IL.var * T.exp) -> T.stm list
	val bind : tbl -> (IL.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 => raise Fail(concat["getDefOf(", V.toString x, ")"])
		| SOME{bind, cnt} => bind
	      (* end case *))

	fun useDefOf tbl x = (case V.Tbl.find tbl x
	       of NONE => raise Fail(concat["useDefOf(", V.toString x, ")"])
		| 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.Var.tmp(cvtTy(V.ty x))
		    in
		      V.Tbl.insert tbl
			(x, {cnt = ref(V.useCount x), bind = T.Expr.var lhs});
		      [T.Stmt.assign(lhs, exp)]
		    end

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

	val clear = V.Tbl.clear

	fun flush tbl = raise Fail "flush"

      end (*  VDefTbl *)

(* FIXME: what about splitting code where the target width doesn't match the
 * source width?
 *)
    fun doRator vtbl = let
	  val exp = VDefTbl.useDefOf vtbl
	  val bind = VDefTbl.bind vtbl
	  fun f (lhs, rator, args) = (case (rator, args)
		 of (Op.Add ty, [a, b])		=> bind (lhs, T.Expr.add(exp a, exp b))
		  | (Op.Sub ty, [a, b])		=> bind (lhs, T.Expr.sub(exp a, exp b))
		  | (Op.Mul ty, [a, b])		=> bind (lhs, T.Expr.mul(exp a, exp b))
		  | (Op.Div ty, [a, b])		=> bind (lhs, T.Expr.divide(exp a, exp b))
		  | (Op.Neg ty, [a])		=> bind (lhs, T.Expr.neg(exp a))
		  | (Op.LT ty, [a, b])		=> bind (lhs, T.Expr.lt(exp a, exp b))
		  | (Op.LTE ty, [a, b])		=> bind (lhs, T.Expr.lte(exp a, exp b))
		  | (Op.EQ ty, [a, b])		=> bind (lhs, T.Expr.equ(exp a, exp b))
		  | (Op.NEQ ty, [a, b])		=> bind (lhs, T.Expr.neq(exp a, exp b))
		  | (Op.GT ty, [a, b])		=> bind (lhs, T.Expr.gt(exp a, exp b))
		  | (Op.GTE ty, [a, b])		=> bind (lhs, T.Expr.gte(exp a, exp b))
		  | (Op.Not, [a])		=> bind (lhs, T.Expr.not(exp a))
		  | (Op.Max, [a, b])		=> bind (lhs, T.Expr.max(exp a, exp b))
		  | (Op.Min, [a, b])		=> bind (lhs, T.Expr.min(exp a, exp b))
		  | (Op.Sin, [a])		=> bind (lhs, T.Expr.sin(exp a))
		  | (Op.Cos, [a])		=> bind (lhs, T.Expr.cos(exp a))
		  | (Op.Pow, [a, b])		=> bind (lhs, T.Expr.pow(exp a, exp b))
		  | (Op.Dot d, [a, b])		=> bind (lhs, T.Expr.dot(exp a, exp b))
		  | (Op.Cross, [a, b])		=> bind (lhs, T.Expr.cross(exp a, exp b))
		  | (Op.Select(ty, i), [a])	=> bind (lhs, T.Expr.select(i, exp a))
		  | (Op.Norm d, [a])		=> bind (lhs, T.Expr.length(exp a))
		  | (Op.Scale d, [a, b])	=> bind (lhs, T.Expr.mul(exp a, exp b))
		  | (Op.InvScale d, [a, b])	=> bind (lhs, T.Expr.divide(exp a, exp b))
		  | (Op.CL, _)			=> raise Fail "CL unimplemented"
		  | (Op.PrincipleEvec ty, _)	=> raise Fail "PrincipleEvec unimplemented"
(*
		  | (Op.Subscript ty,
*)
		  | (Op.Floor d, [a])		=> bind (lhs, T.Expr.floor(exp a))
		  | (Op.IntToReal, [a])		=> bind (lhs, T.Expr.toReal(exp a))
		  | (Op.TruncToInt d, [a])	=> bind (lhs, T.Expr.truncToInt(exp a))
		  | (Op.RoundToInt d, [a])	=> bind (lhs, T.Expr.roundToInt(exp a))
		  | (Op.CeilToInt d, [a])	=> bind (lhs, T.Expr.ceilToInt(exp a))
		  | (Op.FloorToInt d, [a])	=> bind (lhs, T.Expr.floorToInt(exp a))
		  | (Op.ImageAddress, [a])	=> bind (lhs, T.Expr.imageAddr(exp a))
(*
		  | (Op.LoadVoxels(rTy, n), [a]) =>
		  | (Op.PosToImgSpace d, [v, x]) =>
		  | (Op.GradToWorldSpace d, [v, x]) =>
		  | (Op.LoadImage info, [a]) =>
		  | (Op.Inside d, [v, x]) =>
		  | (Op.Input(ty, name), []) =>
		  | (Op.InputWithDefault(ty, name), [a]) =>
*)
		  | _ => raise Fail "incorrect number of arguments"
		(* end case *))
	  in
	    f
	  end

  (* translate a LowIL assignment to a list of zero or more target statements *)
    fun doAssign vtbl = let
	  val setDefOf = VDefTbl.setDefOf vtbl
	  val doRator = doRator vtbl
	  fun assign (lhs, rhs) = let
		fun setDef rhs = (setDefOf (lhs, rhs); [])
		in
		  case rhs
		   of IL.VAR x => setDef (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(lhs, rator, args)
		    | IL.CONS args =>
			VDefTbl.assign vtbl
			  (lhs, T.Expr.vector (List.map (VDefTbl.useDefOf vtbl) args))
		  (* end case *)
		end
	  in
	    assign
	  end

    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, ifCont, stms, !succ)
		  | IL.JOIN{phis, succ, ...} => ifCont (stms, Nd.kind 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, ifCont, stm::stms, !succ)
				  end
			    in
			      doNode (vtbl, kElse, [], !falseBranch)
			    end
		      in
			doNode (vtbl, kThen, [], !trueBranch)
		      end
		  | IL.COM {text, succ, ...} =>
		      doNode (vtbl, ifCont, T.Stmt.comment text :: stms, !succ)
		  | IL.ASSIGN{stm, succ, ...} =>
		      doNode (vtbl, ifCont, doAssign stm @ stms, !succ)
		  | IL.NEW{strand, args, succ, ...} => raise Fail "NEW unimplemented"
		  | IL.DIE _ =>
		      T.Stmt.block (List.rev (T.Stmt.die() :: stms))
		  | IL.STABILIZE _ =>
		      T.Stmt.block (List.rev stms)
		  | IL.EXIT _ => T.Stmt.block (List.rev (T.Stmt.stabilize() :: stms))
		(* end case *))
	  in
	    doNode (vtbl, fn _ => raise Fail "bogus ifCont at JOIN node", [], CFG.entry cfg)
	  end

    fun generate (fileStem, IL.Program{globals, globalInit, strands}) = ()

  end