(* rewrite-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.
 *
 * Infrastructure for rewriting the CFG.
 *)

signature REWRITE = sig
    structure IL : SSA
  (* rewrite the right-hand-side of an assignment; returns NONE if there are no changes
   * and SOME assigns if the assignment is rewritten to the assigns list.
   *)
    val doAssign : IL.assign -> (IL.var * IL.rhs) list option
  (* rewrite the rhs of an multi-assignment *)
    val doMAssign : IL.massign -> (IL.var * IL.rhs) list option
  (* eliminates unused variables; returns true if any changes *)
    val elimUnusedVars : IL.cfg -> bool
  end

functor RewriteFn (R : REWRITE) : sig

    val transform : R.IL.program -> R.IL.program

  end = struct

    open R

    fun useCount (IL.V{useCnt, ...}) = !useCnt

  (* simplify assignment and multi-assignment statements *)
    fun simplify nd = let
	  fun rewrite (SOME[]) = (IL.CFG.deleteNode nd; true)
	    | rewrite (SOME assigns) = let
		val assigns = List.map
		      (fn (y, rhs) => (IL.Var.setBinding(y, IL.VB_RHS rhs); IL.ASSGN(y, rhs)))
			assigns
		in
		  IL.CFG.replaceNodeWithCFG (nd, IL.CFG.mkBlock assigns);
		  true
		end
	    | rewrite NONE = false
	  in
	    case IL.Node.kind nd
	     of IL.ASSIGN{stm=(y, rhs), ...} => if (useCount y = 0)
		  then false (* skip unused assignments *)
		  else rewrite (doAssign (y, rhs))
	      | IL.MASSIGN{stm, ...} => rewrite (doMAssign stm)
	      | _ => false
	    (* end case *)
	  end

    fun loopToFixPt f = let
	  fun loop anyChanges = if f() then loop true else anyChanges
	  in
	    loop false
	  end

    fun transform (prog as IL.Program{props, globalInit, initially, strands}) = let
	  fun simplifyCFG cfg = let
		val changed = ref false
		fun simplify' nd = if simplify nd then changed := true else ()
		in
		  IL.CFG.apply simplify' cfg;
		  !changed
		end
	  fun doCFG cfg = let
		val changes = loopToFixPt (fn () => simplifyCFG cfg)
		val changes = loopToFixPt (fn () => elimUnusedVars cfg) orelse changes
		in
		  changes
		end
	  fun doMethod (IL.Method{body, ...}) = doCFG body
	  fun doStrand (IL.Strand{stateInit, methods, ...}) = let
		val changes = doCFG stateInit
		val changes = List.foldl (fn (m, flg) => doMethod m orelse flg) changes methods
		in
		  changes
		end
	  fun optPass () = let
		val changes = doCFG globalInit
		val changes = List.foldl (fn (s, flg) => doStrand s orelse flg) changes strands
		in
		  changes
		end
	  in
	    loopToFixPt optPass;
(* FIXME: after optimization, we should filter out any globals that are now unused *)
	    IL.Program{
		props = props,
		globalInit = globalInit,
		initially = initially,	(* FIXME: we should optimize this code *)
		strands = strands
	      }
	  end

  end