(* translate-fn.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot.cs.uchicago.edu)
 * All rights reserved.
 *
 * This functor supports the common parts of translating between different
 * instances of the SSA-based ILs (e.g., from HighIL to MidIL).
 *)

signature TRANSLATE_PARAMS =
  sig

    structure SrcIL : SSA
    structure DstIL : SSA

    type var_env = DstIL.var SrcIL.Var.Tbl.hash_table

    val rename : (var_env * SrcIL.var) -> DstIL.var
    val expand : (var_env * SrcIL.assign) -> DstIL.assign list

  end

functor TranslateFn (Params : TRANSLATE_PARAMS) : sig

    structure SrcIL : SSA
    structure DstIL : SSA

    type var_env = DstIL.var SrcIL.Var.Tbl.hash_table

    val translate : var_env * SrcIL.stmt -> DstIL.stmt

  end = struct

    structure SrcIL : SSA = Params.SrcIL
    structure SrcNd = SrcIL.Node
    structure DstIL : SSA = Params.DstIL
    structure DstNd = DstIL.Node
    structure DstStm = DstIL.Stmt

    type var_env = Params.var_env

    datatype env = E of {
	ndMap : DstIL.node Stamp.Tbl.hash_table,
	vMap : var_env
      }

    fun rename (E{vMap, ...}) x = Params.rename(vMap, x)

    fun expand (E{vMap, ...}) (assign, assigns') =
	  Params.expand (vMap, assign) @ assigns'

    fun insertNd (E{ndMap, ...}, id, nd) = Stamp.Tbl.insert ndMap (id, nd)

    fun renameNd (E{ndMap, ...}) (nd as SrcIL.ND{id, ...}) = (
	  case Stamp.Tbl.find ndMap id
	   of SOME nd' => nd'
	    | NONE => raise Fail("unable to find " ^ SrcNd.toString nd)
	  (* end case *))

  (* the first pass creates the nodes of the DstIL CFG and defines
   * the environment that maps from SrcIL nodes and variables to
   * DstIL nodes and variables.
   *)
    fun translateNodes (env, stm) = let
	  fun trans (SrcIL.ND{id, kind, ...}) = let
		val newNd = (case kind
		     of SrcIL.NULL => raise Fail "unexpected NULL node"
		      | SrcIL.ENTRY _ => DstNd.mkENTRY()
		      | SrcIL.JOIN{phis, ...} => let
			  fun cvtPhi (x, xs) =
				(rename env x, List.map (rename env) xs)
			  in
			    DstNd.mkJOIN(List.map cvtPhi (!phis))
			  end
		      | SrcIL.COND{cond, ...} => DstNd.mkCOND{
			    cond = rename env cond,
			    trueBranch = DstNd.dummy,
			    falseBranch = DstNd.dummy
			  }
		      | SrcIL.BLOCK{body, ...} => let
			  val body' = List.foldr (expand env) [] (!body)
			  in
			    DstNd.mkBLOCK body'
			  end
		      | SrcIL.NEW{actor, args, ...} => DstNd.mkNEW{
			    actor = actor,
			    args = List.map (rename env) args
			  }
		      | SrcIL.DIE _ => DstNd.mkDIE()
		      | SrcIL.STABILIZE _ => DstNd.mkSTABILIZE()
		      | SrcIL.EXIT _ => DstNd.mkEXIT()
		   (* end case *))
		in
		  insertNd (env, id, newNd)
		end
	  in
	    SrcIL.applyToNodes trans stm
	  end

  (* the second pass copys the statement tree and sets the CFG edges; it
   * returns the new statement tree.
   *)
    fun translateStmts (env, stm) = let
	  val renameNd = renameNd env
	(* set the CFG edges of the node corresponding to the source node *)
	  fun setEdges (srcNd as SrcIL.ND{kind, ...}) = let
		val dstNd as DstIL.ND{kind=dstKind, ...} = renameNd srcNd
		in
		  case kind
		   of SrcIL.NULL => raise Fail "unexpected NULL node"
		    | SrcIL.ENTRY{succ} => DstNd.setSucc(dstNd, renameNd(!succ))
		    | SrcIL.JOIN{preds, succ, ...} => let
			val DstIL.JOIN{preds=dstPreds, ...} = dstKind
			in
			  dstPreds := List.map renameNd (!preds);
			  DstNd.setSucc (dstNd, renameNd(!succ))
			end
		    | SrcIL.COND{pred, trueBranch, falseBranch, ...} => (
			DstNd.setPred (dstNd, renameNd(!pred));
			DstNd.setTrueBranch (dstNd, renameNd(!trueBranch));
			DstNd.setFalseBranch (dstNd, renameNd(!falseBranch)))
		    | SrcIL.BLOCK{pred, succ, ...} => (
			DstNd.setPred (dstNd, renameNd(!pred));
			DstNd.setSucc (dstNd, renameNd(!succ)))
		    | SrcIL.NEW{pred, succ, ...} => (
			DstNd.setPred (dstNd, renameNd(!pred));
			DstNd.setSucc (dstNd, renameNd(!succ)))
		    | SrcIL.DIE{pred} => DstNd.setPred (dstNd, renameNd(!pred))
		    | SrcIL.STABILIZE{pred} => DstNd.setPred (dstNd, renameNd(!pred))
		    | SrcIL.EXIT{pred} => DstNd.setPred (dstNd, renameNd(!pred))
		  (* end case *)
		end
	(* translate statements *)
	  fun trans (SrcIL.STM{kind, next, ...}) = let
		fun new kind' = DstStm.new(kind', Option.map trans next)
		in
		  case kind
		   of SrcIL.S_SIMPLE nd => new (DstIL.S_SIMPLE(renameNd nd))
		    | SrcIL.S_IF{cond, thenBranch, elseBranch} => new (DstIL.S_IF{
			  cond = renameNd cond,
			  thenBranch = trans thenBranch,
			  elseBranch = trans elseBranch
			})
		    | SrcIL.S_LOOP{hdr, cond, body} => new (DstIL.S_LOOP{
			  hdr = trans hdr,
			  cond = renameNd cond,
			  body = trans body
			})
		  (* end case *)
		end
	  in
	    SrcIL.applyToNodes setEdges stm;
	    trans stm
	  end

    fun translate (vMap, stm) = let
	  val env = E{
		  ndMap = Stamp.Tbl.mkTable (256, Fail "ndMap"),
		  vMap = vMap
		}
	  val _ = translateNodes (env, stm)
	  in
	    translateStmts (env, stm)
	  end

  end