(* value-numbering-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 file contains an implementation of the hash-based value numbering
 * algorithm described in
 *
 *      Value Numbering
 *      by Preston Briggs, Keith Cooper, and Taylor Simpson
 *      CRPC-TR94517-S
 *      November 1994
 *)

functor ValueNumberingFn (D : DOMINANCE_TREE) : sig

    structure IR : SSA

    val transform : IR.program -> IR.program

  end = struct

    structure IR = D.IR
    structure E = ExprFn (IR)
    structure ValueMap = E.Map
    structure ST = Stats

    type expr = E.expr

  (********** Counters for statistics **********)
    val cntMeaninglessPhi       = ST.newCounter (IR.irName ^ ":meaningless-phi")
    val cntRedundantPhi         = ST.newCounter (IR.irName ^ ":redundant-phi")
    val cntRedundantAssign      = ST.newCounter (IR.irName ^ ":redundant-assign")

  (* adjust a variable's use count *)
    fun incUse (IR.V{useCnt, ...}) = (useCnt := !useCnt + 1)
    fun decUse (IR.V{useCnt, ...}) = (useCnt := !useCnt - 1)

    local
    (* property for mapping variables to their value number (VN), which is represented as a
     * SSA variable.  If their VN is different from themselves, then they are redundant.
     *)
      val {getFn=getVN, setFn=setVN, clrFn=clrVN, ...} = IR.Var.newProp (fn x => x)

    (* property for mapping value numbers to hash-consed expressions. *)
      val {getFn=getExp : IR.var -> expr, setFn=setExp, clrFn=clrExp, ...} =
            IR.Var.newProp (fn x => raise Fail(concat["getExp(", IR.Var.toString x, ")"]))

      datatype env = ENV of {
          avail : IR.var ValueMap.map   (* map from expressions to their value numbers, which *)
                                        (* are represented as SSA vars.  The domain are those *)
                                        (* expressions that are available. *)
        }
    in
    val emptyEnv = ENV{avail = ValueMap.empty}
  (* map variables to their hash-consed definition *)
    val getVN = getVN
    val setVN = setVN
    fun varToExp x = getExp(getVN x)
    fun bindVarToExp (ENV{avail}, x, e) = (
(*DEBUG**Log.msg(concat["** bindVarToExp: ", IR.Var.toString x, " --> ", E.toString e, "\n"]);*)
          setVN(x, x); setExp(x, e);
          ENV{avail = ValueMap.insert(avail, e, x)})
    fun expToVN (ENV{avail}, e) = ValueMap.find(avail, e)
  (* rename a variable if it's value number is different than itself *)
    fun rename x = let
          val x' = getVN x
          in
            if IR.Var.same(x, x')
              then x
              else (
(*DEBUG**Log.msg(concat["** rename ", IR.Var.toString x, " to ", IR.Var.toString x', "\n"]);*)
                decUse x; incUse x';
                x')
          end
  (* does a variable change; i.e., get replaced by another variable? *)
    fun changed x = not(IR.Var.same(x, getVN x))
  (* clear the properties of a variable *)
    fun clearVar x = (clrVN x; clrExp x)
  (* clear the properties from the variables of a node *)
    fun clearNode nd = List.app clearVar (IR.Node.defs nd)
    end (* local *)

    fun rewriteCFG cfg = let
        (* rewrite or delete a node, if necessary.  Note that we have already rewritten the JOIN nodes *)
          fun doNode nd = (case IR.Node.kind nd
                 of IR.COND{cond as ref x, ...} =>
                      if changed x
                        then cond := rename x
                        else ()
		  | IR.FOREACH{src as ref x, ...} =>
		      if changed x
			then src := rename x
			else ()
                  | IR.ASSIGN{stm=(y, rhs), succ, ...} =>
                      if changed y
                        then (
                        (* deleting redundant assignment *)
                          IR.RHS.app decUse rhs;
                          IR.CFG.deleteNode nd)
                      else if (List.exists changed (IR.RHS.vars rhs))
                      (* rewrite node to rename variables *)
                        then IR.CFG.replaceNode(nd, IR.Node.mkASSIGN(y, IR.RHS.map rename rhs))
                        else ()
                  | IR.MASSIGN{stm=([], rator, xs), succ, ...} =>
                      if (List.exists changed xs)
                      (* rewrite node to rename variables *)
                        then IR.CFG.replaceNode(nd, IR.Node.mkMASSIGN([], rator, List.map rename xs))
                        else ()
                  | IR.MASSIGN{stm=(ys, rator, xs), succ, ...} =>
                      if List.all changed ys
                        then (
                        (* deleting redundant assignment *)
                          List.app decUse xs;
                          IR.CFG.deleteNode nd)
                      else if (List.exists changed xs)
                      (* rewrite node to rename variables *)
                        then IR.CFG.replaceNode(nd, IR.Node.mkMASSIGN(ys, rator, List.map rename xs))
                        else ()
                  | IR.GASSIGN{lhs, rhs, ...} =>
                      if changed rhs
                        then IR.CFG.replaceNode(nd, IR.Node.mkGASSIGN(lhs, rename rhs))
                        else ()
                  | IR.NEW{strand, args, ...} =>
                      if List.exists changed args
                        then IR.CFG.replaceNode(nd, IR.Node.mkNEW{
                            strand=strand, args=List.map rename args
                          })
                        else ()
                  | IR.SAVE{lhs, rhs, ...} =>
                      if changed rhs
                        then IR.CFG.replaceNode(nd, IR.Node.mkSAVE(lhs, rename rhs))
                        else ()
                  | _ => ()
                (* end case *))
          val _ = List.app doNode (IR.CFG.sort cfg)
          in
            IR.CFG.apply clearNode cfg;
            cfg
          end

    fun transformCFG (liveIn, cfg) = let
          val tbl = E.new()
          val mkGLOBAL = E.mkGLOBAL tbl
          val mkSTATE = E.mkSTATE tbl
          val mkVAR = E.mkVAR tbl
          val mkLIT = E.mkLIT tbl
          val mkOP = E.mkOP tbl
          val mkMULTIOP = E.mkMULTIOP tbl
          val mkCONS = E.mkCONS tbl
          val mkSEQ = E.mkSEQ tbl
          val mkPHI = E.mkPHI tbl
          val mkEINAPP = E.mkEINAPP tbl
        (* convert a list of variables to a list of expressions *)
          fun varsToExp (env, xs) = List.map varToExp xs
        (* convert an SSA RHS into a hash-consed expression *)
          fun mkExp (env, rhs) = (case rhs
                 of IR.GLOBAL x => mkGLOBAL x
                  | IR.STATE x => mkSTATE x
                  | IR.VAR x => varToExp x
                  | IR.LIT l => mkLIT l
                  | IR.OP(rator, args) => mkOP(rator, varsToExp(env, args))
                  | IR.CONS(args, ty) => mkCONS(varsToExp(env, args), ty)
                  | IR.SEQ(args, ty) => mkSEQ(varsToExp(env, args), ty)
		  | IR.EINAPP(ein, args) => mkEINAPP(ein, varsToExp(env, args))
                (* end case *))
	  fun doPhi ((y, xs), (env, phis)) = let
		 val vn::vns = List.mapPartial (Option.map getVN) xs
		 in
		   if List.all (fn vn' => IR.Var.same(vn, vn')) vns
		     then ((* a meaningless phi node; map y to vn *)
(* DEBUG ** Log.msg(concat["** meaningless phi node: ", IR.phiToString (y, xs), "\n"]); *)
		       ST.tick cntMeaninglessPhi;
		       List.app (Option.app decUse) xs;
		       setVN(y, vn);
		       (env, phis))
		     else let
		       val exp = mkPHI(List.mapPartial (Option.map varToExp) xs)
		       in
			 case expToVN(env, exp)
			  of SOME vn' => ((* a redundant phi node *)
(* DEBUG ** Log.msg(concat["** redundant phi node: ", IR.phiToString (y, xs), "\n"]); *)
			       ST.tick cntRedundantPhi;
			       List.app (Option.app decUse) xs;
			       setVN(y, vn');
			       (env, phis))
			   | NONE => let
			       val xs = List.map (Option.map rename) xs
			       in
				 (bindVarToExp(env, y, exp), (y, xs)::phis)
			       end
			 (* end case *)
		       end
		 end
        (* walk the dominator tree computing value numbers *)
          fun vn (env, nd) = let
                val env = (case IR.Node.kind nd
                       of IR.JOIN{phis, ...} => let
			    val (env, remainingPhis) = List.foldl doPhi (env, []) (!phis)
                            in
                              phis := List.rev remainingPhis;
                              env
                            end
			| IR.FOREACH{phis, var, ...} => let
			    val (env, remainingPhis) = List.foldl doPhi (env, []) (!phis)
                            in
                              phis := List.rev remainingPhis;
                              bindVarToExp(env, var, mkVAR var)
                            end
                        | IR.ASSIGN{stm=(y, rhs), ...} => let
                            val exp = mkExp(env, rhs)
                            in
                              case expToVN(env, exp)
                               of SOME vn => ((* y is redundant, so map it to vn *)
(* DEBUG ** Log.msg(concat["** redundant assignment: ", IR.assignToString (y, rhs), *)
(* DEBUG **"; VN[", IR.Var.toString y, "] = ", IR.Var.toString vn, "\n"]); *)
                                    ST.tick cntRedundantAssign;
                                    setVN(y, vn);
                                    env)
                                | NONE => bindVarToExp(env, y, exp)
                              (* end case *)
                            end
                        | IR.MASSIGN{stm=(ys, rator, xs), ...} => let
                            val xs = varsToExp(env, xs)
                            fun mkExps (env, _, []) = env
                              | mkExps (env, i, y::ys) = let
                                  val exp = mkMULTIOP(i, rator, xs)
                                  in
                                    case expToVN(env, exp)
                                     of SOME vn => ((* y is redundant, so map it to vn *)
                                          ST.tick cntRedundantAssign;
                                          setVN(y, vn);
                                          mkExps (env, i+1, ys))
                                      | NONE => mkExps (bindVarToExp(env, y, exp), i+1, ys)
                                    (* end case *)
                                  end
                            in
                              mkExps (env, 0, ys)
                            end
                        | _ => env
                      (* end case *))
                in
                  List.app (fn nd => vn (env, nd)) (D.children nd)
                end
        (* define the initial environment by mapping the liveIn variables to themselves *)
          val env = List.foldl (fn (x, env) => bindVarToExp(env, x, mkVAR x)) emptyEnv liveIn
          in
            D.computeTree cfg;
          (* compute value numbers over the dominance tree *)
            vn (env, IR.CFG.entry cfg);
            D.clear cfg;
          (* delete and rewrite nodes as necessary *)
            rewriteCFG cfg before (List.app clearVar liveIn)
          end

    fun transform prog = let
          val IR.Program{
		  props, consts, inputs, constInit, globals, globalInit, create, strand, update
		} = prog
        (* transform the initial-strand creation code *)
          val create = let
		val IR.Create{dim, code} = create
		in
		  IR.Create{dim = dim, code = transformCFG ([], code)}
		end
        (* transform a strand *)
          fun transformStrand strand = let
		val IR.Strand{name, params, state, stateInit, initM, updateM, stabilizeM} = strand
                val stateInit = transformCFG (params, stateInit)
(* FIXME: what if a state variable becomes redundant? *)
                fun transformMeth body = transformCFG ([], body)
                in
                  IR.Strand{
                      name = name,
                      params = params,
                      state = state,
                      stateInit = stateInit,
		      initM = Option.map transformMeth initM,
		      updateM = transformMeth updateM,
		      stabilizeM = Option.map transformMeth stabilizeM
                    }
                end
          in
            IR.Program{
                props = props,
		consts = consts,
		inputs = inputs,
		constInit = transformCFG ([], constInit),
                globals = globals,
                globalInit = transformCFG ([], globalInit),
                create = create,
                strand = transformStrand strand,
		update = Option.map (fn cfg => transformCFG ([], cfg)) update
              }
          end

  end