(* translate-fn.sml
 *
 * This code is part of the Diderot Project (http://diderot-language.cs.uchicago.edu)
 *
 * COPYRIGHT (c) 2015 The University of Chicago
 * 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).
 *
 * FIXME: the handling of global variables needs more care.  For example, when translating from
 * HighIL to LowIL, field variables go away, but are replaced by the underlying image variable
 * in many uses.  If the image variable was not a user-defined global, then it is not visible
 * to the strand method code, which causes a IL check failure.
 *
 * FIXME: both this code and the Census code sets the bindings of variables (but the translation
 * from SimpleAST does not).
 *)

signature TRANSLATE_PARAMS =
  sig

    structure SrcIL : SSA
    structure DstIL : SSA

    type env

    val mkEnv : unit -> env

    val rename : (env * SrcIL.var) -> DstIL.var
    val renameList : (env * SrcIL.var list) -> DstIL.var list
    val renameSV : (env * SrcIL.state_var) -> DstIL.state_var
    val expand : (env * SrcIL.assign) -> DstIL.cfg
    val mexpand : (env * SrcIL.massign) -> DstIL.cfg

    val insertNd : (env * Stamp.stamp * DstIL.node) -> unit
    val findNd : env -> Stamp.stamp -> DstIL.node option

  end

functor TranslateFn (Params : TRANSLATE_PARAMS) : sig

    structure SrcIL : SSA
    structure DstIL : SSA

    val translate : SrcIL.program -> DstIL.program

  end = struct

    structure SrcIL = Params.SrcIL
    structure SrcNd = SrcIL.Node
    structure VTbl = SrcIL.Var.Tbl
    structure DstIL = Params.DstIL
    structure DstNd = DstIL.Node
    structure DstCFG = DstIL.CFG

    fun rename env x = Params.rename (env, x)

    fun renameList (env, xs) = Params.renameList(env, xs)

    fun renameSV env x = Params.renameSV (env, x)

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

    fun translateCFG (env, SrcIL.CFG{entry, exit}) = let
          val findNd = Params.findNd env
          fun trans (srcNd as SrcIL.ND{id, kind, ...}) = let
                fun newNd nd = (Params.insertNd (env, id, nd); nd)
                in
                  case findNd id
                   of SOME nd => nd
                    | NONE => (case kind
                         of SrcIL.NULL => raise Fail "unexpected NULL node"
                          | SrcIL.ENTRY{succ} => let
                              val nd = newNd (DstNd.mkENTRY())
                              in
                                DstNd.addEdge (nd, trans (!succ));
                                nd
                              end
                          | SrcIL.JOIN{phis, succ, ...} => let
                              fun cvtPhi (x, xs) = let
                                    val x = rename env x
                                    val xs = List.map (rename env) xs
                                    in
                                      DstIL.Var.setBinding (x, DstIL.VB_PHI xs);
                                      (x, xs)
                                    end
                              val nd = newNd (DstNd.mkJOIN(List.map cvtPhi (!phis)))
                              in
                                DstNd.addEdge (nd, trans (!succ));
                                nd
                              end
                          | SrcIL.COND{cond, trueBranch, falseBranch, ...} => let
                              val nd = newNd (DstNd.mkCOND{
                                      cond = rename env cond,
                                      trueBranch = DstNd.dummy,
                                      falseBranch = DstNd.dummy
                                    })
                              val trueB = trans (!trueBranch)
                              val _ = (DstNd.setTrueBranch (nd, trueB); DstNd.setPred(trueB, nd))
                              val falseB = trans (!falseBranch)
                              val _ = (DstNd.setFalseBranch (nd, falseB); DstNd.setPred(falseB, nd))
                              in
                                nd
                              end
                          | SrcIL.COM{text, succ, ...} => let
                              val nd = newNd (DstNd.mkCOM text)
                              in
                                DstNd.addEdge (nd, trans (!succ));
                                nd
                              end
                          | SrcIL.ASSIGN{stm, succ, ...} => let
                              val cfg = Params.expand (env, stm)
                              in
                                if DstCFG.isEmpty cfg
                                  then trans (!succ)
                                  else (
                                    DstNd.addEdge (DstCFG.exit cfg, trans (!succ));
                                    DstCFG.entry cfg)
                              end
                          | SrcIL.MASSIGN{stm, succ, ...} => let
                              val cfg = Params.mexpand (env, stm)
                              in
                                if DstCFG.isEmpty cfg
                                  then trans (!succ)
                                  else (
                                    DstNd.addEdge (DstCFG.exit cfg, trans (!succ));
                                    DstCFG.entry cfg)
                              end
                          | SrcIL.NEW{strand, args, succ, ...} => let
                              val nd = newNd (DstNd.mkNEW{
                                      strand = strand,
                                      args = List.map (rename env) args
                                    })
                              in
                                DstNd.addEdge (nd, trans (!succ));
                                nd
                              end
                          | SrcIL.SAVE{lhs, rhs, succ, ...} => let
                              val nd = newNd (DstNd.mkSAVE (renameSV env lhs, rename env rhs))
                              in
                                DstNd.addEdge (nd, trans (!succ));
                                nd
                              end
                          | SrcIL.EXIT{kind, live, ...} =>
                              newNd (DstNd.mkEXIT(kind, List.map (rename env) live))
                       (* end case *))
                  (* end case *)
                end
          val entry = trans entry
          val exit = (case findNd (SrcNd.id exit)
                 of SOME nd => nd
                  | NONE => DstNd.mkACTIVE()    (* exit is unreachable *)
                (* end case *))
          in
            DstIL.CFG{entry = entry, exit = exit}
          end

    fun translate (SrcIL.Program{props, globalInit, initially, strands}) = let
          val env = Params.mkEnv ()
          fun transInitially (SrcIL.Initially{isArray, rangeInit, iters, create}) = let
                val (argInit, strand, args) = create
                fun trIter (param, lo, hi) = let
                      val param = rename env param
                      in
                        DstIL.Var.setBinding(param, DstIL.VB_PARAM);
                        (param, rename env lo, rename env hi)
                      end
                val iters = List.map trIter iters
                in
                  DstIL.Initially{
                      isArray = isArray,
                      rangeInit = translateCFG (env, rangeInit),
                      create = (translateCFG (env, argInit), strand, renameList(env, args)),
                      iters = iters
                    }
                end
          fun transMethod state (SrcIL.Method{name, body}) = DstIL.Method{
                  name = name,
                  body = translateCFG (env, body)
                }
          fun transStrand (SrcIL.Strand{name, params, state, stateInit, methods}) = let
                val params = renameList (env, params)
                val state = List.map (renameSV env) state
                in
                  List.app (fn x => DstIL.Var.setBinding(x, DstIL.VB_PARAM)) params;
                  DstIL.Strand{
                      name = name,
                      params = params,
                      state = state,
                      stateInit = translateCFG (env, stateInit),
                      methods = List.map (transMethod state) methods
                    }
                end
          val prog = DstIL.Program{
                  props = props,
                  globalInit = translateCFG (env, globalInit),
                  initially = transInitially initially,
                  strands = List.map transStrand strands
                }
          in
            prog
          end

  end